Bind type of EntityTableContext keys to expected values (#46214)

* refactor: EntityTableContext methods now use strongly typed keys, which are bound to type of value you want to extract from context.

* refactor: remove unused attributes

* refactor: extracted EntityTableContext into separate file, fixed xml-doc

Content.IntegrationTests/Tests/EntityTable/EntityTableTest.cs

@@ -300,11 +300,11 @@ public sealed class EntityTableTest : GameTest
         // EntityTableTestEntSelectorWithCost reads cost 10 from the DynamicRuleCostComponent.
 
         // Budget 9 => not enough.
-        var poor = Run(Table("EntityTableTestEntSelectorWithCost"), ctx: new EntityTableContext(new() { ["Budget"] = 9f }));
+        var poor = Run(Table("EntityTableTestEntSelectorWithCost"), ctx: new EntityTableContext { { HasBudgetCondition.BudgetContextKey, 9f } });
         Assert.That(poor, Is.Empty);
 
         // Budget 10 => enough.
-        var rich = Run(Table("EntityTableTestEntSelectorWithCost"), ctx: new EntityTableContext(new() { ["Budget"] = 10f }));
+        var rich = Run(Table("EntityTableTestEntSelectorWithCost"), ctx: new EntityTableContext { { HasBudgetCondition.BudgetContextKey, 10f } });
         Assert.That(rich, Is.EquivalentTo(new[] { new EntProtoId(EntProtoWithCost) }));
     }
 
@@ -313,11 +313,11 @@ public sealed class EntityTableTest : GameTest
     public void RequireAllConditionSemantics()
     {
         // RequireAll = true, one fails => no spawns.
-        var requireAllResult = Run(Table("EntityTableTestEntRequireAll"), ctx: new EntityTableContext(new() { ["Budget"] = 50f }));
+        var requireAllResult = Run(Table("EntityTableTestEntRequireAll"), ctx: new EntityTableContext { { HasBudgetCondition.BudgetContextKey, 50f } });
         Assert.That(requireAllResult, Is.Empty);
 
         // RequireAll = false, one passes => spawns.
-        var requireAnyResult = Run(Table("EntityTableTestEntRequireAny"), ctx: new EntityTableContext(new() { ["Budget"] = 50f }));
+        var requireAnyResult = Run(Table("EntityTableTestEntRequireAny"), ctx: new EntityTableContext { { HasBudgetCondition.BudgetContextKey, 50f } });
         Assert.That(requireAnyResult, Is.EquivalentTo(new[] { new EntProtoId(EntProto1) }));
     }
 
@@ -328,7 +328,7 @@ public sealed class EntityTableTest : GameTest
         // RequireAll = false but every condition fails => no spawns.
         // Regression test: previously the OR-ing was seeded with `true`, so
         // RequireAll = false always passed regardless of the actual results.
-        var result = Run(Table("EntityTableTestEntRequireAnyAllFail"), ctx: new EntityTableContext(new() { ["Budget"] = 50f }));
+        var result = Run(Table("EntityTableTestEntRequireAnyAllFail"), ctx: new EntityTableContext { { HasBudgetCondition.BudgetContextKey, 50f } });
         Assert.That(result, Is.Empty);
     }
 
@@ -337,7 +337,7 @@ public sealed class EntityTableTest : GameTest
     public void DeepComposition_ComplexTree()
     {
         // Ent1's condition fails (Budget 50 < CostOverride 100), so Group only has Ent2.
-        var result = Run(Table("EntityTableTestDeepComposition"), SeededRand(1), new EntityTableContext(new() { ["Budget"] = 50f }));
+        var result = Run(Table("EntityTableTestDeepComposition"), SeededRand(1), new EntityTableContext { { HasBudgetCondition.BudgetContextKey, 50f } });
         Assert.That(result, Is.EqualTo(new[] { new EntProtoId(EntProto2), new EntProtoId(EntProto1) }));
     }
 
