Add compatibility with SoftReferenceableAssets for customLocations

JotunnLib/Entities/CustomLocation.cs

@@ -4,6 +4,7 @@
 using Jotunn.Managers;
 using SoftReferenceableAssets;
 using UnityEngine;
+using Object = UnityEngine.Object;
 
 namespace Jotunn.Entities
 {
@@ -38,6 +39,12 @@ public class CustomLocation : CustomEntity
         ///     Indicator if references from <see cref="Entities.Mock{T}"/>s will be replaced at runtime.
         /// </summary>
         public bool FixReference { get; set; }
+
+        /// <summary>
+        ///     Indicator if location is added from SoftReferenceableAssets.<br />
+        ///     Used to delay mocking prefabs until ZoneSystem.SpawnLocation()
+        /// </summary>
+        public bool SoftReference { get; set; }
 
         /// <summary>
         ///     Custom location from a prefab with a <see cref="LocationConfig"/> attached.<br />
@@ -103,6 +110,38 @@ public CustomLocation(GameObject exteriorPrefab, GameObject interiorPrefab, bool
             FixReference = fixReference;
         }
 
+        /// <summary>
+        ///     Custom location from a prefab with a <see cref="LocationConfig"/> attached. Using SoftReference system.
+        /// </summary>
+        /// <param name="softReferencePrefab">The exterior prefab for this custom location.</param>
+        /// <param name="fixReference">If true references for <see cref="Entities.Mock{T}"/> objects get resolved at runtime by Jötunn.</param>
+        /// <param name="locationConfig">The <see cref="LocationConfig"/> for this custom location.</param>
+        public CustomLocation(SoftReference<GameObject> softReferencePrefab, bool fixReference, LocationConfig locationConfig) : base(Assembly.GetCallingAssembly())
+        {
+            if (!softReferencePrefab.IsValid)
+            {
+                Logger.LogError($"SoftReference invalid for prefab: {softReferencePrefab.Name}");
+                return;
+            }
+
+            var parent = ZoneManager.Instance.LocationContainer.transform;
+            AssetManager.Instance.ResolveMocksOnLoad(softReferencePrefab, parent, OnLocationResolve);
+            Name = softReferencePrefab.Name;
+            ZoneLocation = locationConfig.GetZoneLocation();
+            ZoneLocation.m_prefab = softReferencePrefab;
+            ZoneLocation.m_prefabName = softReferencePrefab.Name;
+            FixReference = fixReference;
+            SoftReference = true;
+        }
+
+        private void OnLocationResolve(GameObject gameObject)
+        {
+            if (gameObject.TryGetComponent<ZoneSystem.ZoneLocation>(out var zoneLocation))
+            {
+                ZoneManager.Instance.PrepareLocation(zoneLocation, SourceMod);
+            }
+        }
+
         /// <summary>
         ///     Helper method to determine if a location prefab with a given name is a custom location created with Jötunn.
         /// </summary>

JotunnLib/Managers/AssetManager.cs

@@ -32,6 +32,9 @@ public class AssetManager : IManager
         private Dictionary<Type, Dictionary<string, AssetID>> mapNameToAssetID;
         internal Dictionary<Type, Dictionary<string, AssetID>> MapNameToAssetID => mapNameToAssetID ??= CreateNameToAssetID();
 
+        private GameObject ResolvedAssetsContainer;
+        private Dictionary<AssetID, MockResolutionContext> assetsToResolve = new Dictionary<AssetID, MockResolutionContext>();
+
         /// <summary>
         ///     Hide .ctor
         /// </summary>
@@ -45,6 +48,11 @@ static AssetManager()
         void IManager.Init()
         {
             Main.LogInit(nameof(AssetManager));
+
+            ResolvedAssetsContainer = new GameObject("Resolved Assets");
+            ResolvedAssetsContainer.transform.parent = Main.RootObject.transform;
+            ResolvedAssetsContainer.SetActive(false);
+
             Main.Harmony.PatchAll(typeof(Patches));
         }
 
@@ -84,6 +92,31 @@ private static IEnumerable<CodeInstruction> AssetBundleLoader_GetAllAssetPathsMa
                     .SetInstruction(new CodeInstruction(OpCodes.Call, addSafeMethod))
                     .InstructionEnumeration();
             }
