Fix antag objectives always overshooting MaxDifficulty (and kill tries20) (#29830)

* The death of try20

* Add integration test for traitor gamerule

* Fix max difficulty being overshot

* Check at least one objective is assigned

* EntProtoId

Content.IntegrationTests/Tests/GameRules/TraitorRuleTest.cs

@@ -0,0 +1,133 @@
+using System.Linq;
+using Content.Server.Antag.Components;
+using Content.Server.GameTicking;
+using Content.Server.GameTicking.Rules;
+using Content.Server.GameTicking.Rules.Components;
+using Content.Server.Mind;
+using Content.Server.Roles;
+using Content.Shared.GameTicking;
+using Content.Shared.GameTicking.Components;
+using Content.Shared.Mind;
+using Content.Shared.NPC.Systems;
+using Content.Shared.Objectives.Components;
+using Robust.Shared.GameObjects;
+using Robust.Shared.Prototypes;
+
+namespace Content.IntegrationTests.Tests.GameRules;
+
+[TestFixture]
+public sealed class TraitorRuleTest
+{
+    private const string TraitorGameRuleProtoId = "Traitor";
+    private const string TraitorAntagRoleName = "Traitor";
+
+    [Test]
+    public async Task TestTraitorObjectives()
+    {
+        await using var pair = await PoolManager.GetServerClient(new PoolSettings()
+        {
+            Dirty = true,
+            DummyTicker = false,
+            Connected = true,
+            InLobby = true,
+        });
+        var server = pair.Server;
+        var client = pair.Client;
+        var entMan = server.EntMan;
+        var protoMan = server.ProtoMan;
+        var compFact = server.ResolveDependency<IComponentFactory>();
+        var ticker = server.System<GameTicker>();
+        var mindSys = server.System<MindSystem>();
+        var roleSys = server.System<RoleSystem>();
+        var factionSys = server.System<NpcFactionSystem>();
+        var traitorRuleSys = server.System<TraitorRuleSystem>();
+
+        // Look up the minimum player count and max total objective difficulty for the game rule
+        var minPlayers = 1;
+        var maxDifficulty = 0f;
+        await server.WaitAssertion(() =>
+        {
+            Assert.That(protoMan.TryIndex<EntityPrototype>(TraitorGameRuleProtoId, out var gameRuleEnt),
+            $"Failed to lookup traitor game rule entity prototype with ID \"{TraitorGameRuleProtoId}\"!");
+
+            Assert.That(gameRuleEnt.TryGetComponent<GameRuleComponent>(out var gameRule, compFact),
+            $"Game rule entity {TraitorGameRuleProtoId} does not have a GameRuleComponent!");
+
+            Assert.That(gameRuleEnt.TryGetComponent<AntagRandomObjectivesComponent>(out var randomObjectives, compFact),
+            $"Game rule entity {TraitorGameRuleProtoId} does not have an AntagRandomObjectivesComponent!");
+
+            minPlayers = gameRule.MinPlayers;
+            maxDifficulty = randomObjectives.MaxDifficulty;
+        });
+
+        // Initially in the lobby
+        Assert.That(ticker.RunLevel, Is.EqualTo(GameRunLevel.PreRoundLobby));
+        Assert.That(client.AttachedEntity, Is.Null);
+        Assert.That(ticker.PlayerGameStatuses[client.User!.Value], Is.EqualTo(PlayerGameStatus.NotReadyToPlay));
+
+        // Add enough dummy players for the game rule
+        var dummies = await pair.Server.AddDummySessions(minPlayers);
+        await pair.RunTicksSync(5);
+
+        // Initially, the players have no attached entities
+        Assert.That(pair.Player?.AttachedEntity, Is.Null);
+        Assert.That(dummies.All(x => x.AttachedEntity == null));
+
+        // Opt-in the player for the traitor role
+        await pair.SetAntagPreference(TraitorAntagRoleName, true);
+
+        // Add the game rule
+        var gameRuleEnt = ticker.AddGameRule(TraitorGameRuleProtoId);
+        Assert.That(entMan.TryGetComponent<TraitorRuleComponent>(gameRuleEnt, out var traitorRule));
+
+        // Ready up
+        ticker.ToggleReadyAll(true);
+        Assert.That(ticker.PlayerGameStatuses.Values.All(x => x == PlayerGameStatus.ReadyToPlay));
+
+        // Start the round
+        await server.WaitPost(() =>
+        {
+            ticker.StartRound();
+            // Force traitor mode to start (skip the delay)
+            ticker.StartGameRule(gameRuleEnt);
+        });
+        await pair.RunTicksSync(10);
+
+        // Game should have started
+        Assert.That(ticker.RunLevel, Is.EqualTo(GameRunLevel.InRound));
+        Assert.That(ticker.PlayerGameStatuses.Values.All(x => x == PlayerGameStatus.JoinedGame));
+        Assert.That(client.EntMan.EntityExists(client.AttachedEntity));
+
+        // Check the player and dummies are spawned
+        var dummyEnts = dummies.Select(x => x.AttachedEntity ?? default).ToArray();
+        var player = pair.Player!.AttachedEntity!.Value;
+        Assert.That(entMan.EntityExists(player));
+        Assert.That(dummyEnts.All(entMan.EntityExists));
+
+        // Make sure the player is a traitor.
+        var mind = mindSys.GetMind(player)!.Value;
+        Assert.That(roleSys.MindIsAntagonist(mind));
+        Assert.That(factionSys.IsMember(player, "Syndicate"), Is.True);
+        Assert.That(factionSys.IsMember(player, "NanoTrasen"), Is.False);
+        Assert.That(traitorRule.TotalTraitors, Is.EqualTo(1));
+        Assert.That(traitorRule.TraitorMinds[0], Is.EqualTo(mind));
+
+        // Check total objective difficulty
+        Assert.That(entMan.TryGetComponent<MindComponent>(mind, out var mindComp));
+        var totalDifficulty = mindComp.Objectives.Sum(o => entMan.GetComponent<ObjectiveComponent>(o).Difficulty);
+        Assert.That(totalDifficulty, Is.AtMost(maxDifficulty),
+            $"MaxDifficulty exceeded! Objectives: {string.Join(", ", mindComp.Objectives.Select(o => FormatObjective(o, entMan)))}");
+        Assert.That(mindComp.Objectives, Is.Not.Empty,
+            $"No objectives assigned!");
+
+
+        await pair.CleanReturnAsync();
+    }
+
+    private static string FormatObjective(Entity<ObjectiveComponent> entity, IEntityManager entMan)
+    {
+        var meta = entMan.GetComponent<MetaDataComponent>(entity);
+        var objective = entMan.GetComponent<ObjectiveComponent>(entity);
+        return $"{meta.EntityName} ({objective.Difficulty})";
+    }
+}

Content.Server/Antag/AntagRandomObjectivesSystem.cs

@@ -39,7 +39,8 @@ public sealed class AntagRandomObjectivesSystem : EntitySystem
 
             for (var pick = 0; pick < set.MaxPicks && ent.Comp.MaxDifficulty > difficulty; pick++)
             {
-                if (_objectives.GetRandomObjective(mindId, mind, set.Groups) is not {} objective)
+                var remainingDifficulty = ent.Comp.MaxDifficulty - difficulty;
+                if (_objectives.GetRandomObjective(mindId, mind, set.Groups, remainingDifficulty) is not { } objective)
                     continue;
 
                 _mind.AddObjective(mindId, mind, objective);

Content.Server/Objectives/ObjectivesSystem.cs

@@ -12,6 +12,7 @@ using Robust.Shared.Random;
 using System.Linq;
 using System.Text;
 using Robust.Server.Player;
+using Robust.Shared.Utility;
 
 namespace Content.Server.Objectives;
 
@@ -180,33 +181,32 @@ public sealed class ObjectivesSystem : SharedObjectivesSystem
         }
     }
 
