EntityManager Singleton methods

RELEASE-NOTES.md

@@ -39,7 +39,9 @@ END TEMPLATE-->
 
 ### New features
 
-*None yet*
+- `IEntityManager` has a new family of 1-4 component methods for working with *singleton entities*, entities which only
+  one instance of exists at a time for state management. Consult `IEntityManager.Single.cs` and the documentation for
+  details.
 
 ### Bugfixes
 

Robust.Shared.IntegrationTests/GameObjects/EntityManagerSingletonTests.cs

@@ -0,0 +1,112 @@
+using NUnit.Framework;
+using Robust.Shared.GameObjects;
+using Robust.Shared.IoC;
+using Robust.Shared.Prototypes;
+using Robust.Shared.Serialization.Manager;
+
+namespace Robust.UnitTesting.Shared.GameObjects;
+
+[TestOf(typeof(EntityManager))]
+internal sealed class EntityManagerSingletonTests : OurRobustUnitTest
+{
+    private const string TestSingleton1 = "T_TestSingleton1";
+
+    private const string Prototypes = $"""
+        - type: entity
+          id: {TestSingleton1}
+          components:
+          - type: Marker1
+          - type: Marker2
+          - type: Marker3
+          - type: Marker4
+        """;
+
+    protected override Type[]? ExtraComponents =>
+    [
+        typeof(Marker1Component), typeof(Marker2Component), typeof(Marker3Component), typeof(Marker4Component)
+    ];
+
+    private PrototypeManager _protoMan = default!;
+    private IEntityManager _entMan = default!;
+
+    [OneTimeSetUp]
+    public void Setup()
+    {
+        IoCManager.Resolve<ISerializationManager>().Initialize();
+
+        _protoMan = (PrototypeManager) IoCManager.Resolve<IPrototypeManager>();
+        _protoMan.RegisterKind(typeof(EntityPrototype), typeof(EntityCategoryPrototype));
+        _protoMan.LoadString(Prototypes);
+        _protoMan.ResolveResults();
+
+        _entMan = IoCManager.Resolve<IEntityManager>();
+    }
+
+    [Test]
+    [Description("""
+    Creates an entity with marker components 1-4, and ensures increasingly constrained Single<>() queries match it.
+    Then creates a second entity, and ensures queries fail.
+    Then deletes both, and ensures Single fails while TrySingle is fine.
+    """)]
+    public void SingleEntity()
+    {
+        var mySingle = _entMan.Spawn(TestSingleton1);
+
+        var single1 = _entMan.Single<Marker1Component>();
+        var single2 = _entMan.Single<Marker1Component, Marker2Component>();
+        var single3 = _entMan.Single<Marker1Component, Marker2Component, Marker3Component>();
+        var single4 = _entMan.Single<Marker1Component, Marker2Component, Marker3Component, Marker4Component>();
+
+        using (Assert.EnterMultipleScope())
+        {
+            Assert.That(mySingle, NUnit.Framework.Is.EqualTo(single1.Owner));
+            Assert.That(mySingle, NUnit.Framework.Is.EqualTo(single2.Owner));
+            Assert.That(mySingle, NUnit.Framework.Is.EqualTo(single3.Owner));
+            Assert.That(mySingle, NUnit.Framework.Is.EqualTo(single4.Owner));
+        }
+
+        var bonusSingle = _entMan.Spawn(TestSingleton1);
+
+        using (Assert.EnterMultipleScope())
+        {
+            Assert.Throws<NonUniqueSingletonException>(() => _entMan.Single<Marker1Component>());
+            Assert.Throws<NonUniqueSingletonException>(() => _entMan.Single<Marker1Component, Marker2Component>());
+            Assert.Throws<NonUniqueSingletonException>(() =>
+                _entMan.Single<Marker1Component, Marker2Component, Marker3Component>());
+            Assert.Throws<NonUniqueSingletonException>(() =>
+                _entMan.Single<Marker1Component, Marker2Component, Marker3Component, Marker4Component>());
+            Assert.Throws<NonUniqueSingletonException>(() => _entMan.TrySingle<Marker1Component>(out _));
+            Assert.Throws<NonUniqueSingletonException>(() =>
+                _entMan.TrySingle<Marker1Component, Marker2Component>(out _));
+            Assert.Throws<NonUniqueSingletonException>(() =>
+                _entMan.TrySingle<Marker1Component, Marker2Component, Marker3Component>(out _));
+            Assert.Throws<NonUniqueSingletonException>(() =>
+                _entMan.TrySingle<Marker1Component, Marker2Component, Marker3Component, Marker4Component>(out _));
+        }
+
+        _entMan.DeleteEntity(bonusSingle);
+        _entMan.DeleteEntity(mySingle);
+
+        using (Assert.EnterMultipleScope())
+        {
+            Assert.Throws<MatchNotFoundException>(() => _entMan.Single<Marker1Component>());
+            Assert.Throws<MatchNotFoundException>(() => _entMan.Single<Marker1Component, Marker2Component>());
+            Assert.Throws<MatchNotFoundException>(() =>
+                _entMan.Single<Marker1Component, Marker2Component, Marker3Component>());
+            Assert.Throws<MatchNotFoundException>(() =>
+                _entMan.Single<Marker1Component, Marker2Component, Marker3Component, Marker4Component>());
+            Assert.That(_entMan.TrySingle<Marker1Component>(out _), NUnit.Framework.Is.False);
+            Assert.That(_entMan.TrySingle<Marker1Component, Marker2Component>(out _), NUnit.Framework.Is.False);
+            Assert.That(_entMan.TrySingle<Marker1Component, Marker2Component, Marker3Component>(out _),
+                NUnit.Framework.Is.False);
+            Assert.That(
+                _entMan.TrySingle<Marker1Component, Marker2Component, Marker3Component, Marker4Component>(out _),
+                NUnit.Framework.Is.False);
+        }
+    }
+}
+
+internal sealed partial class Marker1Component : Component;
+internal sealed partial class Marker2Component : Component;
+internal sealed partial class Marker3Component : Component;
+internal sealed partial class Marker4Component : Component;