+
+            [HarmonyPatch(typeof(AssetLoader), nameof(AssetLoader.InvokeCallbacks)), HarmonyPrefix]
+            private static void SwapResolvedAsset(ref AssetLoader __instance, LoadResult result)
+            {
+                if (result == LoadResult.Succeeded)
+                {
+                    if (Instance.assetsToResolve.TryGetValue(__instance.m_assetID, out var mockedAsset))
+                    {
+                        mockedAsset.InstantiateAndResolveAsset(__instance.m_asset);
+                        __instance.m_asset = mockedAsset.Asset;
+                    }
+                }
+            }
+
+            [HarmonyPatch(typeof(AssetLoader), nameof(AssetLoader.Release)), HarmonyPostfix]
+            private static void AssetLoader_Release(ref AssetLoader __instance)
+            {
+                if (__instance.ReferenceCount == 0)
+                {
+                    if (Instance.assetsToResolve.TryGetValue(__instance.m_assetID, out var mockedAsset))
+                    {
+                        mockedAsset.DestroyAsset();
+                    }
+                }
+            }
         }
 
         /// <summary>
@@ -135,6 +168,59 @@ public AssetID AddAsset(Object asset)
             return AddAsset(asset, null);
         }
 
+        /// <summary>
+        ///     Registers an asset to be instantiated under the given parent and have its mock references resolved on load.<b/>
+        ///     Must be called before the asset is loaded the first time.
+        /// </summary>
+        /// <param name="assetID">The <see cref="AssetID"/> of the asset to instantiate and resolve mocks for on load</param>
+        public void ResolveMocksOnLoad(AssetID assetID)
+        {
+            ResolveMocksOnLoad(assetID, null, null);
+        }
+
+        /// <summary>
+        ///     Registers an asset to be instantiated under the given parent and have its mock references resolved on load.<b/>
+        ///     Must be called before the asset is loaded the first time.
+        /// </summary>
+        /// <param name="assetID">The <see cref="AssetID"/> of the asset to instantiate and resolve mocks for on load</param>
+        /// <param name="parent">Optional transform under which the asset will be instantiated, otherwise a default container is used</param>
+        public void ResolveMocksOnLoad(AssetID assetID, Transform parent)
+        {
+            ResolveMocksOnLoad(assetID, parent, null);
+        }
+
+        /// <summary>
+        ///     Registers an asset to be instantiated under the given parent and have its mock references resolved on load.<b/>
+        ///     Must be called before the asset is loaded the first time.
+        /// </summary>
+        /// <param name="softReference">The <see cref="SoftReference{T}"/> to instantiate and resolve mocks for on load</param>
+        /// <param name="parent">Optional transform under which the asset will be instantiated, otherwise a default container is used</param>
+        /// <param name="resolveCallback">Adds a callback when the asset was resolved and instantiated</param>
+        public void ResolveMocksOnLoad<T>(SoftReference<T> softReference, Transform parent, Action<T> resolveCallback) where T : Object
+        {
+            ResolveMocksOnLoad(softReference.m_assetID, parent, (asset) => resolveCallback?.Invoke(asset as T));
+        }
+
+        /// <summary>
+        ///     Registers an asset to be instantiated under the given parent and have its mock references resolved on load.<b/>
+        ///     Must be called before the asset is loaded the first time.
+        /// </summary>
+        /// <param name="assetID">The <see cref="AssetID"/> of the asset to instantiate and resolve mocks for on load</param>
+        /// <param name="parent">Optional transform under which the asset will be instantiated, otherwise a default container is used</param>
+        /// <param name="resolveCallback">Adds a callback when the asset was resolved and instantiated</param>
+        public void ResolveMocksOnLoad(AssetID assetID, Transform parent, Action<Object> resolveCallback)
+        {
+            if (assetsToResolve.TryGetValue(assetID, out var context))
+            {
+                context.Parent = parent ?? context.Parent ?? ResolvedAssetsContainer.transform;
+                context.ResolveCallback += resolveCallback;
+            }
+            else
+            {
+                assetsToResolve.Add(assetID, new MockResolutionContext(parent ?? ResolvedAssetsContainer.transform, resolveCallback));
+            }
+        }
+
         private static void AddAssetToBundleLoader(AssetBundleLoader assetBundleLoader, AssetID assetID, AssetRef assetRef)
         {
             // create fake bundle, since an AssetBundle can't be created at runtime
@@ -276,7 +362,7 @@ public SoftReference<Object> GetSoftReference(Type type, string name)
         /// <param name="name">Asset name to search for</param>
         /// <typeparam name="T">Asset type to search for</typeparam>
         /// <returns></returns>
-        public SoftReference<T> GetSoftReference<T>(string name) where T: Object
+        public SoftReference<T> GetSoftReference<T>(string name) where T : Object
         {
             AssetID assetID = GetAssetID<T>(name);
             return assetID.IsValid ? new SoftReference<T>(assetID) : default;
@@ -409,5 +495,52 @@ public AssetRef(BepInPlugin sourceMod, Object asset, Object original)
                 this.originalID = original && Instance.IsReady() ? Instance.GetAssetID(original.GetType(), original.name) : default;
             }
         }
