Clean up Marking data structure, add tests for Zombie transformation (#42756)

* Clean up Marking data structure, add tests for Zombie transformation

* empty

* AAAAAAAAAAAAAAAA

---------

Co-authored-by: Princess Cheeseballs <66055347+Pronana@users.noreply.github.com>

Content.Client/Humanoid/LayerMarkingItem.xaml.cs

@@ -103,7 +103,7 @@ public sealed partial class LayerMarkingItem : BoxContainer, ISearchableControl
             ColorsContainer.Visible = false;
         }
 
-        if (_markingsModel.TryGetMarking(_organ, _layer, _markingPrototype.ID) is { } marking &&
+        if (_markingsModel.GetMarking(_organ, _layer, _markingPrototype.ID) is { } marking &&
             _colorSliders is { } sliders)
         {
             for (var i = 0; i < _markingPrototype.Sprites.Count; i++)
@@ -144,7 +144,7 @@ public sealed partial class LayerMarkingItem : BoxContainer, ISearchableControl
         if (_colorSliders is not null)
             return;
 
-        if (_markingsModel.TryGetMarking(_organ, _layer, _markingPrototype.ID) is not { } marking)
+        if (_markingsModel.GetMarking(_organ, _layer, _markingPrototype.ID) is not { } marking)
             return;
 
         _colorSliders = new();

Content.Client/Humanoid/MarkingsViewModel.cs

@@ -98,7 +98,7 @@ public sealed class MarkingsViewModel
                 kvp => kvp.Key,
                 kvp => kvp.Value.ToDictionary(
                     it => it.Key,
-                    it => it.Value.Select(marking => new Marking(marking)).ToList()));
+                    it => it.Value.ShallowClone()));
 
             MarkingsReset?.Invoke();
         }
@@ -138,7 +138,7 @@ public sealed class MarkingsViewModel
         HumanoidVisualLayers layer,
         ProtoId<MarkingPrototype> markingId)
     {
-        return TryGetMarking(organ, layer, markingId) is not null;
+        return GetMarking(organ, layer, markingId) is not null;
     }
 
     public bool IsMarkingColorCustomizable(ProtoId<OrganCategoryPrototype> organ,
@@ -163,7 +163,7 @@ public sealed class MarkingsViewModel
         return !appearance.MatchSkin;
     }
 
