Add support for contextual information in EntityTables (#37737)

* Add context support for entityTables

* fix build fail

* comments

* Update Content.Shared/EntityTable/EntityTableSystem.cs

---------

Co-authored-by: slarticodefast <161409025+slarticodefast@users.noreply.github.com>

Content.Server/Procedural/DungeonJob/DungeonJob.CornerClutter.cs

@@ -42,7 +42,7 @@ public sealed partial class DungeonJob
                 if (random.Prob(gen.Chance))
                 {
                     var coords = _maps.GridTileToLocal(_gridUid, _grid, tile);
-                    var protos = contentsTable.Table.GetSpawns(random, _entManager, _prototype);
+                    var protos = _entTable.GetSpawns(contentsTable, random);
                     _entManager.SpawnEntitiesAttachedTo(coords, protos);
                 }
 

Content.Shared/EntityTable/Conditions/EntityTableCondition.cs

@@ -16,13 +16,13 @@ public abstract partial class EntityTableCondition
     [DataField]
     public bool Invert;
 
-    public bool Evaluate(IEntityManager entMan, IPrototypeManager proto)
+    public bool Evaluate(EntityTableSelector root, IEntityManager entMan, IPrototypeManager proto, EntityTableContext ctx)
     {
-        var res = EvaluateImplementation(entMan, proto);
+        var res = EvaluateImplementation(root, entMan, proto, ctx);
 
         // XOR eval to invert the result.
         return res ^ Invert;
     }
 
-    public abstract bool EvaluateImplementation(IEntityManager entMan, IPrototypeManager proto);
+    protected abstract bool EvaluateImplementation(EntityTableSelector root, IEntityManager entMan, IPrototypeManager proto, EntityTableContext ctx);
 }

Content.Shared/EntityTable/Conditions/PlayerCountCondition.cs

@@ -1,3 +1,4 @@
+using Content.Shared.EntityTable.EntitySelectors;
 using Robust.Shared.Player;
 using Robust.Shared.Prototypes;
 
@@ -22,7 +23,7 @@ public sealed partial class PlayerCountCondition : EntityTableCondition
 
     private static ISharedPlayerManager? _playerManager;
 
-    public override bool EvaluateImplementation(IEntityManager entMan, IPrototypeManager proto)
+    protected override bool EvaluateImplementation(EntityTableSelector root, IEntityManager entMan, IPrototypeManager proto, EntityTableContext ctx)
     {
         // Don't resolve this repeatedly
         _playerManager ??= IoCManager.Resolve<ISharedPlayerManager>();

Content.Shared/EntityTable/EntitySelectors/AllSelector.cs

@@ -12,11 +12,12 @@ public sealed partial class AllSelector : EntityTableSelector
 
     protected override IEnumerable<EntProtoId> GetSpawnsImplementation(System.Random rand,
         IEntityManager entMan,
-        IPrototypeManager proto)
+        IPrototypeManager proto,
+        EntityTableContext ctx)
     {
         foreach (var child in Children)
         {
-            foreach (var spawn in child.GetSpawns(rand, entMan, proto))
+            foreach (var spawn in child.GetSpawns(rand, entMan, proto, ctx))
             {
                 yield return spawn;
             }

Content.Shared/EntityTable/EntitySelectors/EntSelector.cs

@@ -18,7 +18,8 @@ public sealed partial class EntSelector : EntityTableSelector
 
     protected override IEnumerable<EntProtoId> GetSpawnsImplementation(System.Random rand,
         IEntityManager entMan,
-        IPrototypeManager proto)
+        IPrototypeManager proto,
+        EntityTableContext ctx)
     {
         var num = Amount.Get(rand);
         for (var i = 0; i < num; i++)

Content.Shared/EntityTable/EntitySelectors/EntityTableSelector.cs

@@ -42,9 +42,10 @@ public abstract partial class EntityTableSelector
 
     public IEnumerable<EntProtoId> GetSpawns(System.Random rand,
         IEntityManager entMan,
-        IPrototypeManager proto)
+        IPrototypeManager proto,
+        EntityTableContext ctx)
     {
-        if (!CheckConditions(entMan, proto))
+        if (!CheckConditions(entMan, proto, ctx))
             yield break;
 
         var rolls = Rolls.Get(rand);
@@ -53,14 +54,14 @@ public abstract partial class EntityTableSelector
             if (!rand.Prob(Prob))
                 continue;
 
-            foreach (var spawn in GetSpawnsImplementation(rand, entMan, proto))
+            foreach (var spawn in GetSpawnsImplementation(rand, entMan, proto, ctx))
             {
                 yield return spawn;
             }
         }
     }
 
