New requirements fullfilled

Author: Extrys

com.unity.netcode.gameobjects/Runtime/Core/NetworkObject.cs

@@ -1852,31 +1852,6 @@ internal void SpawnInternal(bool destroyWithScene, ulong ownerClientId, bool pla
             }
         }
 
-        /// <summary>
-        /// Serializes the given instantiation data and stores it to be used during the spawn process
-        /// by a compatible <see cref="INetworkPrefabInstanceHandlerWithData{T}"/>
-        /// </summary>
-        public void InjectInstantiationData<T>(T data) where T : struct, INetworkSerializable
-        {
-            if (!NetworkManager.PrefabHandler.TryGetHandlerWithData(GlobalObjectIdHash, out var prefabHandler) || !prefabHandler.HandlesDataType<T>())
-            {
-                throw new Exception("[InstantiationData] Cannot inject data: no compatible handler found for the specified data type.");
-            }
-
-            using var writer = new FastBufferWriter(4, Collections.Allocator.Temp, int.MaxValue);
-            var serializer = new BufferSerializer<BufferSerializerWriter>(new BufferSerializerWriter(writer));
-
-            try
-            {
-                data.NetworkSerialize(serializer);
-                InstantiationData = writer.ToArray();
-            }
-            catch (Exception ex)
-            {
-                NetworkLog.LogError($"[InstantiationData] Failed to serialize instantiation data for {nameof(NetworkObject)} '{gameObject.name}': {ex}");
-            }
-        }
-
         /// <summary>
         /// This invokes <see cref="NetworkSpawnManager.InstantiateAndSpawn(NetworkObject, ulong, bool, bool, bool, Vector3, Quaternion)"/>.
         /// </summary>

com.unity.netcode.gameobjects/Runtime/Spawning/INetworkPrefabInstanceHandlerWithData.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using Unity.Collections;
 using UnityEngine;
 
 namespace Unity.Netcode
@@ -8,47 +9,39 @@ namespace Unity.Netcode
     /// Specialized version of <see cref="INetworkPrefabInstanceHandler"/> that receives
     /// custom instantiation data injected by the server before spawning.
     /// </summary>
-    public interface INetworkPrefabInstanceHandlerWithData<T> : INetworkPrefabInstanceHandlerWithData where T : struct, INetworkSerializable
+    public interface INetworkPrefabInstanceHandlerWithData<T> where T : struct, INetworkSerializable
     {
-        // Propagating custom data through the entire spawn pipeline would add complexity and reduce modularity.
-        // To work around this, deserialization injects the data into this static map, keyed by the handler instance.
-        // The handler then reads the data at instantiation time, keeping the interface clean and stateless from the user's perspective.
-        // This avoids requiring implementers to store the data explicitly or declare additional fields.
-        static readonly Dictionary<INetworkPrefabInstanceHandlerWithData, T> _handlerToData = new();
+        NetworkObject Instantiate(ulong ownerClientId, Vector3 position, Quaternion rotation, T instantiationData);
+        void Destroy(NetworkObject networkObject);
 
-        NetworkObject INetworkPrefabInstanceHandler.Instantiate(ulong ownerClientId, Vector3 position, Quaternion rotation) => Instantiate(ownerClientId, position, rotation, _handlerToData[this]);
-        void INetworkPrefabInstanceHandlerWithData.RemoveDataEntry(INetworkPrefabInstanceHandlerWithData instance) => _handlerToData.Remove(instance);
-        bool INetworkPrefabInstanceHandlerWithData.HandlesDataType<U>() => typeof(T) == typeof(U);
-        void INetworkPrefabInstanceHandlerWithData.ReadInstantiationData<RW>(ref BufferSerializer<RW> serializer)
+        public static void InjectInstantiationDataForObject(GameObject prefab, T instantiationData)
         {
-            _handlerToData.TryGetValue(this, out var value);
-            serializer.SerializeValue(ref value);
-            _handlerToData[this] = value;
+            if (prefab.TryGetComponent<NetworkObject>(out var networkObject))
+            {
+                networkObject.NetworkManager.PrefabHandler.InjectInstantiationData(prefab, instantiationData);
+            }
+            else
+            {
+                Debug.LogError($"Prefab {prefab.name} does not have a NetworkObject component.");
+            }
         }
-
-        NetworkObject Instantiate(ulong ownerClientId, Vector3 position, Quaternion rotation, T instantiationData);
     }
 
-    /// <summary>
-    /// Internal use only. Do not implement directly. Use <see cref="INetworkPrefabInstanceHandlerWithData{T}"/> instead.
-    /// </summary>
-    public interface INetworkPrefabInstanceHandlerWithData : INetworkPrefabInstanceHandler
+    internal interface INetworkPrefabInstanceHandlerWithData : INetworkPrefabInstanceHandler
     {
-        /// <summary>
-        /// Invoked during deserialization to read the instantiation data associated with this prefab instance.
-        /// </summary>
-        void ReadInstantiationData<T>(ref BufferSerializer<T> serializer) where T : IReaderWriter;
+        bool HandlesDataType<T>();
+        void ReadInstantiationData<TReaderWriter>(ref BufferSerializer<TReaderWriter> serializer) where TReaderWriter : IReaderWriter;
+    }
 