@@ -352,12 +352,12 @@ public sealed class EntityTableTest : GameTest
         // EntProto1 already recorded as spawned => blocked.
         var used = new HashSet<EntProtoId> { new(EntProto1) };
         var blocked = Run(Table("EntityTableTestNotRepeating"),
-            ctx: new EntityTableContext(new() { [ExcludeEntitiesFromContextCondition.EntitiesToExclude] = used }));
+            ctx: new EntityTableContext { { ExcludeEntitiesFromContextCondition.EntitiesToExclude, used } });
         Assert.That(blocked, Is.Empty);
 
         // Tracking enabled, but EntProto1 has not been spawned yet => allowed.
         var allowed = Run(Table("EntityTableTestNotRepeating"),
-            ctx: new EntityTableContext(new() { [ExcludeEntitiesFromContextCondition.EntitiesToExclude] = new HashSet<EntProtoId>() }));
+            ctx: new EntityTableContext { { ExcludeEntitiesFromContextCondition.EntitiesToExclude, new HashSet<EntProtoId>()} });
         Assert.That(allowed, Is.EquivalentTo([new EntProtoId(EntProto1)]));
     }
 
@@ -367,7 +367,7 @@ public sealed class EntityTableTest : GameTest
     {
         var used = new HashSet<EntProtoId> { new(EntProto1) };
         var result = Run(Table("EntityTableTestAllNotRepeating"),
-            ctx: new EntityTableContext(new() { [ExcludeEntitiesFromContextCondition.EntitiesToExclude] = used }));
+            ctx: new EntityTableContext { { ExcludeEntitiesFromContextCondition.EntitiesToExclude, used } });
         Assert.That(result, Is.EqualTo(new[] { new EntProtoId(EntProto2) }));
     }
 
@@ -381,7 +381,7 @@ public sealed class EntityTableTest : GameTest
     {
         var used = new HashSet<EntProtoId> { new(EntProto1) };
         var result = Run(Table("EntityTableTestChainNotRepeating"), SeededRand(1),
-            new EntityTableContext(new() { [ExcludeEntitiesFromContextCondition.EntitiesToExclude] = used }));
+            new EntityTableContext { { ExcludeEntitiesFromContextCondition.EntitiesToExclude, used } });
         Assert.That(result, Is.EqualTo(new[] { new EntProtoId(EntProto2) }));
     }
 
@@ -394,10 +394,10 @@ public sealed class EntityTableTest : GameTest
     [RunOnSide(Side.Server)]
     public void NestedSelector_CheckConditions_IncludesNestedTable()
     {
-        var poor = Run(Table("EntityTableTestChainTableWithCost"), ctx: new EntityTableContext(new() { ["Budget"] = 9f }));
+        var poor = Run(Table("EntityTableTestChainTableWithCost"), ctx: new EntityTableContext { { HasBudgetCondition.BudgetContextKey, 9f } });
         Assert.That(poor, Is.Empty);
 
-        var rich = Run(Table("EntityTableTestChainTableWithCost"), ctx: new EntityTableContext(new() { ["Budget"] = 10f }));
+        var rich = Run(Table("EntityTableTestChainTableWithCost"), ctx: new EntityTableContext { { HasBudgetCondition.BudgetContextKey, 10f } });
         Assert.That(rich, Is.EquivalentTo(new[] { new EntProtoId(EntProtoWithCost) }));
     }
 
@@ -409,7 +409,7 @@ public sealed class EntityTableTest : GameTest
     [RunOnSide(Side.Server)]
     public void GroupSelector_ExcludesNestedSelectorWithFailingConditions()
     {
-        var result = Run(Table("EntityTableTestGroupWithCostlyNested"), SeededRand(1), new EntityTableContext(new() { ["Budget"] = 9f }));
+        var result = Run(Table("EntityTableTestGroupWithCostlyNested"), SeededRand(1), new EntityTableContext { { HasBudgetCondition.BudgetContextKey, 9f } });
         Assert.That(result, Is.EqualTo(new[] { new EntProtoId(EntProto2) }));
     }
 
@@ -437,18 +437,18 @@ public sealed class EntityTableTest : GameTest
 
         // Without the injected condition, the UsedSpawns tracking alone has no effect.
         var unconstrained = Run(Table("EntityTableTestEntSelector"),
-            ctx: new EntityTableContext(new() { [ExcludeEntitiesFromContextCondition.EntitiesToExclude] = used }));
+            ctx: new EntityTableContext { { ExcludeEntitiesFromContextCondition.EntitiesToExclude, used } });
         Assert.That(unconstrained, Is.EquivalentTo([new EntProtoId(EntProto1)]));
 
         // Injecting ExcludeEntitiesFromContextCondition gates the selector: EntProto1 is already used => blocked.
