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<ulong> alwaysLoadIds = new List<ulong> { 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<ulong> ids = new SortedSet<ulong>(Data.ScenePartitions.Keys);
            LoadScenePartitions(ids);
        }

        /// <summary>
        /// Discard changes and reload loaded partitions.
        /// </summary>
        public void Reload()
        {
            if (!Data.HasLoadedPartitions) return;
            LoadScenePartitions(new SortedSet<ulong>(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<ulong> partitionIds)
        {
            using (new ProfilerUtility.ProfilerScope($"{nameof(LoadScenePartitions)}"))
            {
                if (!Data.HasCreatedPartitions) return;

                string scenePath = ScenePartitionUtils.GetScenePath(this);

                Data.LoadedScenePartitions.Clear();

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

                // Clear file.
                File.WriteAllText(scenePath, string.Empty);

                // Create scene data.
                try
                {
                    using (FileStream outputStream = new FileStream(scenePath, FileMode.Append, FileAccess.Write))
                    {
                        foreach (ulong id in partitionIds)
                        {
                            ScenePartition p = Data.ScenePartitions[id];

                            using (FileStream inputStream = new FileStream(p.filePath, FileMode.Open, FileAccess.Read))
                            {
                                byte[] buffer = new byte[4096];
                                int bytesRead;
                                while ((bytesRead = inputStream.Read(buffer, 0, buffer.Length)) > 0)
                                {
                                    outputStream.Write(buffer, 0, bytesRead);
                                }
                            }

                            Data.LoadedScenePartitions.Add(p.id, p);
                        }
                    }
                }
                catch (System.Exception ex)
                {
                    Debug.LogException(ex);
                }

                AssetDatabase.Refresh();
            }
        }

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

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

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

                Dictionary<ulong, string> sceneObjectNameById = GetSceneObjectNames(in sceneData);

                // 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], objectIdPattern);

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

                            string extraInfo = "";
                            if (TryGetObjectInfo(in sceneData, in sceneObjectNameById, i, lastIndex, objectId, out string objectInfo))
                            {
                                extraInfo = $"-{objectInfo}";
                            }

                            // Write data to disk.
                            File.WriteAllLines($"{dataPath}/{SceneName}-{id}{extraInfo}.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()
        {
            using (new ProfilerUtility.ProfilerScope($"{nameof(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()
        {
            using (new ProfilerUtility.ProfilerScope($"{nameof(DeleteLoadedPartitions)}"))
            {
                if (!Data.HasLoadedPartitions) return;

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

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

        public void LoadPartitions(ulong[] ids)
        {
            using (new ProfilerUtility.ProfilerScope($"{nameof(LoadPartitions)}"))
            {
                SortedSet<ulong> partitionIds = new SortedSet<ulong>();

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

                LoadScenePartitions(partitionIds);
            }
        }

        public void LoadPartitionsAdditive(ulong[] ids)
        {
            using (new ProfilerUtility.ProfilerScope($"{nameof(LoadPartitions)}"))
            {
                SortedSet<ulong> partitionIds = new SortedSet<ulong>();

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

                // Partitions already loaded.
                if (Data.HasLoadedPartitions)
                {
                    foreach (KeyValuePair<ulong, ScenePartition> item in Data.LoadedScenePartitions)
                    {
                        partitionIds.Add(item.Key);
                    }
                }

                LoadScenePartitions(partitionIds);
            }
        }

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

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

            return partitionIds;
        }

        public void GenerateSceneGridData()
        {
            using (new ProfilerUtility.ProfilerScope($"{nameof(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(ids[i].targetObjectId, rootGameObjects[i].transform.position, ScenePartitionSceneViewEditor.cellSize);
                    }
                    else
                    {
                        Data.SceneGrid.Insert(ids[i].targetPrefabId, rootGameObjects[i].transform.position, ScenePartitionSceneViewEditor.cellSize);
                    }
                }

                Unload();
            }
        }

        public void LoadCell(int gridId)
        {
            if (Data.SceneGrid.Grid.TryGetValue(gridId, out GridList ids))
            {
                LoadPartitions(ids.list.ToArray());
            }
        }

        public void LoadCellAdditive(int gridId)
        {
            if (Data.SceneGrid.Grid.TryGetValue(gridId, out GridList ids))
            {
                LoadPartitionsAdditive(ids.list.ToArray());
            }
        }

        public void ClearCache()
        {
            data = null;
        }

        private Dictionary<ulong, string> GetSceneObjectNames(in string[] sceneData)
        {
            using ProfilerUtility.ProfilerScope profilerScope = new(nameof(GetSceneObjectNames));

            Dictionary<ulong, string> sceneObjectNameById = new Dictionary<ulong, string>();

            const string objectIdPattern = @"&(\d+)";

            int lastIndex = sceneData.Length;
            for (int i = sceneData.Length - 1; i >= 0; i--)
            {
                bool foundName = false;

                {   // GameObjects.
                    const string gameObjectHeaderPattern = "--- !u!1 &";
                    if (!foundName && sceneData[i].StartsWith(gameObjectHeaderPattern))
                    {

                        Match match = Regex.Match(sceneData[i], objectIdPattern);
                        if (!match.Success) continue;
                        if (!ulong.TryParse(match.Groups[1].Value, out ulong id)) continue;

                        for (int j = i; j < lastIndex; j++)
                        {
                            const string namePattern = "m_Name: ";
                            int nameStartIndex = sceneData[j].LastIndexOf(namePattern);
                            if (nameStartIndex < 0) continue;

                            nameStartIndex += namePattern.Length;

                            string name = sceneData[j][nameStartIndex..];
                            sceneObjectNameById.Add(id, name);
                            foundName = true;
                            break;
                        }
                    }
                }

                {   // Prefabs.
                    const string prefabHeaderPattern = "--- !u!1001 &";
                    if (!foundName && sceneData[i].StartsWith(prefabHeaderPattern))
                    {
                        Match match = Regex.Match(sceneData[i], objectIdPattern);
                        if (!match.Success) continue;
                        if (!ulong.TryParse(match.Groups[1].Value, out ulong id)) continue;

                        // Get name form property override in Scene.
                        for (int j = i; j < lastIndex; j++)
                        {
                            const string propertyPattern = "propertyPath: m_Name";
                            int propertyStartIndex = sceneData[j].LastIndexOf(propertyPattern);
                            if (propertyStartIndex < 0) continue;

                            const string valuePattern = "value: ";
                            int valueStartIndex = sceneData[j + 1].LastIndexOf(valuePattern);
                            if (valueStartIndex < 0) continue;

                            valueStartIndex += valuePattern.Length;

                            string name = sceneData[j + 1][valueStartIndex..];
                            sceneObjectNameById.Add(id, name);
                            foundName = true;
                            break;
                        }

                        // Get name from prefab in AssetDatabase.
                        if (!foundName)
                        {
                            // Find the match using regex
                            const string guidPattern = @"guid:\s(\w+)";
                            Match guidMatch = Regex.Match(sceneData[lastIndex - 1], guidPattern);

                            if (guidMatch.Success)
                            {
                                // Extract the GUID
                                string guid = guidMatch.Groups[1].Value;

                                string path = AssetDatabase.GUIDToAssetPath(guid);
                                string name = Path.GetFileNameWithoutExtension(path);
                                sceneObjectNameById.Add(id, name);
                                foundName = true;
                            }
                        }
                    }
                }

                if (sceneData[i].StartsWith("---"))
                {
                    lastIndex = i;
                    continue;
                }
            }

            return sceneObjectNameById;
        }

        private bool TryGetObjectTypeName(string data, out string objectTypeName)
        {
            objectTypeName = data.Remove(data.Length - 1);
            return true;
        }

        private bool TryGetObjectInfo(in string[] sceneData, in Dictionary<ulong, string> sceneObjectNameById, int index, int lastIndex, ulong objectId, out string objectInfo)
        {
            using ProfilerUtility.ProfilerScope profilerScope = new(nameof(TryGetObjectInfo));

            objectInfo = "";

            // Try get scene object name.
            bool foundSceneObjectName = false;

            // If it is a gameObject or prefab try to get the name directly.
            const string gameObjectHeaderPattern = "--- !u!1 &";
            const string prefabHeaderPattern = "--- !u!1001 &";
            if (!foundSceneObjectName && sceneData[index].StartsWith(gameObjectHeaderPattern) || sceneData[index].StartsWith(prefabHeaderPattern))
            {
                if (sceneObjectNameById.TryGetValue(objectId, out string sceneObjectName))
                {
                    objectInfo += $"-{sceneObjectName}";
                    foundSceneObjectName = true;
                }
            }

            if (!foundSceneObjectName)
            {
                // Most components have a property that links to the GameObject.
                for (int j = index; j < lastIndex; j++)
                {
                    const string gameObjectPattern = @"m_GameObject:\s{fileID:\s(\d+)}";

                    // Find the match using regex
                    Match gameObjectMatch = Regex.Match(sceneData[j], gameObjectPattern);

                    if (!gameObjectMatch.Success) continue;

                    if (ulong.TryParse(gameObjectMatch.Groups[1].Value, out ulong fileNumber))
                    {
                        if (fileNumber == 0) break; // 0 == nothing.
                        if (!sceneObjectNameById.TryGetValue(fileNumber, out string sceneObjectName)) break;

                        objectInfo += $"-{sceneObjectName}";
                        foundSceneObjectName = true;
                    }

                    break;
                }
            }

            if (!foundSceneObjectName)
            {
                // Some have a property that links to the PrefabInstance.
                for (int j = index; j < lastIndex; j++)
                {
                    const string prefabInstancePattern = @"m_PrefabInstance:\s{fileID:\s(\d+)}";

                    // Find the match using regex
                    Match prefabInstanceMatch = Regex.Match(sceneData[j], prefabInstancePattern);

                    if (!prefabInstanceMatch.Success) continue;

                    if (ulong.TryParse(prefabInstanceMatch.Groups[1].Value, out ulong fileNumber))
                    {
                        if (fileNumber == 0) break; // 0 == nothing.
                        if (!sceneObjectNameById.TryGetValue(fileNumber, out string sceneObjectName)) break;

                        objectInfo += $"-{sceneObjectName}";
                        foundSceneObjectName = true;
                    }

                    break;
                }
            }

            bool foundObjectTypeName = false;

            // Try get object type name.
            if (TryGetObjectTypeName(sceneData[index + 1], out string objectTypeName))
            {
                objectInfo += $"-{objectTypeName}";
                foundObjectTypeName = true;
            }

            return foundSceneObjectName || foundObjectTypeName;
        }
    }
}