-    public bool CheckConditions(IEntityManager entMan, IPrototypeManager proto)
+    public bool CheckConditions(IEntityManager entMan, IPrototypeManager proto, EntityTableContext ctx)
     {
         if (Conditions.Count == 0)
             return true;
@@ -68,7 +69,7 @@ public abstract partial class EntityTableSelector
         var success = false;
         foreach (var condition in Conditions)
         {
-            var res = condition.Evaluate(entMan, proto);
+            var res = condition.Evaluate(this, entMan, proto, ctx);
 
             if (RequireAll && !res)
                 return false; // intentional break out of loop and function
@@ -84,5 +85,6 @@ public abstract partial class EntityTableSelector
 
     protected abstract IEnumerable<EntProtoId> GetSpawnsImplementation(System.Random rand,
         IEntityManager entMan,
-        IPrototypeManager proto);
+        IPrototypeManager proto,
+        EntityTableContext ctx);
 }

Content.Shared/EntityTable/EntitySelectors/GroupSelector.cs

@@ -13,13 +13,14 @@ public sealed partial class GroupSelector : EntityTableSelector
 
     protected override IEnumerable<EntProtoId> GetSpawnsImplementation(System.Random rand,
         IEntityManager entMan,
-        IPrototypeManager proto)
+        IPrototypeManager proto,
+        EntityTableContext ctx)
     {
         var children = new Dictionary<EntityTableSelector, float>(Children.Count);
         foreach (var child in Children)
         {
             // Don't include invalid groups
-            if (!child.CheckConditions(entMan, proto))
+            if (!child.CheckConditions(entMan, proto, ctx))
                 continue;
 
             children.Add(child, child.Weight);
@@ -27,6 +28,6 @@ public sealed partial class GroupSelector : EntityTableSelector
 
         var pick = SharedRandomExtensions.Pick(children, rand);
 
-        return pick.GetSpawns(rand, entMan, proto);
+        return pick.GetSpawns(rand, entMan, proto, ctx);
     }
 }

Content.Shared/EntityTable/EntitySelectors/NestedSelector.cs

@@ -13,8 +13,9 @@ public sealed partial class NestedSelector : EntityTableSelector
 
     protected override IEnumerable<EntProtoId> GetSpawnsImplementation(System.Random rand,
         IEntityManager entMan,
-        IPrototypeManager proto)
+        IPrototypeManager proto,
+        EntityTableContext ctx)
     {
-        return proto.Index(TableId).Table.GetSpawns(rand, entMan, proto);
+        return proto.Index(TableId).Table.GetSpawns(rand, entMan, proto, ctx);
     }
 }

Content.Shared/EntityTable/EntitySelectors/NoneSelector.cs

@@ -9,7 +9,8 @@ public sealed partial class NoneSelector : EntityTableSelector
 {
     protected override IEnumerable<EntProtoId> GetSpawnsImplementation(System.Random rand,
         IEntityManager entMan,
-        IPrototypeManager proto)
+        IPrototypeManager proto,
+        EntityTableContext ctx)
     {
         yield break;
     }

Content.Shared/EntityTable/EntityTableSystem.cs

@@ -1,4 +1,6 @@
+using System.Diagnostics.CodeAnalysis;
 using Content.Shared.EntityTable.EntitySelectors;
+using JetBrains.Annotations;
 using Robust.Shared.Prototypes;
 using Robust.Shared.Random;
 
@@ -9,18 +11,55 @@ public sealed class EntityTableSystem : EntitySystem
     [Dependency] private readonly IPrototypeManager _prototypeManager = default!;
     [Dependency] private readonly IRobustRandom _random = default!;
 
-    public IEnumerable<EntProtoId> GetSpawns(EntityTablePrototype entTableProto, System.Random? rand = null)
+    public IEnumerable<EntProtoId> GetSpawns(EntityTablePrototype entTableProto, System.Random? rand = null, EntityTableContext? ctx = null)
     {
         // convenient
-        return GetSpawns(entTableProto.Table, rand);
+        return GetSpawns(entTableProto.Table, rand, ctx);
     }
 
-    public IEnumerable<EntProtoId> GetSpawns(EntityTableSelector? table, System.Random? rand = null)
+    public IEnumerable<EntProtoId> GetSpawns(EntityTableSelector? table, System.Random? rand = null, EntityTableContext? ctx = null)
     {
         if (table == null)
             return new List<EntProtoId>();
 
         rand ??= _random.GetRandom();
-        return table.GetSpawns(rand, EntityManager, _prototypeManager);
+        ctx ??= new EntityTableContext();
+        return table.GetSpawns(rand, EntityManager, _prototypeManager, 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;
     }
 }