-        var ctx = new EntityTableContext(new() { [ExcludeEntitiesFromContextCondition.EntitiesToExclude] = used });
+        var ctx = new EntityTableContext { { ExcludeEntitiesFromContextCondition.EntitiesToExclude, used } };
         ctx.SetData(EntityTableSelector.AdditionalConditionsKey,
             new List<EntityTableCondition> { new ExcludeEntitiesFromContextCondition() });
         var blocked = Run(Table("EntityTableTestEntSelector"), ctx: ctx);
         Assert.That(blocked, Is.Empty);
 
         // With the condition injected but EntProto1 not yet used, the spawn is allowed.
-        var fresh = new EntityTableContext(new() { [ExcludeEntitiesFromContextCondition.EntitiesToExclude] = new HashSet<EntProtoId>() });
+        var fresh = new EntityTableContext { { ExcludeEntitiesFromContextCondition.EntitiesToExclude, new HashSet<EntProtoId>() } };
         fresh.SetData(EntityTableSelector.AdditionalConditionsKey,
             new List<EntityTableCondition> { new ExcludeEntitiesFromContextCondition() });
         var allowed = Run(Table("EntityTableTestEntSelector"), ctx: fresh);
@@ -460,11 +460,11 @@ public sealed class EntityTableTest : GameTest
     public void AdditionalConditions_RemainLocal()
     {
         var used = new HashSet<EntProtoId> { new(EntProto2) };
-        var ctx = new EntityTableContext(new() { [ExcludeEntitiesFromContextCondition.EntitiesToExclude] = used });
+        var ctx = new EntityTableContext { { ExcludeEntitiesFromContextCondition.EntitiesToExclude, used } };
 
         var result = Run(Table("EntityTableTestLocalizedChildConditions"), ctx: ctx);
 
-        ctx.TryGetData<object>(EntityTableSelector.AdditionalConditionsKey, out var empty);
+        ctx.TryGetData(EntityTableSelector.AdditionalConditionsKey, out var empty);
         using (Assert.EnterMultipleScope())
         {
             Assert.That(empty, Is.Null);
@@ -482,7 +482,7 @@ public sealed class EntityTableTest : GameTest
 
         // An empty container succeeds.
         var container = _sContainer.MakeContainer<ContainerSlot>(SSpawn(EntProto1), "containerId");
-        var ctx = new EntityTableContext(new() { [EmptyContainerCondition.ContainerContextKey] = container });
+        var ctx = new EntityTableContext { {EmptyContainerCondition.ContainerContextKey, container } };
 
         result = Run(Table("EntityTableTestContainerCondition"), ctx: ctx);
         Assert.That(result, Is.EquivalentTo(new [] { EntProto1 }));

Content.Server/GameTicking/Rules/DynamicRuleSystem.cs

@@ -76,10 +76,10 @@ public sealed partial class DynamicRuleSystem : GameRuleSystem<DynamicRuleCompon
     private IEnumerable<EntProtoId> GetRuleSpawns(Entity<DynamicRuleComponent> entity)
     {
         UpdateBudget((entity.Owner, entity.Comp));
-        var ctx = new EntityTableContext(new Dictionary<string, object>
+        var ctx = new EntityTableContext
         {
             { HasBudgetCondition.BudgetContextKey, entity.Comp.Budget },
-        });
+        };
 
         return _entityTable.GetSpawns(entity.Comp.Table, ctx: ctx);
     }

Content.Server/Xenoarchaeology/Artifact/XenoArtifactSystem.ProcGen.cs

@@ -250,10 +250,10 @@ public sealed partial class XenoArtifactSystem
         public TriggerPoolData(int requestedSize)
         {
             _usedTriggers = new(requestedSize);
-            Context = new EntityTableContext(new Dictionary<string, object>
+            Context = new EntityTableContext
             {
-                [ExcludeEntitiesFromContextCondition.EntitiesToExclude] = _usedTriggers
-            });
+                { ExcludeEntitiesFromContextCondition.EntitiesToExclude, _usedTriggers }
+            };
         }
 
         public readonly EntityTableContext Context;

Content.Shared/Containers/ContainerFillSystem.cs

@@ -76,12 +76,10 @@ public sealed partial class ContainerFillSystem : EntitySystem
             if (ent.Comp.ContextContainers)
             {
                 // Pass in the container being filled so conditions can look at it
-                ctx = new EntityTableContext(
-                    new Dictionary<string, object>
-                    {
-                        [EmptyContainerCondition.ContainerContextKey] = container,
-                    }
-                );
+                ctx = new EntityTableContext
+                {
+                    { EmptyContainerCondition.ContainerContextKey, container },
+                };
             }
 
             var spawns = _entityTable.GetSpawns(table, ctx: ctx).ToList();

