Auto-generate prototype variants at runtime (space-wizards#6692)

Author: chromiumboy

Robust.Shared.IntegrationTests/Prototypes/PrototypeVariantizationTest.cs

@@ -0,0 +1,169 @@
+using JetBrains.Annotations;
+using NUnit.Framework;
+using Robust.Shared.GameObjects;
+using Robust.Shared.IoC;
+using Robust.Shared.Prototypes;
+using Robust.Shared.Serialization.Manager;
+using Robust.Shared.Serialization.Manager.Attributes;
+
+namespace Robust.UnitTesting.Shared.Prototypes;
+
+[UsedImplicitly]
+[TestFixture]
+internal sealed partial class PrototypeVariantizationTest : OurRobustUnitTest
+{
+    private const string TestProtoId = "TestPrototype";
+    private const string TestProtoVariantAId = "TestPrototypeVariantA";
+    private const string TestProtoVariantBId = "TestPrototypeVariantB";
+
+    private IPrototypeManager protoManager = default!;
+
+    protected override Type[] ExtraComponents => new[] { typeof(PrototypeVariantizationTestComponent) };
+
+    [OneTimeSetUp]
+    public void Setup()
+    {
+        IoCManager.Resolve<ISerializationManager>().Initialize();
+        protoManager = IoCManager.Resolve<IPrototypeManager>();
+        protoManager.Initialize();
+        protoManager.LoadString(DOCUMENT, changed: new());
+        protoManager.ResolveResults();
+    }
+
+    /// <summary>
+    /// Tests that the prototypes defined in the test YAML document, as well as their expected variants,
+    /// are properly generated and can be resolved from the prototype manager.
+    /// </summary>
+    [Test]
+    public void TestPrototypesExist()
+    {
+        Assert.Multiple(() =>
+        {
+            Assert.That(protoManager.Resolve<EntityPrototype>(TestProtoId, out var _));
+            Assert.That(protoManager.Resolve<EntityPrototype>(TestProtoVariantAId, out _));
+            Assert.That(protoManager.Resolve<EntityPrototype>(TestProtoVariantBId, out _));
+        });
+    }
+
+    /// <summary>
+    /// Tests that the value modifications defined in the test YAML document are properly applied to all variants.
+    /// </summary>
+    [Test]
+    public void TestValueModification()
+    {
+        // Original
+        Assert.Multiple(() =>
+        {
+            var testProto = protoManager.Index<EntityPrototype>(TestProtoId);
+            Assert.That(testProto.Components, Contains.Key("PrototypeVariantizationTest"));
+
+            var comp = testProto.Components["PrototypeVariantizationTest"].Component as PrototypeVariantizationTestComponent;
+            Assert.That(comp, Is.Not.Null);
+            Assert.That(comp!.NumValue, Is.EqualTo(1));
+            Assert.That(comp!.EnumValue, Is.EqualTo(PrototypeVariantizationTestEnum.First));
+            Assert.That(comp!.StringValue, Is.EqualTo("string-1"));
+            Assert.That(comp!.StringArray.SequenceEqual(["el-1"]));
+
+            Assert.That(comp!.StringDict.TryGetValue("key1", out var value1));
+            Assert.That(value1, Is.EqualTo("string-1a"));
+            Assert.That(comp!.StringDict.TryGetValue("key2", out var value2));
+            Assert.That(value2, Is.EqualTo("string-1b"));
+
+            Assert.That(comp!.RecursiveRecord?.Child?.Child?.Value, Is.EqualTo("child-1"));
+        });
+
+        // First variant
+        Assert.Multiple(() =>
+        {
+            var testProtoVariantA = protoManager.Index<EntityPrototype>(TestProtoVariantAId);
+            Assert.That(testProtoVariantA.Components, Contains.Key("PrototypeVariantizationTest"));
+
+            var compVariantA = testProtoVariantA.Components["PrototypeVariantizationTest"].Component as PrototypeVariantizationTestComponent;
+            Assert.That(compVariantA, Is.Not.Null);
+            Assert.That(compVariantA!.NumValue, Is.EqualTo(1.5f));
+            Assert.That(compVariantA!.EnumValue, Is.EqualTo(PrototypeVariantizationTestEnum.Second));
+            Assert.That(compVariantA!.StringValue, Is.EqualTo("string-2"));
+            Assert.That(compVariantA!.StringArray.SequenceEqual(["el-1", "el-2"]));
+
+            Assert.That(compVariantA!.StringDict.TryGetValue("key1", out var value1));
+            Assert.That(value1, Is.EqualTo("string-2a"));
+            Assert.That(compVariantA!.StringDict.TryGetValue("key2", out var value2));
+            Assert.That(value2, Is.EqualTo("string-2b"));
+
+            Assert.That(compVariantA!.RecursiveRecord?.Child?.Child?.Value, Is.EqualTo("child-2"));
+        });
+
+        // Second variant
+        Assert.Multiple(() =>
+        {
+            var testProtoVariantB = protoManager.Index<EntityPrototype>(TestProtoVariantBId);
+            Assert.That(testProtoVariantB.Components, Contains.Key("PrototypeVariantizationTest"));
+
+            var compVariantB = testProtoVariantB.Components["PrototypeVariantizationTest"].Component as PrototypeVariantizationTestComponent;
+            Assert.That(compVariantB, Is.Not.Null);
+            Assert.That(compVariantB!.NumValue, Is.EqualTo(2));
+            Assert.That(compVariantB!.EnumValue, Is.EqualTo(PrototypeVariantizationTestEnum.Third));
+            Assert.That(compVariantB!.StringValue, Is.EqualTo("string-3"));
+            Assert.That(compVariantB!.StringArray.SequenceEqual(["el-1", "el-2", "el-3"]));
+
+            Assert.That(compVariantB!.StringDict.TryGetValue("key1", out var value1));
+            Assert.That(value1, Is.EqualTo("string-3a"));
+            Assert.That(compVariantB!.StringDict.TryGetValue("key2", out var value2));
+            Assert.That(value2, Is.EqualTo("string-3b"));
+
+            Assert.That(compVariantB!.RecursiveRecord?.Child?.Child?.Value, Is.EqualTo("child-3"));
+        });
+    }
+
+    const string DOCUMENT = $@"
+- type: entity
+  id: !type:CreateVariants
+    values: [ TestPrototype, TestPrototypeVariantA, TestPrototypeVariantB ]  
+  components:
+  - type: PrototypeVariantizationTest
+    numValue: !type:CreateVariants
+      values: [ 1, 1.5, 2 ] 
+    enumValue: !type:CreateVariants
+      values: [ 0, 1, 2 ]
+    stringValue: !type:CreateVariants
+      values: [ string-1, string-2, string-3 ]
+    stringArray: !type:CreateVariants
+      sequences: [ [ el-1 ], [ el-1, el-2 ], [ el-1, el-2, el-3 ] ]
+    stringDict:
+      key1: !type:CreateVariants
+        values: [ string-1a, string-2a, string-3a ]
+      key2: !type:CreateVariants
+        values: [ string-1b, string-2b, string-3b ]
+    recursiveRecord: !type:RecursionTestRecord
+      child:
+        child:
+          value: !type:CreateVariants
+            values: [ child-1, child-2, child-3 ]
+";
+}
+
+internal sealed partial class PrototypeVariantizationTestComponent : Component
+{
+    [DataField] public float NumValue = -1;
+    [DataField] public PrototypeVariantizationTestEnum EnumValue = PrototypeVariantizationTestEnum.Invalid;
+    [DataField] public string StringValue = string.Empty;
+    [DataField] public string[] StringArray = Array.Empty<string>();
+    [DataField] public Dictionary<string, string> StringDict = new();
+    [DataField] public RecursionTestRecord RecursiveRecord = new();
+}
+
+[DataDefinition]
+internal sealed partial record RecursionTestRecord
+{
+    [DataField] public RecursionTestRecord? Child = null;
+    [DataField] public string Value = string.Empty;
+}
+
+public enum PrototypeVariantizationTestEnum : int
+{
+    Invalid = -1,
+    First = 0,
+    Second = 1,
+    Third = 2,
+}
+

Robust.Shared/Prototypes/CreateVariants.cs

@@ -0,0 +1,101 @@
+using Robust.Shared.Serialization.Manager.Attributes;
+
+namespace Robust.Shared.Prototypes;
+
+/// <summary>
+/// This class is part of a mechanism that is used to automatically generate variants of prototypes at run time.
+/// It is substituted in the place of a value or linear collection (e.g., lists and arrays) in a YAML prototype definition.
+/// When the prototype manager reads this class, it will use the arrays defined here to dynamically populate the associated value/collection.
+/// To generate prototype variants, the ID of the prototype must also be replaced with a reference to this class.
+/// </summary>
+/// <example>
+/// Here is example YAML code that will auto-generate three different prototypes, each with its own distinct set of values, at runtime:
+/// 
+/// - type: entity
+///   id: !type:CreateVariants
+///     variants: [ TestEntityA, TestEntityB, TestEntityC ]
+///   components:
+///   - type: Test
+///     floatValue: !type:CreateVariants
+///       values: [ 0.5, 1, 1.5 ]
+///     enumValue: !type:CreateVariants
+///       values: [ 1, 2, 3 ]
+///     stringArray: !type:CreateVariants
+///       sequences: [ [ string-1 ], [ string-1, string-2 ], [ string-1, string-2, string-3 ] ]
+///   - type: Sprite
+///     sprite: path/to/overlay.rsi
+///     layers:
+///     - sprite: !type:CreateVariants
+///         values: [ path/to/base.rsi, path/to/base_alt1.rsi, path/to/base_alt2.rsi ]
+///       state: baseState
+///       map: ["firstLayer"]
+///     - state: overlayState
+///       map: ["secondLayer"]
+///
+/// The second prototype to be generated will have the following values:
+///
+/// - type: entity
+///   id: TestEntityB
+///   components:
+///   - type: Test
+///     floatValue: 1
+///     enumValue: 2
+///     stringArray: [ string-1, string-2 ]
+///   - type: Sprite
+///     sprite: path/to/overlay.rsi
+///     layers:
+///     - sprite: path/to/base_alt1.rsi
+///       state: baseState
+///       map: ["firstLayer"]
+///     - state: overlayState
+///       map: ["secondLayer"]
+/// </example>
+[DataDefinition]
+public sealed partial class CreateVariants
+{
+    /// <summary>
+    /// Use this field when replacing a single value (e.g., a number or string) in a prototype at runtime.
+    /// </summary>
+    [VariantValuesField]
+    public string[]? Values;
+
+    /// <summary>
+    /// Use this field when replacing a collection of values (e.g., an array) in a prototype at runtime.
+    /// </summary>
+    [VariantSequencesField]
+    public string[][]? Sequences;
+}
+
+/// <summary>
+/// Denotes an array that will be used to populate variants of a prototype at runtime.
+/// </summary>
+public sealed class VariantValuesFieldAttribute : DataFieldAttribute
+{
+    public const string Name = "values";
+
+    /// <summary>
+    /// Denotes an array that will be used to populate variants of a prototype at runtime.
+    /// </summary>
+    /// <param name="priority">See <see cref="DataFieldBaseAttribute.Priority"/>.</param>
+    public VariantValuesFieldAttribute(int priority = 1) :
+        base(Name, false, priority, false, false)
+    {
+    }
+}
+
+/// <summary>
+/// Denotes an array that will be used to populate variants of a prototype at runtime.
+/// </summary>
+public sealed class VariantSequencesFieldAttribute : DataFieldAttribute
+{
+    public const string Name = "sequences";
+
+    /// <summary>
+    /// Denotes an array that will be used to populate variants of a prototype at runtime.
+    /// </summary>
+    /// <param name="priority">See <see cref="DataFieldBaseAttribute.Priority"/>.</param>
+    public VariantSequencesFieldAttribute(int priority = 1) :
+        base(Name, false, priority, false, false)
+    {
+    }
+}

Robust.Shared/Prototypes/IPrototypeManager.cs

@@ -573,6 +573,14 @@ void ReloadPrototypes(
     /// <param name="prototypes">List of entity prototypes that form part category or null.</param>
     /// <returns>True if the provided <see cref="EntityCategoryPrototype"/> id has a matching list of <see cref="EntityPrototype"/> False otherwise.</returns>
     bool TryGetEntityPrototypesByCategory(ProtoId<EntityCategoryPrototype> category, [NotNullWhen(true)] out IReadOnlyList<EntityPrototype>? prototypes);
+
+    /// <summary>
+    /// Tries to get the list of all associated variants for a given prototype. 
+    /// </summary>
+    /// <param name="collectionMember">The prototype being indexed.</param>
+    /// <param name="collectionVariants">The collection of variants this prototype belongs to.</param>
+    /// <returns>Returns true if the prototype is part of a variant collection, false otherwise.</returns>
+    bool TryGetVariantCollection<T>(ProtoId<T> collectionMember, [NotNullWhen(true)] out List<ProtoId<T>>? collectionVariants) where T : class, IPrototype;
 }
 
 internal interface IPrototypeManagerInternal : IPrototypeManager