-    public Marking? TryGetMarking(ProtoId<OrganCategoryPrototype> organ,
+    public Marking? GetMarking(ProtoId<OrganCategoryPrototype> organ,
         HumanoidVisualLayers layer,
         ProtoId<MarkingPrototype> markingId)
     {
@@ -173,7 +173,7 @@ public sealed class MarkingsViewModel
         if (!markingSet.TryGetValue(layer, out var markings))
             return null;
 
-        return markings.FirstOrDefault(it => it.MarkingId == markingId);
+        return markings.FirstOrNull(it => it.MarkingId == markingId);
     }
 
     public bool TrySelectMarking(ProtoId<OrganCategoryPrototype> organ,
@@ -202,8 +202,7 @@ public sealed class MarkingsViewModel
 
         var colors = _previousColors.GetValueOrDefault(markingId) ??
                      MarkingColoring.GetMarkingLayerColors(markingProto, profileData.SkinColor, profileData.EyeColor, layerMarkings);
-        var newMarking = new Marking(markingId, colors);
-        newMarking.Forced = AnyEnforcementsLifted;
+        var newMarking = new Marking(markingId, colors) { Forced = AnyEnforcementsLifted };
 
         var limits = groupPrototype.Limits.GetValueOrDefault(layer);
         if (limits is null || !EnforceLimits)
@@ -234,13 +233,9 @@ public sealed class MarkingsViewModel
     public List<Marking>? SelectedMarkings(ProtoId<OrganCategoryPrototype> organ,
         HumanoidVisualLayers layer)
     {
-        if (!_markings.TryGetValue(organ, out var organMarkings))
-            return null;
-
-        if (!organMarkings.TryGetValue(layer, out var layerMarkings))
-            return null;
-
-        return layerMarkings;
+        return !_markings.TryGetValue(organ, out var organMarkings)
+            ? null
+            : organMarkings.GetValueOrDefault(layer);
     }
 
     public bool TryDeselectMarking(ProtoId<OrganCategoryPrototype> organ,
@@ -292,10 +287,10 @@ public sealed class MarkingsViewModel
         if (!markingSet.TryGetValue(layer, out var markings))
             return;
 
-        if (markings.FirstOrDefault(it => it.MarkingId == markingId) is not { } marking)
+        if (markings.FindIndex(it => it.MarkingId == markingId) is var markingIdx && markingIdx >= 0)
             return;
 
-        marking.SetColor(colorIndex, color);
+        markings[markingIdx] = markings[markingIdx].WithColorAt(colorIndex, color);
         MarkingsChanged?.Invoke(organ, layer);
     }
 

Content.IntegrationTests/Tests/Zombie/ZombieMarkingTests.cs

@@ -0,0 +1,95 @@
+using System.Linq;
+using Content.IntegrationTests.Tests.Interaction;
+using Content.Server.Zombies;
+using Content.Shared.Body;
+using Content.Shared.Zombies;
+
+namespace Content.IntegrationTests.Tests.Zombie;
+
+[TestOf(typeof(ZombieSystem))]
+public sealed class ZombieMarkingTests : InteractionTest
+{
+    protected override string PlayerPrototype => "MobVulpkanin";
+
+    [Test]
+    public async Task ProfileApplication()
+    {
+        await Server.WaitAssertion(() =>
+        {
+            var zombie = SEntMan.System<ZombieSystem>();
+            var visualBody = SEntMan.System<SharedVisualBodySystem>();
+            zombie.ZombifyEntity(SPlayer);
+            var comp = SEntMan.GetComponent<ZombieComponent>(SPlayer);
+
+            if (!visualBody.TryGatherMarkingsData(SPlayer,
+                    null,
+                    out var profiles,
+                    out _,
+                    out _))
+            {
+                Assert.Fail($"Failed to gather markings data for {SEntMan.ToPrettyString(SPlayer):SPlayer}");
+            }
+
+            foreach (var (organ, profile) in profiles)
+            {
+                Assert.That(profile.SkinColor, Is.EqualTo(comp.SkinColor), $"Organ {organ} has non-zombified skin color");
+                Assert.That(profile.EyeColor, Is.EqualTo(comp.EyeColor), $"Organ {organ} has non-zombified skin color");
+            }
+        });
+    }
+
+    [Test]
+    public async Task MarkingApplication()
+    {
+        await Server.WaitAssertion(() =>
+        {
+            var visualBody = SEntMan.System<SharedVisualBodySystem>();
+            if (!visualBody.TryGatherMarkingsData(SPlayer,
+                    null,
+                    out _,
+                    out _,
+                    out var preZombieMarkings))
+            {
+                Assert.Fail($"Failed to gather pre-zombie markings data for {SEntMan.ToPrettyString(SPlayer):SPlayer}");
+            }
+
+            var zombie = SEntMan.System<ZombieSystem>();
+            zombie.ZombifyEntity(SPlayer);
+            var comp = SEntMan.GetComponent<ZombieComponent>(SPlayer);
+
+            if (!visualBody.TryGatherMarkingsData(SPlayer,
+                    null,
+                    out _,
+                    out _,
+                    out var postZombieMarkings))
+            {
+                Assert.Fail($"Failed to gather post-zombie markings data for {SEntMan.ToPrettyString(SPlayer):SPlayer}");
+            }
+
+            foreach (var (organ, layers) in postZombieMarkings)
+            {
+                Assert.That(preZombieMarkings, Does.ContainKey(organ), "Zombification added organs (it shouldn't)");
+                Assert.That(preZombieMarkings[organ], Is.Not.SameAs(layers), "Zombification shouldn't mutate the existing data structures");
+
+                foreach (var (layer, markingSet) in layers)
+                {
+                    Assert.That(preZombieMarkings[organ], Does.ContainKey(layer), "Zombification added layers (it shouldn't)");
+                    Assert.That(preZombieMarkings[organ][layer], Is.Not.SameAs(markingSet), "Zombification shouldn't mutate the existing data structures");
+                    Assert.That(preZombieMarkings[organ][layer], Has.Count.EqualTo(markingSet.Count), "Zombification shouldn't change the amount of markings");
+
+                    if (!ZombieSystem.AdditionalZombieLayers.Contains(layer))
+                        continue;
+
+                    foreach (var (preMarking, postMarking) in preZombieMarkings[organ][layer].Zip(markingSet))
+                    {
+                        Assert.That(preMarking, Is.Not.EqualTo(postMarking), $"Zombification should change marking {postMarking.MarkingId} on layer {layer}");
+                        foreach (var color in postMarking.MarkingColors)
+                        {
+                            Assert.That(color, Is.EqualTo(comp.SkinColor), $"Zombification should change {postMarking.MarkingId} on layer {layer} to the skin color");
+                        }
+                    }
+                }
+            }
+        });
+    }
+}

Content.Server/Database/ServerDbBase.cs

@@ -254,7 +254,7 @@ namespace Content.Server.Database
 
                     if (parsed is null) continue;
 
-                    markingsList.Add(parsed);
+                    markingsList.Add(parsed.Value);
                 }
 
                 if (Marking.ParseFromDbString($"{profile.HairName}@{profile.HairColor}") is { } facialMarking)
@@ -348,7 +348,7 @@ namespace Content.Server.Database
             var legacyMarkings = appearance.Markings
                 .SelectMany(organ => organ.Value.Values)
                 .SelectMany(i => i)
-                .Select(marking => marking.ToString())
+                .Select(marking => marking.ToLegacyDbString())
                 .ToList();
             var flattenedMarkings = appearance.Markings.SelectMany(it => it.Value)
                 .ToDictionary(kvp => kvp.Key, kvp => kvp.Value);

Content.Server/Wagging/WaggingSystem.cs

@@ -98,9 +98,9 @@ public sealed class WaggingSystem : EntitySystem
                 }
                 else
                 {
-                    if (currentMarkingId.EndsWith(ent.Comp.Suffix))
+                    if (currentMarkingId.Id.EndsWith(ent.Comp.Suffix))
                     {
-                        newMarkingId = currentMarkingId[..^ent.Comp.Suffix.Length];
+                        newMarkingId = currentMarkingId.Id[..^ent.Comp.Suffix.Length];
                     }
                     else
                     {

Content.Server/Zombies/ZombieSystem.Transform.cs

@@ -45,6 +45,7 @@ using Robust.Shared.Prototypes;
 using Content.Shared.NPC.Prototypes;
 using Content.Shared.Roles;
 using Content.Shared.Temperature.Components;
+using Robust.Shared.Utility;
 
 namespace Content.Server.Zombies;
 
@@ -78,7 +79,7 @@ public sealed partial class ZombieSystem
     private static readonly ProtoId<NpcFactionPrototype> ZombieFaction = "Zombie";
     private static readonly string MindRoleZombie = "MindRoleZombie";
     private static readonly List<ProtoId<AntagPrototype>> BannableZombiePrototypes = ["Zombie"];
-    private static readonly HashSet<HumanoidVisualLayers> AdditionalZombieLayers = [HumanoidVisualLayers.Tail, HumanoidVisualLayers.HeadSide, HumanoidVisualLayers.HeadTop, HumanoidVisualLayers.Snout];
+    internal static readonly HashSet<HumanoidVisualLayers> AdditionalZombieLayers = [HumanoidVisualLayers.Tail, HumanoidVisualLayers.HeadSide, HumanoidVisualLayers.HeadTop, HumanoidVisualLayers.Snout];
 
     /// <summary>
     /// Handles an entity turning into a zombie when they die or go into crit
@@ -201,27 +202,33 @@ public sealed partial class ZombieSystem
                 kvp => kvp.Key,
                 kvp => kvp.Value.ToDictionary(
                     it => it.Key,
-                    it => it.Value.Select(marking => new Marking(marking)).ToList()));
+                    it => it.Value.ShallowClone()));
 
             var zombifiedProfiles = profiles.ToDictionary(pair => pair.Key,
                 pair => pair.Value with { EyeColor = zombiecomp.EyeColor, SkinColor = zombiecomp.SkinColor });
             _visualBody.ApplyProfiles(target, zombifiedProfiles);
 