Robust.Shared/GameObjects/Docs.xml

@@ -11,4 +11,23 @@
         This is also preferable if you may have already looked up the component, saving on lookup time.
     </remarks>
   </entry>
+  <entry name="SingleOrSpawnRemark">
+    <remarks>
+      <para>
+          This does not return the entity it spawns, it tries to look it up, so if spawning that entity violates
+          singleton status (or it lacks the necessary component) this will throw immediately after instead of later
+          on next call.
+
+          This will also still throw <see cref="MatchNotFoundException"/> if the spawned entity doesn't match at all.
+      </para>
+    </remarks>
+  </entry>
+  <entry name="SingletonPauseRemark">
+    <remarks>
+      <para>
+        This API does not obey <see cref="MetaDataComponent.EntityPaused"/>.
+        If you need to "pause" singleton entities, use an extra marker component that you remove when "pausing" them.
+      </para>
+    </remarks>
+  </entry>
 </entries>

Robust.Shared/GameObjects/EntityManager.Single.cs

@@ -0,0 +1,239 @@
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Linq;
+using Robust.Shared.Map;
+using Robust.Shared.Prototypes;
+using Robust.Shared.Utility;
+
+namespace Robust.Shared.GameObjects;
+
+public abstract partial class EntityManager
+{
+    public Entity<TComp1> Single<TComp1>()
+        where TComp1: IComponent
+    {
+        var index = _entTraitArray[CompIdx.ArrayIndex<TComp1>()];
+
+        if (index.Keys.FirstOrNull() is { } ent && index.Count == 1)
+        {
+            return new Entity<TComp1>(ent, (TComp1)index[ent]);
+        }
+
+        if (index.Count > 1)
+        {
+            throw new NonUniqueSingletonException(index.Keys.ToArray(), typeof(TComp1));
+        }
+        else
+        {
+            // 0.
+            throw new MatchNotFoundException(typeof(TComp1));
+        }
+    }
+
+    public Entity<TComp1, TComp2> Single<TComp1, TComp2>()
+        where TComp1: IComponent
+        where TComp2: IComponent
+    {
+        var query = AllEntityQueryEnumerator<TComp1, TComp2>();
+
+        if (!query.MoveNext(out var ent, out var comp1, out var comp2))
+            throw new MatchNotFoundException(typeof(TComp1), typeof(TComp2));
+
+        if (query.MoveNext(out var ent2, out _, out _))
+        {
+            var list = new List<EntityUid> { ent, ent2 };
+
+            while (query.MoveNext(out var ent3, out _, out _))
+            {
+                list.Add(ent3);
+            }
+
+            throw new NonUniqueSingletonException(list.ToArray(), typeof(TComp1), typeof(TComp2));
+        }
+
+        return new(ent, comp1, comp2);
+    }
+
+    public Entity<TComp1, TComp2, TComp3> Single<TComp1, TComp2, TComp3>()
+        where TComp1: IComponent
+        where TComp2: IComponent
+        where TComp3: IComponent
+    {
+        var query = AllEntityQueryEnumerator<TComp1, TComp2, TComp3>();
+
+        if (!query.MoveNext(out var ent, out var comp1, out var comp2, out var comp3))
+            throw new MatchNotFoundException(typeof(TComp1), typeof(TComp2), typeof(TComp3));
+
+        if (query.MoveNext(out var ent2, out _, out _, out _))
+        {
+            var list = new List<EntityUid> { ent, ent2 };
+
+            while (query.MoveNext(out var ent3, out _, out _, out _))
+            {
+                list.Add(ent3);
+            }
+
+            throw new NonUniqueSingletonException(list.ToArray(), typeof(TComp1), typeof(TComp2), typeof(TComp3));
+        }
+
+        return new(ent, comp1, comp2, comp3);
+    }
+
+    public Entity<TComp1, TComp2, TComp3, TComp4> Single<TComp1, TComp2, TComp3, TComp4>()
+        where TComp1: IComponent
+        where TComp2: IComponent
+        where TComp3: IComponent
+        where TComp4: IComponent
+    {
+        var query = AllEntityQueryEnumerator<TComp1, TComp2, TComp3, TComp4>();
+
+        if (!query.MoveNext(out var ent, out var comp1, out var comp2, out var comp3, out var comp4))
+            throw new MatchNotFoundException(typeof(TComp1), typeof(TComp2), typeof(TComp3), typeof(TComp4));
+
+        if (query.MoveNext(out var ent2, out _, out _, out _, out _))
+        {
+            var list = new List<EntityUid> { ent, ent2 };
+
+            while (query.MoveNext(out var ent3, out _, out _, out _, out _))
+            {
+                list.Add(ent3);
+            }
+
+            throw new NonUniqueSingletonException(list.ToArray(), typeof(TComp1), typeof(TComp2), typeof(TComp3), typeof(TComp4));
+        }
+
+        return new(ent, comp1, comp2, comp3, comp4);
+    }
+
+    public bool TrySingle<TComp1>([NotNullWhen(true)] out Entity<TComp1>? entity)
+        where TComp1: IComponent
+    {
+        var index = _entTraitArray[CompIdx.ArrayIndex<TComp1>()];
+
+        if (index.Keys.FirstOrNull() is { } ent && index.Count == 1)
+        {
+            entity = new Entity<TComp1>(ent, (TComp1)index[ent]);
+            return true;
+        }
+        else
+        {
+            entity = null;
+        }
+
+        if (index.Count > 1)
+        {
+            throw new NonUniqueSingletonException(index.Keys.ToArray(), typeof(TComp1));
+        }
+
+        return false;
+    }
+
+    public bool TrySingle<TComp1, TComp2>([NotNullWhen(true)] out Entity<TComp1, TComp2>? entity)
+        where TComp1: IComponent
+        where TComp2: IComponent
+    {
+        var query = AllEntityQueryEnumerator<TComp1, TComp2>();
+
+        if (!query.MoveNext(out var ent, out var comp1, out var comp2))
+        {
+            entity = null;
+            return false;
+        }
+
+        if (query.MoveNext(out var ent2, out _, out _))
+        {
+            var list = new List<EntityUid> { ent, ent2 };
+
+            while (query.MoveNext(out var ent3, out _, out _))
+            {
+                list.Add(ent3);
+            }
+
+            throw new NonUniqueSingletonException(list.ToArray(), typeof(TComp1), typeof(TComp2));
+        }
+
+        entity = new(ent, comp1, comp2);
+        return true;
+    }
+
+    public bool TrySingle<TComp1, TComp2, TComp3>([NotNullWhen(true)] out Entity<TComp1, TComp2, TComp3>? entity)
+        where TComp1: IComponent
+        where TComp2: IComponent
+        where TComp3: IComponent
+    {
+        var query = AllEntityQueryEnumerator<TComp1, TComp2, TComp3>();
+
+        if (!query.MoveNext(out var ent, out var comp1, out var comp2, out var comp3))
+        {
+            entity = null;
+            return false;
+        }
+
+        if (query.MoveNext(out var ent2, out _, out _, out _))
+        {
+            var list = new List<EntityUid> { ent, ent2 };
+
+            while (query.MoveNext(out var ent3, out _, out _, out _))
+            {
+                list.Add(ent3);
+            }
+
+            throw new NonUniqueSingletonException(list.ToArray(), typeof(TComp1), typeof(TComp2), typeof(TComp3));
+        }
+
+        entity = new(ent, comp1, comp2, comp3);
+        return true;
+    }
+
+    public bool TrySingle<TComp1, TComp2, TComp3, TComp4>([NotNullWhen(true)] out Entity<TComp1, TComp2, TComp3, TComp4>? entity)
+        where TComp1: IComponent
+        where TComp2: IComponent
+        where TComp3: IComponent
+        where TComp4: IComponent
+    {
+        var query = AllEntityQueryEnumerator<TComp1, TComp2, TComp3, TComp4>();
+
+        if (!query.MoveNext(out var ent, out var comp1, out var comp2, out var comp3, out var comp4))
+        {
+            entity = null;
+            return false;
+        }
+
+        if (query.MoveNext(out var ent2, out _, out _, out _, out _))
+        {
+            var list = new List<EntityUid> { ent, ent2 };
+
+            while (query.MoveNext(out var ent3, out _, out _, out _, out _))
+            {
+                list.Add(ent3);
+            }
+
+            throw new NonUniqueSingletonException(list.ToArray(), typeof(TComp1), typeof(TComp2), typeof(TComp3), typeof(TComp4));
+        }
+
+        entity = new(ent, comp1, comp2, comp3, comp4);
+        return true;
+    }
+}
+
+/// <summary>
+///     Exception for when <see cref="float"/> and co cannot find a unique match.
+/// </summary>
+/// <param name="matches">The set of matching entities.</param>
+/// <param name="components">The set of components you tried to match over.</param>
+public sealed class NonUniqueSingletonException(EntityUid[] matches, params Type[] components) : Exception
+{
+    public override string Message =>
+        $"Expected precisely one entity to match the component set {string.Join(", ", components)}, but found {matches.Length}: {string.Join(", ", matches)}";
+}
+
+/// <summary>
+///     Exception for when <see cref="float"/> and co cannot find any match.
+/// </summary>
+/// <param name="components">The set of components you tried to match over.</param>
+public sealed class MatchNotFoundException(params Type[] components) : Exception
+{
+    public override string Message =>
+        $"Expected precisely one entity to match the component set {string.Join(", ", components)}, but found none.";
+}