Add condition support to entity tables (#36819)

Content.Shared/EntityTable/Conditions/EntityTableCondition.cs

@@ -0,0 +1,28 @@
+using Content.Shared.EntityTable.EntitySelectors;
+using JetBrains.Annotations;
+using Robust.Shared.Prototypes;
+
+namespace Content.Shared.EntityTable.Conditions;
+
+/// <summary>
+/// Used for implementing conditional logic for <see cref="EntityTableSelector"/>.
+/// </summary>
+[ImplicitDataDefinitionForInheritors, UsedImplicitly(ImplicitUseTargetFlags.WithInheritors)]
+public abstract partial class EntityTableCondition
+{
+    /// <summary>
+    /// If true, inverts the result of the condition.
+    /// </summary>
+    [DataField]
+    public bool Invert;
+
+    public bool Evaluate(IEntityManager entMan, IPrototypeManager proto)
+    {
+        var res = EvaluateImplementation(entMan, proto);
+
+        // XOR eval to invert the result.
+        return res ^ Invert;
+    }
+
+    public abstract bool EvaluateImplementation(IEntityManager entMan, IPrototypeManager proto);
+}

Content.Shared/EntityTable/Conditions/PlayerCountCondition.cs

@@ -0,0 +1,34 @@
+using Robust.Shared.Player;
+using Robust.Shared.Prototypes;
+
+namespace Content.Shared.EntityTable.Conditions;
+
+/// <summary>
+/// Condition that passes only if the server player count is within a certain range.
+/// </summary>
+public sealed partial class PlayerCountCondition : EntityTableCondition
+{
+    /// <summary>
+    /// Minimum players of needed for this condition to succeed. Inclusive.
+    /// </summary>
+    [DataField]
+    public int Min = int.MinValue;
+
+    /// <summary>
+    /// Maximum numbers of players there can be for this condition to succeed. Inclusive.
+    /// </summary>
+    [DataField]
+    public int Max = int.MaxValue;
+
+    private static ISharedPlayerManager? _playerManager;
+
+    public override bool EvaluateImplementation(IEntityManager entMan, IPrototypeManager proto)
+    {
+        // Don't resolve this repeatedly
+        _playerManager ??= IoCManager.Resolve<ISharedPlayerManager>();
+
+        var playerCount = _playerManager.PlayerCount;
+
+        return playerCount >= Min && playerCount <= Max;
+    }
+}

Content.Shared/EntityTable/EntitySelectors/EntityTableSelector.cs

@@ -1,3 +1,4 @@
+using Content.Shared.EntityTable.Conditions;
 using Content.Shared.EntityTable.ValueSelector;
 using JetBrains.Annotations;
 using Robust.Shared.Prototypes;
@@ -26,10 +27,26 @@ public abstract partial class EntityTableSelector
     [DataField]
     public double Prob = 1;
 
+    /// <summary>
+    /// A list of conditions that must evaluate to 'true' for the selector to apply.
+    /// </summary>
+    [DataField]
+    public List<EntityTableCondition> Conditions = new();
+
+    /// <summary>
+    /// If true, all the conditions must be successful in order for the selector to process.
+    /// Otherwise, only one of them must be.
+    /// </summary>
+    [DataField]
+    public bool RequireAll = true;
+
     public IEnumerable<EntProtoId> GetSpawns(System.Random rand,
         IEntityManager entMan,
         IPrototypeManager proto)
     {
+        if (!CheckConditions(entMan, proto))
+            yield break;
+
         var rolls = Rolls.Get(rand);
         for (var i = 0; i < rolls; i++)
         {
@@ -43,6 +60,28 @@ public abstract partial class EntityTableSelector
         }
     }
 
+    public bool CheckConditions(IEntityManager entMan, IPrototypeManager proto)
+    {
+        if (Conditions.Count == 0)
+            return true;
+
+        var success = false;
+        foreach (var condition in Conditions)
+        {
+            var res = condition.Evaluate(entMan, proto);
+
+            if (RequireAll && !res)
+                return false; // intentional break out of loop and function
+
+            success |= res;
+        }
+
+        if (RequireAll)
+            return true;
+
+        return success;
+    }
+
     protected abstract IEnumerable<EntProtoId> GetSpawnsImplementation(System.Random rand,
         IEntityManager entMan,
         IPrototypeManager proto);

Content.Shared/EntityTable/EntitySelectors/GroupSelector.cs

@@ -18,6 +18,10 @@ public sealed partial class GroupSelector : EntityTableSelector
         var children = new Dictionary<EntityTableSelector, float>(Children.Count);
         foreach (var child in Children)
         {
+            // Don't include invalid groups
+            if (!child.CheckConditions(entMan, proto))
+                continue;
+
             children.Add(child, child.Weight);
         }
 

Resources/Prototypes/Catalog/Fills/Lockers/heads.yml

@@ -51,6 +51,10 @@
     - id: SpaceCash1000
     - id: WeaponDisabler
     - id: ClothingEyesGlassesCommand
+    - id: HeadSkeleton # A skull to accompany your skeleton crew
+      conditions:
+      - !type:PlayerCountCondition
+        max: 15
 
 # No laser table + Laser table
 - type: entityTable