-            foreach (var markingSet in markings.Values)
+            var newMarkings = markings.ToDictionary(
+                kvp => kvp.Key,
+                kvp => kvp.Value.ToDictionary(
+                    it => it.Key,
+                    it => it.Value.ShallowClone()));
+
+            foreach (var markingSet in newMarkings.Values)
             {
                 foreach (var (layer, layerMarkings) in markingSet)
                 {
                     if (!AdditionalZombieLayers.Contains(layer))
                         continue;
 
-                    foreach (var marking in layerMarkings)
+                    for (var i = 0; i < layerMarkings.Count; i++)
                     {
-                        marking.SetColor(zombiecomp.SkinColor);
+                        layerMarkings[i] = layerMarkings[i].WithColor(zombiecomp.SkinColor);
                     }
                 }
             }
 
-            _visualBody.ApplyMarkings(target, markings);
+            _visualBody.ApplyMarkings(target, newMarkings);
         }
 
         //We have specific stuff for humanoid zombies because they matter more

Content.Shared/Body/SharedVisualBodySystem.cs

@@ -58,7 +58,7 @@ public abstract partial class SharedVisualBodySystem : EntitySystem
             };
             if (appearances.GetValueOrDefault(prototype.BodyPart) is { MatchSkin: true } appearance && skinColor is { } color)
             {
-                markingWithColor.SetColor(color.WithAlpha(appearance.LayerAlpha));
+                markingWithColor = markingWithColor.WithColor(color.WithAlpha(appearance.LayerAlpha));
             }
             ret.Add(markingWithColor);
         }

