Construction graph improvements (#17960)

Content.Shared/Construction/ConstructionGraphNode.cs

@@ -1,4 +1,6 @@
 using System.Diagnostics.CodeAnalysis;
+using Content.Shared.Construction.NodeEntities;
+using Content.Shared.Construction.Serialization;
 using Robust.Shared.Prototypes;
 using Robust.Shared.Serialization.TypeSerializers.Implementations.Custom.Prototype;
 
@@ -23,8 +25,11 @@ namespace Content.Shared.Construction
         [ViewVariables]
         public IReadOnlyList<IGraphAction> Actions => _actions;
 
-        [DataField("entity", customTypeSerializer:typeof(PrototypeIdSerializer<EntityPrototype>))]
-        public string? Entity { get; private set; }
+        [DataField("transform")]
+        public IGraphTransform[] TransformLogic = Array.Empty<IGraphTransform>();
+
+        [DataField("entity", customTypeSerializer: typeof(GraphNodeEntitySerializer), serverOnly:true)]
+        public IGraphNodeEntity Entity { get; private set; } = new NullNodeEntity();
 
         public ConstructionGraphEdge? GetEdge(string target)
         {

Content.Shared/Construction/IGraphNodeEntity.cs

@@ -0,0 +1,29 @@
+using Robust.Shared.Prototypes;
+
+namespace Content.Shared.Construction;
+
+public interface IGraphNodeEntity
+{
+    /// <summary>
+    ///     Gets the <see cref="EntityPrototype"/> ID for a node, given the <see cref="EntityUid"/> of both the
+    ///     construction entity and the user entity.
+    ///     If the construction entity is null, then we are dealing with a "start construction" for an entity that
+    ///     does not exist yet.
+    ///     If the user entity is null, this node was reached through means other some sort of "user interaction".
+    /// </summary>
+    /// <param name="uid">Uid of the construction entity.</param>
+    /// <param name="userUid">Uid of the user that caused the transition to the node.</param>
+    /// <param name="args">Arguments with useful instances, etc.</param>
+    /// <returns></returns>
+    public string? GetId(EntityUid? uid, EntityUid? userUid, GraphNodeEntityArgs args);
+}
+
+public readonly struct GraphNodeEntityArgs
+{
+    public readonly IEntityManager EntityManager;
+
+    public GraphNodeEntityArgs(IEntityManager entityManager)
+    {
+        EntityManager = entityManager;
+    }
+}

Content.Shared/Construction/IGraphTransform.cs

@@ -0,0 +1,16 @@
+namespace Content.Shared.Construction;
+
+public interface IGraphTransform
+{
+    public void Transform(EntityUid oldUid, EntityUid newUid, EntityUid? userUid, GraphTransformArgs args);
+}
+
+public readonly struct GraphTransformArgs
+{
+    public readonly IEntityManager EntityManager;
+
+    public GraphTransformArgs(IEntityManager entityManager)
+    {
+        EntityManager = entityManager;
+    }
+}

Content.Shared/Construction/NodeEntities/NullNodeEntity.cs

@@ -0,0 +1,13 @@
+using JetBrains.Annotations;
+
+namespace Content.Shared.Construction.NodeEntities;
+
+[UsedImplicitly]
+[DataDefinition]
+public sealed class NullNodeEntity : IGraphNodeEntity
+{
+    public string? GetId(EntityUid? uid, EntityUid? userUid, GraphNodeEntityArgs args)
+    {
+        return null;
+    }
+}

Content.Shared/Construction/NodeEntities/StaticNodeEntity.cs

@@ -0,0 +1,28 @@
+using JetBrains.Annotations;
+using Robust.Shared.Prototypes;
+using Robust.Shared.Serialization.TypeSerializers.Implementations.Custom.Prototype;
+
+namespace Content.Shared.Construction.NodeEntities;
+
+[UsedImplicitly]
+[DataDefinition]
+public sealed class StaticNodeEntity : IGraphNodeEntity
+{
+    [DataField("id", customTypeSerializer:typeof(PrototypeIdSerializer<EntityPrototype>))]
+    [ViewVariables(VVAccess.ReadWrite)]
+    public string? Id { get; }
+
+    public StaticNodeEntity()
+    {
+    }
+
+    public StaticNodeEntity(string id)
+    {
+        Id = id;
+    }
+
+    public string? GetId(EntityUid? uid, EntityUid? userUid, GraphNodeEntityArgs args)
+    {
+        return Id;
+    }
+}

Content.Shared/Construction/Serialization/GraphNodeEntitySerializer.cs

@@ -0,0 +1,53 @@
+using Content.Shared.Construction.NodeEntities;
+using Robust.Shared.Prototypes;
+using Robust.Shared.Serialization;
+using Robust.Shared.Serialization.Manager;
+using Robust.Shared.Serialization.Markdown;
+using Robust.Shared.Serialization.Markdown.Mapping;
+using Robust.Shared.Serialization.Markdown.Validation;
+using Robust.Shared.Serialization.Markdown.Value;
+using Robust.Shared.Serialization.TypeSerializers.Interfaces;
+
+namespace Content.Shared.Construction.Serialization;
+
+public sealed class GraphNodeEntitySerializer : ITypeSerializer<IGraphNodeEntity, ValueDataNode>, ITypeSerializer<IGraphNodeEntity, MappingDataNode>
+{
+    public ValidationNode Validate(ISerializationManager serializationManager, ValueDataNode node,
+        IDependencyCollection dependencies, ISerializationContext? context = null)
+    {
+        var id = node.Value;
+
+        var protoMan = dependencies.Resolve<IPrototypeManager>();
+
+        if (!protoMan.HasIndex<EntityPrototype>(id))
+        {
+            return new ErrorNode(node, $"Entity Prototype {id} was not found!");
+        }
+
+        return new ValidatedValueNode(node);
+    }
+
+    public IGraphNodeEntity Read(ISerializationManager serializationManager, ValueDataNode node,
+        IDependencyCollection dependencies, SerializationHookContext hookCtx, ISerializationContext? context = null, ISerializationManager.InstantiationDelegate<IGraphNodeEntity>? instanceProvider = null)
+    {
+        return new StaticNodeEntity(node.Value);
+    }
+
+    public ValidationNode Validate(ISerializationManager serializationManager, MappingDataNode node,
+        IDependencyCollection dependencies, ISerializationContext? context = null)
+    {
+        return serializationManager.ValidateNode<IGraphNodeEntity>(node, context);
+    }
+
+    public IGraphNodeEntity Read(ISerializationManager serializationManager, MappingDataNode node,
+        IDependencyCollection dependencies, SerializationHookContext hookCtx, ISerializationContext? context = null, ISerializationManager.InstantiationDelegate<IGraphNodeEntity>? instanceProvider = null)
+    {
+        return serializationManager.Read(node, hookCtx, context, instanceProvider, false);
+    }
+
+    public DataNode Write(ISerializationManager serializationManager, IGraphNodeEntity value, IDependencyCollection dependencies,
+        bool alwaysWrite = false, ISerializationContext? context = null)
+    {
+        return serializationManager.WriteValue(value, alwaysWrite, context, false);
+    }
+}

Content.Shared/Storage/EntitySystems/SharedItemMapperSystem.cs

@@ -32,11 +32,14 @@ namespace Content.Shared.Storage.EntitySystems
                 val.Layer = layerName;
             }
 
-            if (EntityManager.TryGetComponent(component.Owner, out AppearanceComponent? appearanceComponent))
+            if (EntityManager.TryGetComponent(uid, out AppearanceComponent? appearanceComponent))
             {
                 var list = new List<string>(component.MapLayers.Keys);
-                _appearance.SetData(component.Owner, StorageMapVisuals.InitLayers, new ShowLayerData(list), appearanceComponent);
+                _appearance.SetData(uid, StorageMapVisuals.InitLayers, new ShowLayerData(list), appearanceComponent);
             }
+
+            // Ensure appearance is correct with current contained entities.
+            UpdateAppearance(uid, component);
         }
 
         private void MapperEntityRemoved(EntityUid uid, ItemMapperComponent itemMapper,
@@ -45,7 +48,7 @@ namespace Content.Shared.Storage.EntitySystems
             if (itemMapper.ContainerWhitelist != null && !itemMapper.ContainerWhitelist.Contains(args.Container.ID))
                 return;
 
-            UpdateAppearance(uid, itemMapper, args);
+            UpdateAppearance(uid, itemMapper);
         }
 
         private void MapperEntityInserted(EntityUid uid, ItemMapperComponent itemMapper,
@@ -54,40 +57,43 @@ namespace Content.Shared.Storage.EntitySystems
             if (itemMapper.ContainerWhitelist != null && !itemMapper.ContainerWhitelist.Contains(args.Container.ID))
                 return;
 
-            UpdateAppearance(uid, itemMapper, args);
+            UpdateAppearance(uid, itemMapper);
         }
 
