ScenePartition/Editor/ScenePartitionSO.cs

using System.Collections.Generic;
using System.IO;
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 ScenePartitionDataSO scenePartitionData = null;

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

        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);

            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<uint> ids = new SortedSet<uint>(scenePartitionData.scenePartitions.Keys);
            LoadScenePartitions(ids);
        }

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

            scenePartitionData.scenePartitions = new SerializableSortedList<uint, ScenePartition>();

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

        private void LoadScenePartitions(SortedSet<uint> partitionIds)
        {
            if (!scenePartitionData.hasCreatedPartitions) return;

            string scenePath = ScenePartitionUtils.GetScenePath(this);

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

            scenePartitionData.loadedScenePartitions.Clear();

            // Add default ids.
            for (int i = 0; i < alwaysLoadIds.Count; i++)
            {
                if (scenePartitionData.scenePartitions.ContainsKey(alwaysLoadIds[i]))
                {
                    partitionIds.Add(alwaysLoadIds[i]);
                }
            }

            // Create scene data.
            foreach (uint id in partitionIds)
            {
                ScenePartition p = scenePartitionData.scenePartitions[id];
                data.AddRange(p.data);
                scenePartitionData.loadedScenePartitions.Add(p.id, p);
            }

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

            AssetDatabase.Refresh();
        }

        /// <summary>
        /// Convert scene to partitions and save them to disk.
        /// </summary>
        public void Save()
        {
            if (!scenePartitionData.hasCreatedPartitions) return;
            if (!scenePartitionData.hasLoadedPartitions) return;

            DeleteLoadedPartitions(); // Delete the loaded partitions from disk so we can write the new ones.

            string dataPath = ScenePartitionUtils.GetDataPath(this);
            string scenePath = ScenePartitionUtils.GetScenePath(this);

            string[] data = File.ReadAllLines(scenePath);

            int lastIndex = data.Length;
            for (int i = data.Length - 1; i >= 0; i--)
            {
                if (data[i].StartsWith("---")) // --- is the start of a new yaml document.
                {
                    int idStartIndex = data[i].IndexOf(" &") + 2; // & is the start of the object id.
                    int idLength = data[i].Length;

                    File.WriteAllLines($"{dataPath}/{sceneName}-{data[i][idStartIndex..idLength]}.yaml", data[i..lastIndex]);

                    lastIndex = i;
                }
            }

            File.WriteAllLines($"{dataPath}/{sceneName}.yaml", data[0..lastIndex]);

            LoadAll(); // Reload.
        }

        /// <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);

            scenePartitionData.scenePartitions.Clear();

            AssetDatabase.Refresh();
        }

        private void DeleteLoadedPartitions()
        {
            if (!scenePartitionData.hasCreatedPartitions) return;

            foreach (KeyValuePair<uint, ScenePartition> scenePartition in scenePartitionData.scenePartitions)
            {
                if (!File.Exists(scenePartition.Value.filePath)) continue;

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

            scenePartitionData.loadedScenePartitions.Clear();
        }

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

            for (int i = 0; i < ids.Length; i++)
            {
                SortedSet<uint> connections = ScenePartitionUtils.FindDeeplyLinkedObjects(scenePartitionData.scenePartitions, ids[i]);

                foreach (uint c in connections)
                {
                    partitionIds.Add(c);
                }
            }

            LoadScenePartitions(partitionIds);
        }
    }
}