Content.Shared/EntityTable/Conditions/EmptyContainerCondition.cs

@@ -12,7 +12,7 @@ public sealed partial class EmptyContainerCondition : EntityTableCondition
     /// <summary>
     /// Key for <see cref="EntityTableContext"/> to store container that should be checked by this condition.
     /// </summary>
-    public const string ContainerContextKey = "Container";
+    public static readonly EntityTableContextKey<BaseContainer> ContainerContextKey = new("Container");
 
     /// <inheritdoc/>>
     protected override bool EvaluateImplementation(
@@ -22,7 +22,7 @@ public sealed partial class EmptyContainerCondition : EntityTableCondition
         EntityTableContext ctx
     )
     {
-        if (!ctx.TryGetData<BaseContainer>(ContainerContextKey, out var container))
+        if (!ctx.TryGetData(ContainerContextKey, out var container))
             return false;
 
         return container.Count == 0;

Content.Shared/EntityTable/Conditions/ExcludeEntitiesFromContextCondition.cs

@@ -13,7 +13,7 @@ public sealed partial class ExcludeEntitiesFromContextCondition : EntityTableCon
     /// Context key used to track which entity prototypes should not be spawned.
     /// Value should be <see cref="HashSet{EntityUid}"/>.
     /// </summary>
-    public const string EntitiesToExclude = "EntitiesToExclude";
+    public static readonly EntityTableContextKey<HashSet<EntProtoId>> EntitiesToExclude = new("EntitiesToExclude");
 
     /// <inheritdoc/>>
     protected override bool EvaluateImplementation(
@@ -23,7 +23,7 @@ public sealed partial class ExcludeEntitiesFromContextCondition : EntityTableCon
         EntityTableContext ctx
     )
     {
-        if (!ctx.TryGetData<HashSet<EntProtoId>>(EntitiesToExclude, out var used))
+        if (!ctx.TryGetData(EntitiesToExclude, out var used))
             return true;
 
         if (root is not EntSelector entSelector)

Content.Shared/EntityTable/Conditions/HasBudgetCondition.cs

@@ -1,5 +1,4 @@
 using Content.Shared.EntityTable.EntitySelectors;
-using Content.Shared.GameTicking.Rules;
 using Content.Shared.GameTicking.Rules.Components;
 using Robust.Shared.Prototypes;
 
@@ -10,7 +9,7 @@ namespace Content.Shared.EntityTable.Conditions;
 /// </summary>
 public sealed partial class HasBudgetCondition : EntityTableCondition
 {
-    public const string BudgetContextKey = "Budget";
+    public static readonly EntityTableContextKey<float> BudgetContextKey = new("Budget");
 
     /// <summary>
     /// Used for determining the cost for the budget.

Content.Shared/EntityTable/EntitySelectors/EntityTableSelector.cs

@@ -21,7 +21,7 @@ public abstract partial class EntityTableSelector
     /// <summary>
     /// Key for <see cref="EntityTableContext"/>, under which additional scoped conditions should be stored.
     /// </summary>
-    public const string AdditionalConditionsKey = "AdditionalConditions";
+    public static readonly EntityTableContextKey<List<EntityTableCondition>> AdditionalConditionsKey = new("AdditionalConditions");
 
     /// <summary>
     /// The number of times this selector is run
@@ -152,7 +152,7 @@ public abstract partial class EntityTableSelector
     /// </summary>
     private bool TryGetConditions(EntityTableContext ctx, [NotNullWhen(true)] out List<EntityTableCondition>? conditions)
     {
-        var hasAdditionalConditions = ctx.TryGetData<List<EntityTableCondition>>(AdditionalConditionsKey, out var additionalConditions);
+        var hasAdditionalConditions = ctx.TryGetData(AdditionalConditionsKey, out var additionalConditions);
 
         if (Conditions.Count == 0 && !hasAdditionalConditions)
         {

Content.Shared/EntityTable/EntityTableContext.cs

@@ -0,0 +1,75 @@
+using System.Collections;
+using System.Diagnostics.CodeAnalysis;
+using JetBrains.Annotations;
+
+namespace Content.Shared.EntityTable;
+
+/// <summary>
+/// Context used by selectors and conditions to evaluate in generic gamestate information.
+/// </summary>
+public sealed class EntityTableContext : IEnumerable
+{
+    private readonly Dictionary<string, object> _data = new();
+
+    /// <summary>
+    /// Retrieves an arbitrary piece of data from the context based on a provided key.
+    /// </summary>
+    /// <param name="key">A string key that corresponds to the value we are searching for. </param>
+    /// <param name="value">The value we are trying to extract from the context object</param>
+    /// <typeparam name="T">The type of <see cref="value"/> that we are trying to retrieve</typeparam>
+    /// <returns>If <see cref="key"/> has a corresponding value of type <see cref="T"/></returns>
+    [PublicAPI]
+    public bool TryGetData<T>(EntityTableContextKey<T> key, [NotNullWhen(true)] out T? value)
+    {
+        value = default;
+        if (!_data.TryGetValue(key.Key, out var valueData) || valueData is not T castValueData)
+            return false;
+
+        value = castValueData;
+        return true;
+    }
+
+    /// <summary>
+    /// Sets data into context using provided key.
+    /// </summary>
+    [PublicAPI]
+    public void SetData<T>(EntityTableContextKey<T> key, T data) where T : notnull
+    {
+        _data[key.Key] = data;
+    }
+
+    /// <summary>
+    /// Add method duplicates <see cref="SetData{T}"/>, is used for object initializer.
+    /// </summary>
+    /// <exception cref="ArgumentException ">Thrown if there is already a value stored under same key.</exception>
+    [PublicAPI]
+    public void Add<T>(EntityTableContextKey<T> key, T data) where T : notnull
+    {
+        _data.Add(key.Key, data);
+    }
+
+    /// <summary>
+    /// Removes data from the context, if the key exists.
+    /// </summary>
+    [PublicAPI]
+    public void RemoveData<T>(EntityTableContextKey<T> key)
+    {
+        _data.Remove(key.Key);
+    }
+
+    /// <inheritdoc/>
+    public IEnumerator GetEnumerator()
+    {
+        return _data.GetEnumerator();
+    }
+}
+
+/// <summary>
+/// Key for <see cref="EntityTableContext"/>, used to strongly type calls for setting and getting entities,
+/// should usually be static readonly fields as anchors.
+/// </summary>
+/// <typeparam name="T">Type of value to store/extract to/from context.</typeparam>
+/// <param name="Key">
+/// String value under which data will be stored inside. They should not have collisions, otherwise data will be rewritten.
+/// </param>
+public readonly record struct EntityTableContextKey<T>(string Key);

Content.Shared/EntityTable/EntityTableSystem.cs

@@ -1,6 +1,4 @@
-using System.Diagnostics.CodeAnalysis;
 using Content.Shared.EntityTable.EntitySelectors;
-using JetBrains.Annotations;
 using Robust.Shared.Prototypes;
 using Robust.Shared.Random;
 using Robust.Shared.Utility;
@@ -100,57 +98,3 @@ public sealed partial class EntityTableSystem : EntitySystem
         return table.AverageSpawns(EntityManager, ProtoMan, ctx);
     }
 }
-
-/// <summary>
-/// Context used by selectors and conditions to evaluate in generic gamestate information.
-/// </summary>
-public sealed class EntityTableContext
-{
-    private readonly Dictionary<string, object> _data = new();
-
-    public EntityTableContext()
-    {
-
-    }
-
-    public EntityTableContext(Dictionary<string, object> data)
-    {
-        _data = data;
-    }
-
-    /// <summary>
-    /// Retrieves an arbitrary piece of data from the context based on a provided key.
-    /// </summary>
-    /// <param name="key">A string key that corresponds to the value we are searching for. </param>
-    /// <param name="value">The value we are trying to extract from the context object</param>
-    /// <typeparam name="T">The type of <see cref="value"/> that we are trying to retrieve</typeparam>
-    /// <returns>If <see cref="key"/> has a corresponding value of type <see cref="T"/></returns>
-    [PublicAPI]
-    public bool TryGetData<T>([ForbidLiteral] string key, [NotNullWhen(true)] out T? value)
-    {
-        value = default;
-        if (!_data.TryGetValue(key, out var valueData) || valueData is not T castValueData)
-            return false;
-
-        value = castValueData;
-        return true;
-    }
-
-    /// <summary>
-    /// Sets data into context using provided key.
-    /// </summary>
-    [PublicAPI]
-    public void SetData<T>([ForbidLiteral] string key, T data) where T : notnull
-    {
-        _data[key] = data;
-    }
-
-    /// <summary>
-    /// Removes data from the context, if the key exists.
-    /// </summary>
-    [PublicAPI]
-    public void RemoveData([ForbidLiteral] string key)
-    {
-        _data.Remove(key);
-    }
-}