-        private void UpdateAppearance(EntityUid uid, ItemMapperComponent itemMapper, ContainerModifiedMessage message)
+        private void UpdateAppearance(EntityUid uid, ItemMapperComponent? itemMapper = null)
         {
-            if (EntityManager.TryGetComponent(itemMapper.Owner, out AppearanceComponent? appearanceComponent)
-                && TryGetLayers(message, itemMapper, out var containedLayers))
+            if(!Resolve(uid, ref itemMapper))
+                return;
+
+            if (EntityManager.TryGetComponent(uid, out AppearanceComponent? appearanceComponent)
+                && TryGetLayers(uid, itemMapper, out var containedLayers))
             {
-                _appearance.SetData(itemMapper.Owner, StorageMapVisuals.LayerChanged, new ShowLayerData(containedLayers), appearanceComponent);
+                _appearance.SetData(uid, StorageMapVisuals.LayerChanged, new ShowLayerData(containedLayers), appearanceComponent);
             }
         }
 
         /// <summary>
-        /// Method that iterates over storage of the entity in <paramref name="msg"/> and sets <paramref name="containedLayers"/> according to
+        /// Method that iterates over storage of the entity in <paramref name="uid"/> and sets <paramref name="containedLayers"/> according to
         /// <paramref name="itemMapper"/> definition. It will have O(n*m) time behavior (n - number of entities in container, and m - number of
         /// definitions in <paramref name="containedLayers"/>.
         /// </summary>