-    public EntityUid? GetRandomObjective(EntityUid mindId, MindComponent mind, string objectiveGroupProto)
+    public EntityUid? GetRandomObjective(EntityUid mindId, MindComponent mind, ProtoId<WeightedRandomPrototype> objectiveGroupProto, float maxDifficulty)
     {
-        if (!_prototypeManager.TryIndex<WeightedRandomPrototype>(objectiveGroupProto, out var groups))
+        if (!_prototypeManager.TryIndex(objectiveGroupProto, out var groupsProto))
         {
             Log.Error($"Tried to get a random objective, but can't index WeightedRandomPrototype {objectiveGroupProto}");
             return null;
         }
 
-        // TODO replace whatever the fuck this is with a proper objective selection system
-        // yeah the old 'preventing infinite loops' thing wasn't super elegant either and it mislead people on what exactly it did
-        var tries = 0;
-        while (tries < 20)
-        {
-            var groupName = groups.Pick(_random);
+        // Make a copy of the weights so we don't trash the prototype by removing entries
+        var groups = groupsProto.Weights.ShallowClone();
 
+        while (_random.TryPickAndTake(groups, out var groupName))
+        {
             if (!_prototypeManager.TryIndex<WeightedRandomPrototype>(groupName, out var group))
             {
                 Log.Error($"Couldn't index objective group prototype {groupName}");
                 return null;
             }
 
-            var proto = group.Pick(_random);
-            var objective = TryCreateObjective(mindId, mind, proto);
-            if (objective != null)
+            var objectives = group.Weights.ShallowClone();
+            while (_random.TryPickAndTake(objectives, out var objectiveProto))
+            {
+                if (TryCreateObjective((mindId, mind), objectiveProto, out var objective)
+                    && Comp<ObjectiveComponent>(objective.Value).Difficulty <= maxDifficulty)
                     return objective;
-
-            tries++;
+            }
         }
 
         return null;

Content.Shared/Objectives/Systems/SharedObjectivesSystem.cs

@@ -1,5 +1,5 @@
+using System.Diagnostics.CodeAnalysis;
 using Content.Shared.Mind;
-using Content.Shared.Objectives;
 using Content.Shared.Objectives.Components;
 using Robust.Shared.Prototypes;
 using Robust.Shared.Utility;
@@ -40,7 +40,7 @@ public abstract class SharedObjectivesSystem : EntitySystem
         if (comp.Unique)
         {
             var proto = _metaQuery.GetComponent(uid).EntityPrototype?.ID;
-            foreach (var objective in mind.AllObjectives)
+            foreach (var objective in mind.Objectives)
             {
                 if (_metaQuery.GetComponent(objective).EntityPrototype?.ID == proto)
                     return false;
@@ -92,7 +92,18 @@ public abstract class SharedObjectivesSystem : EntitySystem
     }
 
     /// <summary>
-    /// Get the title, description, icon and progress of an objective using <see cref="ObjectiveGetProgressEvent"/>.
+    /// Spawns and assigns an objective for a mind.
+    /// The objective is not added to the mind's objectives, mind system does that in TryAddObjective.
+    /// If the objective could not be assigned the objective is deleted and false is returned.
+    /// </summary>
+    public bool TryCreateObjective(Entity<MindComponent> mind, EntProtoId proto, [NotNullWhen(true)] out EntityUid? objective)
+    {
+        objective = TryCreateObjective(mind.Owner, mind.Comp, proto);
+        return objective != null;
+    }
+
+    /// <summary>
+    /// Get the title, description, icon and progress of an objective using <see cref="ObjectiveGetInfoEvent"/>.
     /// If any of them are null it is logged and null is returned.
     /// </summary>
     /// <param name="uid"/>ID of the condition entity</param>

Content.Shared/Random/Helpers/SharedRandomExtensions.cs

@@ -1,3 +1,4 @@
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using Content.Shared.Dataset;
 using Content.Shared.FixedPoint;
@@ -87,6 +88,26 @@ namespace Content.Shared.Random.Helpers
             throw new InvalidOperationException("Invalid weighted pick");
         }
 
+        public static T PickAndTake<T>(this IRobustRandom random, Dictionary<T, float> weights)
+            where T : notnull
+        {
+            var pick = Pick(random, weights);
+            weights.Remove(pick);
+            return pick;
+        }
+
+        public static bool TryPickAndTake<T>(this IRobustRandom random, Dictionary<T, float> weights, [NotNullWhen(true)] out T? pick)
+            where T : notnull
+        {
+            if (weights.Count == 0)
+            {
+                pick = default;
+                return false;
+            }
+            pick = PickAndTake(random, weights);
+            return true;
+        }
+
         public static (string reagent, FixedPoint2 quantity) Pick(this WeightedRandomFillSolutionPrototype prototype, IRobustRandom? random = null)
         {
             var randomFill = prototype.PickRandomFill(random);