ScenePartition/Editor/ScenePartitionSO.cs

using System.Collections.Generic;
using System.IO;
using System.Text.RegularExpressions;
using UnityEditor;
using UnityEditor.SceneManagement;
using UnityEngine;
using UnityEngine.SceneManagement;

namespace VertexColor.ScenePartition.Editor
{
    [CreateAssetMenu(fileName = "Scene", menuName = "Max/ScenePartitionSO")]
    public class ScenePartitionSO : ScriptableObject
    {
        [field: SerializeField]
        public SceneAsset SceneAsset { get; private set; } = null;
        public string SceneName => SceneAsset == null ? name : SceneAsset.name;

        public List<long> alwaysLoadIds = new List<long> { 0, 1, 2, 3, 4 };

        public ScenePartitionData Data
        {
            get
            {
                // Load data from the ScriptableSingleton.
                data ??= ScenePartitionSS.instance.GetScenePartitionData(this);
                return data;
            }
        }
        private ScenePartitionData data = null;

        public void CreateScene()
        {
            if (SceneAsset != null) return;

            string scenePath = Path.Combine(AssetDatabase.GetAssetPath(this), $"../{this.name}.unity");

            Scene scene = EditorSceneManager.NewScene(NewSceneSetup.EmptyScene, NewSceneMode.Single);
            EditorSceneManager.SaveScene(scene, scenePath);
            Save();

            AssetDatabase.SaveAssets();
            AssetDatabase.Refresh();

            SceneAsset = AssetDatabase.LoadAssetAtPath<SceneAsset>(scenePath);

            AssetDatabase.SaveAssets();
            AssetDatabase.Refresh();
        }

        /// <summary>
        /// Load partitions from disk and construct the scene file.
        /// </summary>
        public void LoadAll()
        {
            CreateScenePartitions();
            SortedSet<long> ids = new SortedSet<long>(Data.ScenePartitions.Keys);
            LoadScenePartitions(ids);
        }

        /// <summary>
        /// Discard changes and reload loaded partitions.
        /// </summary>
        public void Reload()
        {
            if (!Data.HasLoadedPartitions) return;
            LoadScenePartitions(new SortedSet<long>(Data.LoadedScenePartitions.Keys));
        }

        private void CreateScenePartitions()
        {
            string dataPath = ScenePartitionUtils.GetDataPath(this);
            string[] files = Directory.GetFiles(dataPath);

            Data.ScenePartitions = new ScenePartitionSortedList();

            for (int i = 0; i < files.Length; i++)
            {
                ScenePartition scenePartition = new ScenePartition(files[i]);
                Data.ScenePartitions.Add(scenePartition.id, scenePartition);
            }
        }

        private void LoadScenePartitions(SortedSet<long> partitionIds)
        {
            if (!Data.HasCreatedPartitions) return;

            string scenePath = ScenePartitionUtils.GetScenePath(this);

            List<string> sceneData = new List<string>();

            Data.LoadedScenePartitions.Clear();

            // Add always load ids.
            SortedSet<long> baseIds = GetAlwaysLoadIds();
            foreach (var id in baseIds)
            {
                partitionIds.Add(id);
            }

            // Create scene data.
            foreach (long id in partitionIds)
            {
                ScenePartition p = Data.ScenePartitions[id];
                sceneData.AddRange(File.ReadAllLines(p.filePath));
                Data.LoadedScenePartitions.Add(p.id, p);
            }

            // Create scene.
            File.WriteAllLines(scenePath, sceneData);

            AssetDatabase.Refresh();
        }

        /// <summary>
        /// Convert scene to partitions and save them to disk.
        /// </summary>
        public void Save()
        {
            DeleteLoadedPartitions(); // Delete the loaded partitions from disk so we can write the new ones.

            string pattern = @"&(\d+)";
            string dataPath = ScenePartitionUtils.GetDataPath(this);
            string scenePath = ScenePartitionUtils.GetScenePath(this);

            // Read the data from the scene file.
            string[] sceneData = File.ReadAllLines(scenePath);

            // Split it into blocks.
            int lastIndex = sceneData.Length;
            for (int i = sceneData.Length - 1; i >= 0; i--)
            {
                if (sceneData[i].StartsWith("---")) // --- is the start of a new yaml document.
                {
                    Match match = Regex.Match(sceneData[i], pattern);

                    if (match.Success)
                    {
                        // Extract the file number
                        string id = match.Groups[1].Value;

                        // Write data to disk.
                        File.WriteAllLines($"{dataPath}/{SceneName}-{id}.yaml", sceneData[i..lastIndex]);
                    }

                    lastIndex = i;
                }
            }

            // Write header to disk.
            File.WriteAllLines($"{dataPath}/{SceneName}.yaml", sceneData[0..lastIndex]);
        }

        /// <summary>
        /// Empty the scene and save it (so it has no changes in source control).
        /// </summary>
        public void Unload()
        {
            string dataPath = ScenePartitionUtils.GetDataPath(this);
            string scenePath = ScenePartitionUtils.GetScenePath(this);

            Scene scene = EditorSceneManager.OpenScene(scenePath, OpenSceneMode.Single);

            GameObject[] gameObjects = scene.GetRootGameObjects();

            for (int i = gameObjects.Length - 1; i >= 0; i--)
            {
                DestroyImmediate(gameObjects[i]);
            }

            EditorSceneManager.SaveScene(scene);

            Data.LoadedScenePartitions.Clear();

            AssetDatabase.Refresh();
        }

