diff --git a/Content.Server/Antag/AntagSelectionPlayerPool.cs b/Content.Server/Antag/AntagSelectionPlayerPool.cs index 87873e96d1a..f778b580471 100644 --- a/Content.Server/Antag/AntagSelectionPlayerPool.cs +++ b/Content.Server/Antag/AntagSelectionPlayerPool.cs @@ -2,6 +2,7 @@ using System.Linq; using Robust.Shared.Player; using Robust.Shared.Random; +using Robust.Shared.Utility; namespace Content.Server.Antag; @@ -23,5 +24,46 @@ public bool TryPickAndTake(IRobustRandom random, [NotNullWhen(true)] out ICommon return session != null; } + // EE + public bool TryGetItems(IRobustRandom random, + [NotNullWhen(true)] out ICommonSession[]? sessions, + int count, + bool allowDuplicates = true) + { + DebugTools.Assert(count > 0, $"The count {nameof(count)} of requested sessions must be greater than zero!"); + + sessions = null; + List session_list = []; + + foreach (var pool in orderedPools) + { + if (pool.Count == 0) + continue; + + var picked = random.GetItems(pool, count - session_list.Count, allowDuplicates); + session_list.AddRange(picked); + if (session_list.Count < count) + { + continue; + } + sessions = session_list.ToArray(); + break; + } + + return sessions != null; + } + + // EE + public AntagSelectionPlayerPool Where(Func predicate) + { + var newPools = orderedPools.Select( + (pool) => + { + return pool.Where(predicate).ToList(); + } + ); + + return new AntagSelectionPlayerPool(newPools.ToList()); + } public int Count => orderedPools.Sum(p => p.Count); } diff --git a/Content.Server/Antag/AntagSelectionSystem.cs b/Content.Server/Antag/AntagSelectionSystem.cs index 8efdc2738b9..34448dd9530 100644 --- a/Content.Server/Antag/AntagSelectionSystem.cs +++ b/Content.Server/Antag/AntagSelectionSystem.cs @@ -204,20 +204,68 @@ public void ChooseAntags(Entity ent, IList failed = []; + + while (ent.Comp.SelectedSessions.Count < count && retry < maxRetries) { - if (!playerPool.TryPickAndTake(RobustRandom, out session)) + var sessions = (ICommonSession[]?) null; + if (!playerPool.TryGetItems(RobustRandom, + out sessions, + count - ent.Comp.SelectedSessions.Count, + false)) + break; // Ends early if there are no eligible sessions + + foreach (var session in sessions) + { + MakeAntag(ent, session, def); + if (!ent.Comp.SelectedSessions.Contains(session)) + { + failed.Add(session); + } + } + // In case we're done + if (ent.Comp.SelectedSessions.Count >= count) break; - if (ent.Comp.SelectedSessions.Contains(session)) - continue; + playerPool = playerPool.Where((session_) => + { + return !ent.Comp.SelectedSessions.Contains(session_) && + !failed.Contains(session_); + }); + retry++; } + } + + // This preserves previous behavior for when def.PickPlayer + // was not satisfied. This behavior is not that obvious to + // read from the previous code. + // It may otherwise process leftover slots if maxRetries have + // been reached. - MakeAntag(ent, session, def); + for (var i = ent.Comp.SelectedSessions.Count; i < count; i++) + { + MakeAntag(ent, null, def); } + ///// End of Einstein Engines changes ///// } /// @@ -261,17 +309,16 @@ public void MakeAntag(Entity ent, ICommonSession? sessi { var getEntEv = new AntagSelectEntityEvent(session, ent); RaiseLocalEvent(ent, ref getEntEv, true); - - if (!getEntEv.Handled) - { - throw new InvalidOperationException($"Attempted to make {session} antagonist in gamerule {ToPrettyString(ent)} but there was no valid entity for player."); - } - antagEnt = getEntEv.Entity; } if (antagEnt is not { } player) + { + Log.Error($"Attempted to make {session} antagonist in gamerule {ToPrettyString(ent)} but there was no valid entity for player."); + if (session != null) + ent.Comp.SelectedSessions.Remove(session); return; + } var getPosEv = new AntagSelectLocationEvent(session, ent); RaiseLocalEvent(ent, ref getPosEv, true); @@ -282,11 +329,15 @@ public void MakeAntag(Entity ent, ICommonSession? sessi _transform.SetMapCoordinates((player, playerXform), pos); } + // If we want to just do a ghost role spawner, set up data here and then return early. + // This could probably be an event in the future if we want to be more refined about it. if (isSpawner) { if (!TryComp(player, out var spawnerComp)) { - Log.Error("Antag spawner with GhostRoleAntagSpawnerComponent."); + Log.Error($"Antag spawner {player} does not have a GhostRoleAntagSpawnerComponent."); + if (session != null) + ent.Comp.SelectedSessions.Remove(session); return; }