Content.Shared/Humanoid/Markings/Marking.cs

@@ -1,140 +1,103 @@
 using System.Linq;
-using Content.Shared.Humanoid.Prototypes;
 using Robust.Shared.Prototypes;
 using Robust.Shared.Serialization;
+using Robust.Shared.Utility;
 
-namespace Content.Shared.Humanoid.Markings
+namespace Content.Shared.Humanoid.Markings;
+
+/// <summary>
+///     Represents a marking ID and its colors
+/// </summary>
+[DataDefinition, Serializable, NetSerializable]
+public partial record struct Marking
 {
-    [DataDefinition]
-    [Serializable, NetSerializable]
-    public sealed partial class Marking : IEquatable<Marking>, IComparable<Marking>, IComparable<string>
+    /// <summary>
+    /// The <see cref="MarkingPrototype"/> referred to by this marking
+    /// </summary>
+    [DataField(required: true)]
+    public ProtoId<MarkingPrototype> MarkingId;
+
+    [DataField("markingColor")]
+    private List<Color> _markingColors;
+
+    /// <summary>
+    /// The colors taken on by the marking
+    /// </summary>
+    public IReadOnlyList<Color> MarkingColors => _markingColors;
+
+    /// <summary>
+    /// Whether the marking is forced regardless of points
+    /// </summary>
+    public bool Forced;
+
+    public Marking()
     {
-        [DataField("markingColor")]
-        private List<Color> _markingColors = new();
+        _markingColors = new();
+    }
 
-        private Marking()
+    public Marking(ProtoId<MarkingPrototype> markingId, IEnumerable<Color> colors)
+    {
+        MarkingId = markingId;
+        _markingColors = colors.ToList();
+    }
+
+    public Marking(ProtoId<MarkingPrototype> markingId, int colorsCount) : this(markingId,
+        Enumerable.Repeat(Color.White, colorsCount).ToList())
+    {
+    }
+
+    public bool Equals(Marking other)
+    {
+        return MarkingId.Equals(other.MarkingId)
+            && MarkingColors.SequenceEqual(other.MarkingColors)
+            && Forced.Equals(other.Forced);
+    }
+
+    public override int GetHashCode()
+    {
+        return HashCode.Combine(MarkingId, MarkingColors, Forced);
+    }
+
+    public Marking WithColor(Color color) =>
+        this with { _markingColors = Enumerable.Repeat(color, MarkingColors.Count).ToList() };
+
+    public Marking WithColorAt(int index, Color color)
+    {
+        var newColors = _markingColors.ShallowClone();
+        newColors[index] = color;
+        return this with { _markingColors = newColors };
+    }
+
+    // look this could be better but I don't think serializing
+    // colors is the correct thing to do
+    //
+    // this is still janky imo but serializing a color and feeding
+    // it into the default JSON serializer (which is just *fine*)
+    // doesn't seem to have compatible interfaces? this 'works'
+    // for now but should eventually be improved so that this can,
+    // in fact just be serialized through a convenient interface
+    public string ToLegacyDbString()
+    {
+        // reserved character
+        string sanitizedName = MarkingId.Id.Replace('@', '_');
+        List<string> colorStringList = new();
+        foreach (var color in MarkingColors)
+            colorStringList.Add(color.ToHex());
+
+        return $"{sanitizedName}@{String.Join(',', colorStringList)}";
+    }
+
+    public static Marking? ParseFromDbString(string input)
+    {
+        if (input.Length == 0) return null;
+        var split = input.Split('@');
+        if (split.Length != 2) return null;
+        List<Color> colorList = new();
+        foreach (string color in split[1].Split(','))
         {
+            colorList.Add(Color.FromHex(color));
         }
 
-        public Marking(string markingId,
-            List<Color> markingColors)
-        {
-            MarkingId = markingId;
-            _markingColors = markingColors;
-        }
-
-        public Marking(string markingId,
-            IReadOnlyList<Color> markingColors)
-            : this(markingId, new List<Color>(markingColors))
-        {
-        }
-
-        public Marking(string markingId, int colorCount)
-        {
-            MarkingId = markingId;
-            List<Color> colors = new();
-            for (int i = 0; i < colorCount; i++)
-                colors.Add(Color.White);
-            _markingColors = colors;
-        }
-
-        public Marking(Marking other)
-        {
-            MarkingId = other.MarkingId;
-            _markingColors = new(other.MarkingColors);
-            Forced = other.Forced;
-        }
-
-        /// <summary>
-        ///     ID of the marking prototype.
-        /// </summary>
-        [DataField("markingId", required: true)]
-        public string MarkingId { get; private set; } = default!;
-
-        /// <summary>
-        ///     All colors currently on this marking.
-        /// </summary>
-        [ViewVariables]
-        public IReadOnlyList<Color> MarkingColors => _markingColors;
-
-        /// <summary>
-        ///     If this marking should be forcefully applied, regardless of points.
-        /// </summary>
-        [ViewVariables]
-        public bool Forced;
-
-        public void SetColor(int colorIndex, Color color) =>
-            _markingColors[colorIndex] = color;
-
-        public void SetColor(Color color)
-        {
-            for (int i = 0; i < _markingColors.Count; i++)
-            {
-                _markingColors[i] = color;
-            }
-        }
-
-        public int CompareTo(Marking? marking)
-        {
-            if (marking == null)
-            {
-                return 1;
-            }
-
-            return string.Compare(MarkingId, marking.MarkingId, StringComparison.Ordinal);
-        }
-
-        public int CompareTo(string? markingId)
-        {
-            if (markingId == null)
-                return 1;
-
-            return string.Compare(MarkingId, markingId, StringComparison.Ordinal);
-        }
-
-        public bool Equals(Marking? other)
-        {
-            if (other == null)
-            {
-                return false;
-            }
-            return MarkingId.Equals(other.MarkingId)
-                && _markingColors.SequenceEqual(other._markingColors)
-                && Forced.Equals(other.Forced);
-        }
-
-        // VERY BIG TODO: TURN THIS INTO JSONSERIALIZER IMPLEMENTATION
-
-
-        // look this could be better but I don't think serializing
-        // colors is the correct thing to do
-        //
-        // this is still janky imo but serializing a color and feeding
-        // it into the default JSON serializer (which is just *fine*)
-        // doesn't seem to have compatible interfaces? this 'works'
-        // for now but should eventually be improved so that this can,
-        // in fact just be serialized through a convenient interface
-        new public string ToString()
-        {
-            // reserved character
-            string sanitizedName = this.MarkingId.Replace('@', '_');
-            List<string> colorStringList = new();
-            foreach (Color color in _markingColors)
-                colorStringList.Add(color.ToHex());
-
-            return $"{sanitizedName}@{String.Join(',', colorStringList)}";
-        }
-
-        public static Marking? ParseFromDbString(string input)
-        {
-            if (input.Length == 0) return null;
-            var split = input.Split('@');
-            if (split.Length != 2) return null;
-            List<Color> colorList = new();
-            foreach (string color in split[1].Split(','))
-                colorList.Add(Color.FromHex(color));
-
-            return new Marking(split[0], colorList);
-        }
+        return new Marking(split[0], colorList);
     }
 }