        private void DeleteLoadedPartitions()
        {
            if (!Data.HasLoadedPartitions) return;

            foreach (KeyValuePair<long, ScenePartition> scenePartition in Data.LoadedScenePartitions)
            {
                if (!File.Exists(scenePartition.Value.filePath)) continue;

                File.Delete(scenePartition.Value.filePath);
            }
        }

        public void LoadPartitions(long[] ids)
        {
            SortedSet<long> partitionIds = new SortedSet<long>();

            for (int i = 0; i < ids.Length; i++)
            {
                SortedSet<long> connections = ScenePartitionUtils.FindDeeplyLinkedObjects(Data.ScenePartitions, ids[i]);
                foreach (long c in connections)
                {
                    partitionIds.Add(c);
                }
            }

            LoadScenePartitions(partitionIds);
        }

        private SortedSet<long> GetAlwaysLoadIds()
        {
            SortedSet<long> partitionIds = new SortedSet<long>();

            foreach (long id in alwaysLoadIds)
            {
                SortedSet<long> connections = ScenePartitionUtils.FindDeeplyLinkedObjects(Data.ScenePartitions, id);
                foreach (long c in connections)
                {
                    partitionIds.Add(c);
                }
            }

            return partitionIds;
        }

        public void GenerateSceneGridData()
        {
            LoadAll();

            if (!Data.HasCreatedPartitions) return;

            Scene scene = EditorSceneManager.OpenScene(ScenePartitionUtils.GetScenePath(this), OpenSceneMode.Single);

            GameObject[] rootGameObjects = scene.GetRootGameObjects();

            Data.SceneGrid.Grid.Clear();

            //// Maybe later switch to getting the data from disk instead of loading the scene and then unloading it again.
            //foreach (GameObject gameObject in rootGameObjects)
            //{
            //    // https://forum.unity.com/threads/how-to-get-the-local-identifier-in-file-for-scene-objects.265686/
            //    PropertyInfo inspectorModeInfo = typeof(SerializedObject).GetProperty("inspectorMode", BindingFlags.NonPublic | BindingFlags.Instance);
            //    SerializedObject serializedObject = new SerializedObject(gameObject.transform);
            //    inspectorModeInfo.SetValue(serializedObject, InspectorMode.Debug, null);
            //    SerializedProperty localIdProp = serializedObject.FindProperty("m_LocalIdentfierInFile");

            //    long localId = localIdProp.longValue;

            //    if (localId == 0)
            //    {
            //        if (PrefabUtility.IsPartOfPrefabInstance(gameObject))
            //        {
            //            GlobalObjectId id = GlobalObjectId.GetGlobalObjectIdSlow(gameObject); // We could use this funtion for all objects. Might be a bit slower but is also simple.
            //            localId = long.Parse(id.targetPrefabId.ToString());

            //            Debug.Log($"{id.assetGUID} | {id.identifierType} | {id.targetObjectId} | {id.targetPrefabId}");

            //            if (id.targetObjectId == 0 && id.targetPrefabId == 0)
            //            {
            //                Debug.LogWarning($"Could not find LocalIdentfierInFile for {gameObject.transform} {gameObject.name} {gameObject.transform.GetInstanceID()}");
            //                continue;
            //            }
            //        }
            //        else
            //        {
            //            Debug.LogWarning($"Could not find LocalIdentfierInFile for {gameObject.transform} {gameObject.name} {gameObject.transform.GetInstanceID()}");
            //            continue;
            //        }
            //    }

            //    if (!Data.ScenePartitions.ContainsKey(localId))
            //    {
            //        Debug.LogWarning($"Could not find LocalIdentfierInFile for {gameObject.transform} {gameObject.name} {gameObject.transform.GetInstanceID()}");
            //        continue;
            //    }

            //    Data.SceneGrid.Insert(localId, gameObject.transform.position);
            //}

            GlobalObjectId[] ids = new GlobalObjectId[rootGameObjects.Length];
            GlobalObjectId.GetGlobalObjectIdsSlow(rootGameObjects, ids);

            for (int i = 0; i < ids.Length; i++)
            {
                Debug.Log($"{ids[i].assetGUID} | {ids[i].identifierType} | {ids[i].targetObjectId} | {ids[i].targetPrefabId}");

                if (ids[i].targetPrefabId == 0) // 0 = no prefab.
                {
                    Data.SceneGrid.Insert((long)ids[i].targetObjectId, rootGameObjects[i].transform.position);
                }
                else
                {
                    Data.SceneGrid.Insert((long)ids[i].targetPrefabId, rootGameObjects[i].transform.position);
                }
            }

            Unload();
        }

        public void LoadCell(int gridId)
        {
            if (Data.SceneGrid.Grid.TryGetValue(gridId, out List<long> ids))
            {
                LoadPartitions(ids.ToArray());
            }
        }

        public void ClearCache()
        {
            data = null;
        }
    }
}