Fix reloading prototypes with AlwaysPushInheritance (#5612)

Author: ElectroJr

RELEASE-NOTES.md

@@ -43,7 +43,7 @@ END TEMPLATE-->
 
 ### Bugfixes
 
-*None yet*
+* Fixed prototype reloading/hotloading not properly handling data-fields with the `AlwaysPushInheritanceAttribute`
 
 ### Other
 

Robust.Client/Replays/Loading/ReplayLoadManager.Checkpoints.cs

@@ -352,6 +352,9 @@ private void LoadPrototype(
                 // prototype changes when jumping around in time. This also requires reworking how the initial
                 // implicit state data is generated, because we can't simply cache it anymore.
                 // Also, does reloading prototypes in release mode modify existing entities?
+                // Yes, yes it does. See PrototypeReloadSystem.UpdateEntity()
+                // Its just not supported ATM.
+                // TBH it'd be easier if overriding existing prototypes in release mode was just forbidden.
 
                 var msg = $"Overwriting an existing prototype! Kind: {kind.Name}. Ids: {string.Join(", ", ids)}";
                 if (_confMan.GetCVar(CVars.ReplayIgnoreErrors))
@@ -361,7 +364,6 @@ private void LoadPrototype(
             }
         }
 
-        _protoMan.ResolveResults();
         _protoMan.ReloadPrototypes(changed);
         _locMan.ReloadLocalizations();
     }

Robust.Shared/Prototypes/IPrototypeManager.cs