-        /// <summary>
-        /// Removes the data entry for the given instance.
-        /// Is important to call this when the instance isnt referenced to avoid memory leaks.
-        /// </summary>
-        /// <param name="instance"></param>
-        void RemoveDataEntry(INetworkPrefabInstanceHandlerWithData instance);
+    internal class HandlerWrapper<T> : INetworkPrefabInstanceHandlerWithData where T : struct, INetworkSerializable
+    {
+        private readonly INetworkPrefabInstanceHandlerWithData<T> _impl;
+        private T _payload;
 
-        /// <summary>
-        /// Returns true if <typeparamref name="T"/> matches the expected instantiation data type for this handler.
-        /// </summary>
-        bool HandlesDataType<T>() where T : struct, INetworkSerializable;
+        public HandlerWrapper(INetworkPrefabInstanceHandlerWithData<T> impl) => _impl = impl;
+        public bool HandlesDataType<U>() => typeof(T) == typeof(U);
+        public void ReadInstantiationData<TReaderWriter>(ref BufferSerializer<TReaderWriter> serializer) where TReaderWriter : IReaderWriter => serializer.SerializeValue(ref _payload);
+        public NetworkObject Instantiate(ulong ownerClientId, Vector3 position, Quaternion rotation) => _impl.Instantiate(ownerClientId, position, rotation, _payload);
+        public void Destroy(NetworkObject networkObject) => _impl.Destroy(networkObject);
     }
-}
+}

com.unity.netcode.gameobjects/Runtime/Spawning/NetworkPrefabHandler.cs

@@ -84,6 +84,10 @@ public bool AddHandler(GameObject networkPrefabAsset, INetworkPrefabInstanceHand
         {
             return AddHandler(networkPrefabAsset.GetComponent<NetworkObject>().GlobalObjectIdHash, instanceHandler);
         }
+        public bool AddHandler<T>(GameObject networkPrefabAsset, INetworkPrefabInstanceHandlerWithData<T> instanceHandler) where T : struct, INetworkSerializable
+        {
+            return AddHandler(networkPrefabAsset.GetComponent<NetworkObject>().GlobalObjectIdHash, instanceHandler);
+        }
 
         /// <summary>
         /// Use a  <see cref="NetworkObject"/> to register a class that implements the <see cref="INetworkPrefabInstanceHandler"/> interface with the <see cref="NetworkPrefabHandler"/>
@@ -95,6 +99,10 @@ public bool AddHandler(NetworkObject prefabAssetNetworkObject, INetworkPrefabIns
         {
             return AddHandler(prefabAssetNetworkObject.GlobalObjectIdHash, instanceHandler);
         }
+        public bool AddHandler<T>(NetworkObject prefabAssetNetworkObject, INetworkPrefabInstanceHandlerWithData<T> instanceHandler) where T : struct, INetworkSerializable
+        {
+            return AddHandler(prefabAssetNetworkObject.GlobalObjectIdHash, instanceHandler);
+        }
 
         /// <summary>
         /// Use a <see cref="NetworkObject.GlobalObjectIdHash"/> to register a class that implements the <see cref="INetworkPrefabInstanceHandler"/> interface with the <see cref="NetworkPrefabHandler"/>
@@ -116,6 +124,38 @@ public bool AddHandler(uint globalObjectIdHash, INetworkPrefabInstanceHandler in
 
             return false;
         }
+        public bool AddHandler<T>(uint globalObjectIdHash, INetworkPrefabInstanceHandlerWithData<T> instanceHandler) where T : struct, INetworkSerializable
+        {
+            if (!m_PrefabAssetToPrefabHandler.ContainsKey(globalObjectIdHash))
+                return AddHandler(globalObjectIdHash, new HandlerWrapper<T>(instanceHandler));
+            return false;
+        }
+
+        public void InjectInstantiationData<T>(GameObject gameObject, T instantiationData) where T : struct, INetworkSerializable
+        {
+            if (gameObject.TryGetComponent<NetworkObject>(out var networkObject))
+                InjectInstantiationData(networkObject,instantiationData);
+        }
+        public void InjectInstantiationData<T>(NetworkObject networkObject, T data) where T : struct, INetworkSerializable
+        {
+            if (!TryGetHandlerWithData(networkObject.GlobalObjectIdHash, out var prefabHandler) || !prefabHandler.HandlesDataType<T>())
+            {
+                throw new Exception("[InstantiationData] Cannot inject data: no compatible handler found for the specified data type.");
+            }
+
+            using var writer = new FastBufferWriter(4, Collections.Allocator.Temp, int.MaxValue);
+            var serializer = new BufferSerializer<BufferSerializerWriter>(new BufferSerializerWriter(writer));
+
+            try
+            {
+                data.NetworkSerialize(serializer);
+                networkObject.InstantiationData = writer.ToArray();
+            }
+            catch (Exception ex)
+            {
+                NetworkLog.LogError($"[InstantiationData] Failed to serialize instantiation data for {nameof(NetworkObject)} '{networkObject.name}': {ex}");
+            }
+        }
 
         /// <summary>
         /// HOST ONLY!
@@ -214,7 +254,6 @@ public bool RemoveHandler(uint globalObjectIdHash)
 
             if(m_PrefabAssetToPrefabHandlerWithData.TryGetValue(globalObjectIdHash, out var handlerWithData))
             {
-                handlerWithData.RemoveDataEntry(handlerWithData);
                 m_PrefabAssetToPrefabHandlerWithData.Remove(globalObjectIdHash);
             }
             return m_PrefabAssetToPrefabHandler.Remove(globalObjectIdHash);

com.unity.netcode.gameobjects/Tests/Runtime/Prefabs/NetworkPrefabHandlerWithDataTests.cs

@@ -107,7 +107,7 @@ private void RegisterPrefab(NetworkManager manager, out PrefabInstanceHandlerWit
         private NetworkObject SpawnPrefabWithData(NetworkSerializableTest data)
         {
             var instance = GameObject.Instantiate(_prefab).GetComponent<NetworkObject>();
-            instance.InjectInstantiationData(data);
+            server.PrefabHandler.InjectInstantiationData(instance, data);
             instance.Spawn();
             return instance;
         }