+
+        internal class MockResolutionContext
+        {
+            public Object Asset { get; private set; }
+            public Transform Parent { get; set; }
+            public Action<Object> ResolveCallback { get; set; }
+
+            public MockResolutionContext(Transform parent, Action<Object> resolveCallback)
+            {
+                this.Parent = parent;
+                this.ResolveCallback += resolveCallback;
+            }
+
+            public bool IsResolved => (bool)Asset;
+
+            public void InstantiateAndResolveAsset(Object realAsset)
+            {
+                if (IsResolved)
+                {
+                    return;
+                }
+
+                Asset = Object.Instantiate(realAsset, Parent);
+                Asset.name = realAsset.name;
+
+                if (Asset is GameObject gameObject)
+                {
+                    gameObject.FixReferences(true);
+                }
+                else
+                {
+                    Asset.FixReferences();
+                }
+
+                ResolveCallback?.Invoke(Asset);
+            }
+
+            public void DestroyAsset()
+            {
+                if (Asset)
+                {
+                    Object.Destroy(Asset);
+                }
+
+                Asset = null;
+            }
+        }
     }
 }

JotunnLib/Managers/PrefabManager.cs

@@ -373,6 +373,12 @@ public void RegisterToZNetScene(GameObject gameObject)
             if (znet)
             {
                 string name = gameObject.name;
+
+                if (gameObject.name.StartsWith(MockManager.JVLMockPrefix))
+                {
+                    return;
+                }
+
                 int hash = name.GetStableHashCode();
 
                 if (znet.m_namedPrefabs.ContainsKey(hash))

JotunnLib/Managers/RenderManager.cs

@@ -451,8 +451,18 @@ private static List<Type> GetTopolocialSort(Transform tranform)
             HashSet<Type> visitedNodes = new HashSet<Type>();
             HashSet<Type> recursionStack = new HashSet<Type>();
 
+            if (!tranform || !tranform.gameObject)
+            {
+                return result;
+            }
+
             foreach (Component component in tranform.gameObject.GetComponents<Component>())
             {
+                if (!component)
+                {
+                    continue;
+                }
+
                 if (!TopologicalSortUtil(component.GetType(), visitedNodes, recursionStack, result))
                 {
                     Logger.LogWarning($"Cycles detected in component dependencies for type {component.GetType()}. Unable to determine deletion order for {tranform}.");

JotunnLib/Managers/ZoneManager.cs

@@ -236,11 +236,14 @@ public bool AddCustomLocation(CustomLocation customLocation)
                 return false;
             }
 
-            customLocation.Prefab.transform.SetParent(LocationContainer.transform);
-
-            // The root prefab needs to be active, otherwise ZNetViews are not prepared correctly
-            customLocation.Prefab.SetActive(true);
-
+            // Skip if location uses softReference system.
+            // If using softReference prefab is in assetbundle and therefore not immutable
+            if (!customLocation.SoftReference)
+            {
+                // The root prefab needs to be active, otherwise ZNetViews are not prepared correctly
+                customLocation.Prefab.SetActive(true);
+            }
+
             Locations.Add(customLocation.Name, customLocation);
             return true;
         }
@@ -504,24 +507,25 @@ private void RegisterLocations(ZoneSystem self)
                 {
                     try
                     {
-                        Logger.LogDebug(
-                            $"Adding custom location {customLocation} in {string.Join(", ", GetMatchingBiomes(customLocation.ZoneLocation.m_biome))}");
+                        Logger.LogDebug($"Adding custom location {customLocation} in {string.Join(", ", GetMatchingBiomes(customLocation.ZoneLocation.m_biome))}");
 
-                        // Fix references if needed
-                        if (customLocation.FixReference)
+                        if (customLocation.FixReference && !customLocation.SoftReference)
                         {
                             customLocation.Prefab.FixReferences(true);
                             customLocation.FixReference = false;
                         }
 
-                        var zoneLocation = customLocation.ZoneLocation;
+                        if (!customLocation.SoftReference)
+                        {
+                            PrepareLocation(customLocation.ZoneLocation, customLocation.SourceMod);
+                        }
 
-                        RegisterLocationInZoneSystem(self, zoneLocation, customLocation.SourceMod);
+                        RegisterLocationInZoneSystem(self, customLocation.ZoneLocation);
                     }
                     catch (Exception ex)
                     {
                         Logger.LogWarning(customLocation?.SourceMod, $"Exception caught while adding location: {ex}");
-                        toDelete.Add(customLocation.Name);
+                        toDelete.Add(customLocation?.Name);
                     }
                 }
 
@@ -581,30 +585,19 @@ private void RegisterVegetation(ZoneSystem self)
         ///     No mock references are fixed.
         /// </summary>
         /// <param name="zoneLocation"><see cref="ZoneLocation"/> to add to the <see cref="ZoneSystem"/></param>
-        public void RegisterLocationInZoneSystem(ZoneLocation zoneLocation) =>
-            RegisterLocationInZoneSystem(ZoneSystem.instance, zoneLocation, BepInExUtils.GetSourceModMetadata());
+        public void RegisterLocationInZoneSystem(ZoneLocation zoneLocation)
+        {
+            PrepareLocation(zoneLocation, BepInExUtils.GetSourceModMetadata());
+            RegisterLocationInZoneSystem(ZoneSystem.instance, zoneLocation);
+        }
 
-        /// <summary>
-        ///     Internal method for adding a ZoneLocation to a specific ZoneSystem.
-        /// </summary>
-        /// <param name="zoneSystem"><see cref="ZoneSystem"/> the location should be added to</param>
-        /// <param name="zoneLocation"><see cref="ZoneLocation"/> to add</param>
-        /// <param name="sourceMod"><see cref="BepInPlugin"/> which created the location</param>
-        private void RegisterLocationInZoneSystem(ZoneSystem zoneSystem, ZoneLocation zoneLocation, BepInPlugin sourceMod)
+        internal void PrepareLocation(ZoneLocation zoneLocation, BepInPlugin sourceMod)
         {
             zoneLocation.m_prefab.Load();
 
             foreach (var znet in global::Utils.GetEnabledComponentsInChildren<ZNetView>(zoneLocation.m_prefab.Asset))
             {
-                string prefabName = znet.GetPrefabName();
-                if (!ZNetScene.instance.m_namedPrefabs.ContainsKey(prefabName.GetStableHashCode()))
-                {
-                    var prefab = Object.Instantiate(znet.gameObject, PrefabManager.Instance.PrefabContainer.transform);
-                    prefab.name = prefabName;
-                    CustomPrefab customPrefab = new CustomPrefab(prefab, sourceMod);
-                    PrefabManager.Instance.AddPrefab(customPrefab);
-                    PrefabManager.Instance.RegisterToZNetScene(customPrefab.Prefab);
-                }
+                RegisterUnknownPrefab(sourceMod, znet);
             }
 
             RandomSpawn[] randomSpawns = global::Utils.GetEnabledComponentsInChildren<RandomSpawn>(zoneLocation.m_prefab.Asset);
@@ -615,17 +608,30 @@ private void RegisterLocationInZoneSystem(ZoneSystem zoneSystem, ZoneLocation zo
 
             foreach (var znet in randomSpawns.SelectMany(x => x.m_childNetViews))
             {
-                string prefabName = znet.GetPrefabName();
-                if (!ZNetScene.instance.m_namedPrefabs.ContainsKey(prefabName.GetStableHashCode()))
-                {
-                    var prefab = Object.Instantiate(znet.gameObject, PrefabManager.Instance.PrefabContainer.transform);
-                    prefab.name = prefabName;
-                    CustomPrefab customPrefab = new CustomPrefab(prefab, sourceMod);
-                    PrefabManager.Instance.AddPrefab(customPrefab);
-                    PrefabManager.Instance.RegisterToZNetScene(customPrefab.Prefab);
-                }
+                RegisterUnknownPrefab(sourceMod, znet);
             }
+        }
 
+        private static void RegisterUnknownPrefab(BepInPlugin sourceMod, ZNetView znet) {
+            string prefabName = znet.GetPrefabName();
+
+            if (prefabName.StartsWith(MockManager.JVLMockPrefix))
+            {
+                return;
+            }
+
+            if (!ZNetScene.instance.m_namedPrefabs.ContainsKey(prefabName.GetStableHashCode()))
+            {
+                var prefab = Object.Instantiate(znet.gameObject, PrefabManager.Instance.PrefabContainer.transform);
+                prefab.name = prefabName;
+                CustomPrefab customPrefab = new CustomPrefab(prefab, sourceMod);
+                PrefabManager.Instance.AddPrefab(customPrefab);
+                PrefabManager.Instance.RegisterToZNetScene(customPrefab.Prefab);
+            }
+        }
+
+        private void RegisterLocationInZoneSystem(ZoneSystem zoneSystem, ZoneLocation zoneLocation)
+        {
             if (!zoneSystem.m_locationsByHash.ContainsKey(zoneLocation.Hash))
             {
                 zoneSystem.m_locationsByHash.Add(zoneLocation.Hash, zoneLocation);