-        /// <param name="msg">event with EntityUid used to search the storage</param>
+        /// <param name="uid">EntityUid used to search the storage</param>
         /// <param name="itemMapper">component that contains definition used to map <see cref="Content.Shared.Whitelist.EntityWhitelist">whitelist</see> in
         /// <c>mapLayers</c> to string.
         /// </param>
         /// <param name="containedLayers">list of <paramref name="itemMapper"/> layers that should be visible</param>
         /// <returns>false if <c>msg.Container.Owner</c> is not a storage, true otherwise.</returns>
-        private bool TryGetLayers(ContainerModifiedMessage msg,
+        private bool TryGetLayers(EntityUid uid,
             ItemMapperComponent itemMapper,
             out List<string> showLayers)
         {
-            var containedLayers = _container.GetAllContainers(msg.Container.Owner)
+            var containedLayers = _container.GetAllContainers(uid)
                 .Where(c => itemMapper.ContainerWhitelist?.Contains(c.ID) ?? true).SelectMany(cont => cont.ContainedEntities).ToArray();
 
             var list = new List<string>();
             foreach (var mapLayerData in itemMapper.MapLayers.Values)
             {
-                var count = containedLayers.Count(uid => mapLayerData.ServerWhitelist.IsValid(uid));
+                var count = containedLayers.Count(ent => mapLayerData.ServerWhitelist.IsValid(ent));
                 if (count >= mapLayerData.MinCount && count <= mapLayerData.MaxCount)
                 {
                     list.Add(mapLayerData.Layer);