@@ -284,15 +284,21 @@ Dictionary<string, HashSet<ErrorNode>> ValidateDirectory(ResPath path,
     void LoadDefaultPrototypes(Dictionary<Type, HashSet<string>>? loaded = null);
 
     /// <summary>
-    /// Syncs all inter-prototype data. Call this when operations adding new prototypes are done.
+    /// Call this when operations adding new prototypes are done.
+    /// This will handle prototype inheritance, instance creation, and update entity categories.
+    /// When loading extra prototypes, or reloading a subset of existing prototypes, you should probably use
+    /// <see cref="ReloadPrototypes"/> instead.
     /// </summary>
     void ResolveResults();
 
     /// <summary>
-    /// Invokes <see cref="PrototypesReloaded"/> with information about the modified prototypes.
-    /// When built with development tools, this will also push inheritance for reloaded prototypes/
+    /// This should be called after new or updated prototypes ahve been loaded.
+    /// This will handle prototype inheritance, instance creation, and update entity categories.
+    /// It will also invoke <see cref="PrototypesReloaded"/> and raise a <see cref="PrototypesReloadedEventArgs"/>
+    /// event with information about the modified prototypes.
     /// </summary>
-    void ReloadPrototypes(Dictionary<Type, HashSet<string>> modified,
+    void ReloadPrototypes(
+        Dictionary<Type, HashSet<string>> modified,
         Dictionary<Type, HashSet<string>>? removed = null);
 
     /// <summary>

Robust.Shared/Prototypes/PrototypeManager.YamlLoad.cs

@@ -209,10 +209,10 @@ private void MergeMapping(
 
         var kindData = _kinds[kind];
 
-        if (!overwrite && kindData.Results.ContainsKey(id))
+        if (!overwrite && kindData.RawResults.ContainsKey(id))
             throw new PrototypeLoadException($"Duplicate ID: '{id}' for kind '{kind}");
 
-        kindData.Results[id] = data;
+        kindData.RawResults[id] = data;
 
         if (kindData.Inheritance is { } inheritance)
         {
@@ -295,6 +295,7 @@ public void RemoveString(string prototypes)
                 kindData.UnfrozenInstances ??= kindData.Instances.ToDictionary();
                 kindData.UnfrozenInstances.Remove(id);
                 kindData.Results.Remove(id);
+                kindData.RawResults.Remove(id);
                 modified.Add(kindData);
             }
         }

Robust.Shared/Prototypes/PrototypeManager.cs

@@ -335,98 +335,117 @@ protected void ReloadPrototypes(IEnumerable<ResPath> filePaths)
         }
 
         /// <inheritdoc />
-        public void ReloadPrototypes(Dictionary<Type, HashSet<string>> modified,
+        public void ReloadPrototypes(
+            Dictionary<Type, HashSet<string>> modified,
             Dictionary<Type, HashSet<string>>? removed = null)
         {
-#if TOOLS
             var prototypeTypeOrder = modified.Keys.ToList();
             prototypeTypeOrder.Sort(SortPrototypesByPriority);
 
-            var pushed = new Dictionary<Type, HashSet<string>>();
+            var byType = new Dictionary<Type, PrototypesReloadedEventArgs.PrototypeChangeSet>();
             var modifiedKinds = new HashSet<KindData>();
+            var toProcess = new HashSet<string>();
+            var processQueue = new Queue<string>();
 
             foreach (var kind in prototypeTypeOrder)
             {
+                var modifiedInstances = new Dictionary<string, IPrototype>();
                 var kindData = _kinds[kind];
-                if (!kind.IsAssignableTo(typeof(IInheritingPrototype)))
-                {
-                    foreach (var id in modified[kind])
-                    {
-                        var prototype = (IPrototype)_serializationManager.Read(kind, kindData.Results[id])!;
-                        kindData.UnfrozenInstances ??= kindData.Instances.ToDictionary();
-                        kindData.UnfrozenInstances[id] = prototype;
-                        modifiedKinds.Add(kindData);
-                    }
 
-                    continue;
-                }
+                var tree = kindData.Inheritance;
+                toProcess.Clear();
+                processQueue.Clear();
+
+                DebugTools.AssertEqual(kind.IsAssignableTo(typeof(IInheritingPrototype)), tree != null);
+                DebugTools.Assert(tree != null || kindData.RawResults == kindData.Results);
 
-                var tree = kindData.Inheritance!;
-                var processQueue = new Queue<string>();
                 foreach (var id in modified[kind])
                 {
+                    AddToQueue(id);
+                }
+
+                void AddToQueue(string id)
+                {
+                    if (!toProcess.Add(id))
+                        return;
                     processQueue.Enqueue(id);
+
+                    if (tree == null)
+                        return;
+
+                    if (!tree.TryGetChildren(id, out var children))
+                        return;
+
+                    foreach (var child in children!)
+                    {
+                        AddToQueue(child);
+                    }
                 }
 
                 while (processQueue.TryDequeue(out var id))
                 {
-                    var pushedSet = pushed.GetOrNew(kind);
-
-                    if (tree.TryGetParents(id, out var parents))
+                    DebugTools.Assert(toProcess.Contains(id));
+                    if (tree != null)
                     {
-                        var nonPushedParent = false;
-                        foreach (var parent in parents)
+                        if (tree.TryGetParents(id, out var parents))
                         {
-                            //our parent has been reloaded and has not been added to the pushedSet yet
-                            if (modified[kind].Contains(parent) && !pushedSet.Contains(parent))
+                            DebugTools.Assert(parents.Length > 0);
+                            var nonPushedParent = false;
+                            foreach (var parent in parents)
                             {
-                                //we re-queue ourselves at the end of the queue
+                                if (!toProcess.Contains(parent))
+                                    continue;
+
+                                // our parent has been modified, but has not yet been processed.
+                                // we re-queue ourselves at the end of the queue.
+                                DebugTools.Assert(processQueue.Contains(parent));
                                 processQueue.Enqueue(id);
                                 nonPushedParent = true;
                                 break;
                             }
-                        }
 
-                        if (nonPushedParent)
-                            continue;
+                            if (nonPushedParent)
+                                continue;
 
-                        var parentMaps = new MappingDataNode[parents.Length];
-                        for (var i = 0; i < parentMaps.Length; i++)
+                            var parentMaps = new MappingDataNode[parents.Length];
+                            for (var i = 0; i < parentMaps.Length; i++)
+                            {
+                                parentMaps[i] = kindData.Results[parents[i]];
+                            }
+
+                            kindData.Results[id] = _serializationManager.PushCompositionWithGenericNode(
+                                kind,
+                                parentMaps,
+                                kindData.RawResults[id]);
+                        }
+                        else
                         {
-                            parentMaps[i] = kindData.Results[parents[i]];
+                            kindData.Results[id] = kindData.RawResults[id];
                         }
-
-                        kindData.Results[id] = _serializationManager.PushCompositionWithGenericNode(
-                            kind,
-                            parentMaps,
-                            kindData.Results[id]);
                     }
 
+                    toProcess.Remove(id);
 
                     var prototype = TryReadPrototype(kind, id, kindData.Results[id], SerializationHookContext.DontSkipHooks);
-                    if (prototype != null)
-                    {
-                        kindData.UnfrozenInstances ??= kindData.Instances.ToDictionary();
-                        kindData.UnfrozenInstances[id] = prototype;
-                        modifiedKinds.Add(kindData);
-                    }
+                    if (prototype == null)
+                        continue;
 
-                    pushedSet.Add(id);
+                    kindData.UnfrozenInstances ??= kindData.Instances.ToDictionary();
+                    kindData.UnfrozenInstances[id] = prototype;
+                    modifiedInstances.Add(id, prototype);
                 }
+
+                if (modifiedInstances.Count == 0)
+                    continue;
+
+                byType.Add(kindData.Type, new(modifiedInstances));
+                modifiedKinds.Add(kindData);
             }
 
             Freeze(modifiedKinds);
+
             if (modifiedKinds.Any(x => x.Type == typeof(EntityPrototype) || x.Type == typeof(EntityCategoryPrototype)))
                 UpdateCategories();
-#endif
-
-            //todo paul i hate it but i am not opening that can of worms in this refactor
-            var byType = modified
-                .ToDictionary(
-                    g => g.Key,
-                    g => new PrototypesReloadedEventArgs.PrototypeChangeSet(
-                        g.Value.Where(x => _kinds[g.Key].Instances.ContainsKey(x))
-                            .ToDictionary(a => a, a => _kinds[g.Key].Instances[a])));
 
             var modifiedTypes = new HashSet<Type>(byType.Keys);
             if (removed != null)
@@ -592,11 +611,9 @@ private async Task PushKindInheritance(Type kind, KindData data)
 
             // var sw = RStopwatch.StartNew();
 
-            var results = new Dictionary<string, InheritancePushDatum>(
-                data.Results.Select(k => new KeyValuePair<string, InheritancePushDatum>(
-                    k.Key,
-                    new InheritancePushDatum(k.Value, tree.GetParentsCount(k.Key))))
-            );
+            var results = data.RawResults.ToDictionary(
+                k => k.Key,
+                k => new InheritancePushDatum(k.Value, tree.GetParentsCount(k.Key)));
 
             using var countDown = new CountdownEvent(results.Count);
 
@@ -1015,6 +1032,8 @@ private void RegisterKind(Type kind, Dictionary<Type, KindData> kinds)
 
             if (kind.IsAssignableTo(typeof(IInheritingPrototype)))
                 kindData.Inheritance = new MultiRootInheritanceGraph<string>();
+            else
+                kindData.Results = kindData.RawResults;
         }
 
         /// <inheritdoc />
@@ -1026,7 +1045,13 @@ private sealed class KindData(Type kind, string name)
 
             public FrozenDictionary<string, IPrototype> Instances = FrozenDictionary<string, IPrototype>.Empty;
 
-            public readonly Dictionary<string, MappingDataNode> Results = new();
+            public Dictionary<string, MappingDataNode> Results = new();
+
+            /// <summary>
+            /// Variant of <see cref="Results"/> prior to inheritance pushing. If the kind does not have inheritance,
+            /// then this is just the same dictionary.
+            /// </summary>
+            public readonly Dictionary<string, MappingDataNode> RawResults = new();
 
             public readonly Type Type = kind;
             public readonly string Name = name;

Robust.Shared/Serialization/Manager/SerializationManager.Composition.cs

@@ -26,14 +26,27 @@ private delegate DataNode PushCompositionDelegate(
 
     public DataNode PushComposition(Type type, DataNode[] parents, DataNode child, ISerializationContext? context = null)
     {
+        if (parents.Length == 0)
+            return child.Copy();
+
         DebugTools.Assert(parents.All(x => x.GetType() == child.GetType()));
 
+        // TODO SERIALIZATION
+        // Change inheritance pushing so that it modifies the passed in child.
+        // I.e., move the child.Clone() statement to the beginning here, then make the delegate modify the clone.
+        // Currently pusing more than one parent requires multiple unnecessary clones.
+
         var pusher = GetOrCreatePushCompositionDelegate(type, child);
 
         var node = child;
-        for (int i = 0; i < parents.Length; i++)
+        foreach (var parent in parents)
         {
-            node = pusher(type, parents[i], node, context);
+            var newNode = pusher(type, parent, node, context);
+
+            // Currently delegate pusher should be returning a new instance, and not modifying the passed in child.
+            DebugTools.Assert(!ReferenceEquals(newNode, node));
+
+            node = newNode;
         }
 
         return node;
@@ -95,7 +108,7 @@ private PushCompositionDelegate GetOrCreatePushCompositionDelegate(Type type, Da
                         Expression.Convert(parentParam, nodeType)),
                     MappingDataNode => Expression.Call(
                         instanceConst,
-                        nameof(PushInheritanceMapping),
+                        nameof(CombineMappings),
                         Type.EmptyTypes,
                         Expression.Convert(childParam, nodeType),
                         Expression.Convert(parentParam, nodeType)),
@@ -117,32 +130,26 @@ private SequenceDataNode PushInheritanceSequence(SequenceDataNode child, Sequenc
         //todo implement different inheritancebehaviours for yamlfield
         // I have NFI what this comment means.
 
-        var result = new SequenceDataNode(child.Count + parent.Count);
+        var result = child.Copy();
         foreach (var entry in parent)
         {
-            result.Add(entry);
-        }
-        foreach (var entry in child)
-        {
-            result.Add(entry);
+            result.Add(entry.Copy());
         }
 
         return result;
     }
 
-    private MappingDataNode PushInheritanceMapping(MappingDataNode child, MappingDataNode parent)
+    public MappingDataNode CombineMappings(MappingDataNode child, MappingDataNode parent)
     {
         //todo implement different inheritancebehaviours for yamlfield
         // I have NFI what this comment means.
+        // I still don't know what it means, but if it's talking about the always/never push inheritance attributes,
+        // make sure it doesn't break entity serialization.
 
-        var result = new MappingDataNode(child.Count + parent.Count);
+        var result = child.Copy();
         foreach (var (k, v) in parent)
         {
-            result[k] = v;
-        }
-        foreach (var (k, v) in child)
-        {
-            result[k] = v;
+            result.TryAddCopy(k, v);
         }
 
         return result;

Robust.Shared/Serialization/Markdown/Mapping/MappingDataNode.cs

@@ -3,6 +3,7 @@
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
+using System.Runtime.InteropServices;
 using System.Threading;
 using Robust.Shared.Serialization.Markdown.Value;
 using Robust.Shared.Utility;
@@ -396,5 +397,20 @@ public bool Remove(KeyValuePair<DataNode, DataNode> item)
 
         public int Count => _children.Count;
         public bool IsReadOnly => false;
+
+        public bool TryAdd(DataNode key, DataNode value)
+        {
+            return _children.TryAdd(key, value);
+        }
+
+        public bool TryAddCopy(DataNode key, DataNode value)
+        {
+            ref var entry = ref CollectionsMarshal.GetValueRefOrAddDefault(_children, key, out var exists);
+            if (exists)
+                return false;
+
+            entry = value.Copy();
+            return true;
+        }
     }
 }

Robust.Shared/Upload/SharedPrototypeLoadManager.cs

@@ -42,7 +42,6 @@ protected virtual void LoadPrototypeData(GamePrototypeLoadMessage message)
 
         var changed = new Dictionary<Type, HashSet<string>>();
         _prototypeManager.LoadString(data, true, changed);
-        _prototypeManager.ResolveResults();
         _prototypeManager.ReloadPrototypes(changed);
         _localizationManager.ReloadLocalizations();