plugin-library/Assets/00.StaryEvo/Editor/YooAsset/ShaderVariantCollector.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.IO;
using UnityEngine;
using UnityEditor;
using UnityEditor.SceneManagement;
using YooAsset.Editor;

public static class ShaderVariantCollector
{
    private enum ESteps
    {
        None,
        Prepare,
        CollectAllMaterial,
        CollectVariants,
        CollectSleeping,
        WaitingDone,
    }

    private const float WaitMilliseconds = 3000f;

    private const float SleepMilliseconds = 3000f;
    // private static string _savePath;
    // private static string _packageName;
    // private static int _processMaxNum;
    // private static Action _completedCallback;
    //
    // private static ESteps _steps = ESteps.None;
    // private static Stopwatch _elapsedTime;
    //private static List<string> _allMaterials;
    // private static List<GameObject> _allSpheres = new List<GameObject>(1000);

    // 新增任务队列
    private static Queue<CollectTask> _taskQueue = new Queue<CollectTask>();
    private static CollectTask _currentTask;
    public static string OriginalScenePath { get; set; }

    static ShaderVariantCollector()
    {
        // 保存当前场景路径
        OriginalScenePath = EditorSceneManager.GetActiveScene().path;
    }


    // 新增初始化方法
    private static void InitializeTask(CollectTask task)
    {
        // 原有Run方法中的初始化逻辑迁移到这里
        if (Path.HasExtension(task.SavePath) == false)
            task.SavePath = $"{task.SavePath}.shadervariants";

        AssetDatabase.DeleteAsset(task.SavePath);
        EditorTools.CreateFileDirectory(task.SavePath);

        task.Steps = ESteps.Prepare;
    }

    /// <summary>
    /// 开始收集
    /// </summary>
    public static void Run(string savePath, string packageName, int processMaxNum, Action completedCallback)
    {
        if (Path.HasExtension(savePath) == false)
            savePath = $"{savePath}.shadervariants";
        if (Path.GetExtension(savePath) != ".shadervariants")
            throw new System.Exception("Shader variant file extension is invalid.");
        if (string.IsNullOrEmpty(packageName))
            throw new System.Exception("Package name is null or empty !");


        // 将新任务加入队列
        _taskQueue.Enqueue(new CollectTask
        {
            SavePath = savePath,
            PackageName = packageName,
            ProcessMaxNum = processMaxNum,
            CompletedCallback = completedCallback
        });
        // 如果当前没有正在执行的任务，开始执行队列
        if (_currentTask == null)
        {
            EditorApplication.update -= EditorUpdate;
            EditorApplication.update += EditorUpdate;

            // 聚焦到游戏窗口
            //EditorTools.FocusUnityGameWindow();

            // 创建临时测试场景
            CreateTempScene();
        }
    }

    private static void EditorUpdate()
    {
        // 当前没有活动任务时从队列获取新任务
        if (_currentTask == null && _taskQueue.Count > 0)
        {
            _currentTask = _taskQueue.Dequeue();
            InitializeTask(_currentTask);
        }

        if (_currentTask == null)
            return;


        if (_currentTask.Steps == ESteps.None)
            return;

        if (_currentTask.Steps == ESteps.Prepare)
        {
            ShaderVariantCollectionHelper.ClearCurrentShaderVariantCollection();
            _currentTask.Steps = ESteps.CollectAllMaterial;
            return; //等待一帧
        }

        if (_currentTask.Steps == ESteps.CollectAllMaterial)
        {
            _currentTask.AllMaterials = GetAllMaterials();
            _currentTask.Steps = ESteps.CollectVariants;
            return; //等待一帧
        }

        if (_currentTask.Steps == ESteps.CollectVariants)
        {
            int count = Mathf.Min(_currentTask.ProcessMaxNum, _currentTask.AllMaterials.Count);
            List<string> range = _currentTask.AllMaterials.GetRange(0, count);
            _currentTask.AllMaterials.RemoveRange(0, count);
            CollectVariants(range);

            if (_currentTask.AllMaterials.Count > 0)
            {
                _currentTask.ElapsedTime = Stopwatch.StartNew();
                _currentTask.Steps = ESteps.CollectSleeping;
            }
            else
            {
                _currentTask.ElapsedTime = Stopwatch.StartNew();
                _currentTask.Steps = ESteps.WaitingDone;
            }
        }

        if (_currentTask.Steps == ESteps.CollectSleeping)
        {
            if (_currentTask.ElapsedTime.ElapsedMilliseconds > SleepMilliseconds)
            {
                DestroyAllSpheres();
                _currentTask.ElapsedTime.Stop();
                _currentTask.Steps = ESteps.CollectVariants;
            }
        }

        if (_currentTask.Steps == ESteps.WaitingDone)
        {
            // 注意：一定要延迟保存才会起效
            if (_currentTask.ElapsedTime.ElapsedMilliseconds > WaitMilliseconds)
            {
                _currentTask.ElapsedTime.Stop();
                _currentTask.Steps = ESteps.None;

                // 保存结果并创建清单
                ShaderVariantCollectionHelper.SaveCurrentShaderVariantCollection(_currentTask.SavePath);
                CreateManifest();

                UnityEngine.Debug.Log($"搜集SVC完毕！");

                _currentTask.CompletedCallback?.Invoke();
                _currentTask = null;

                // 如果队列还有任务，立即触发下一次更新
                if (_taskQueue.Count > 0)
                {
                    EditorApplication.QueuePlayerLoopUpdate();
                }
                else
                {
                    // 恢复原场景
                    if (!string.IsNullOrEmpty(OriginalScenePath))
                    {
                        EditorSceneManager.OpenScene(OriginalScenePath);
                    }
                }
            }
        }
    }

    private static void CreateTempScene()
    {
        EditorSceneManager.NewScene(NewSceneSetup.DefaultGameObjects);
    }

    private static List<string> GetAllMaterials()
    {
        // 获取所有打包的资源
        CollectResult collectResult =
            AssetBundleCollectorSettingData.Setting.BeginCollect(_currentTask.PackageName, false, false);

        // 搜集所有材质球
        int progressValue = 0;
        HashSet<string> result = new HashSet<string>();
        foreach (var collectAssetInfo in collectResult.CollectAssets)
        {
            if (collectAssetInfo.AssetInfo.AssetType == typeof(UnityEngine.Material))
            {
                string assetPath = collectAssetInfo.AssetInfo.AssetPath;
                if (result.Contains(assetPath) == false)
                    result.Add(assetPath);
            }

            foreach (var dependAssetInfo in collectAssetInfo.DependAssets)
            {
                if (dependAssetInfo.AssetType == typeof(UnityEngine.Material))
                {
                    string assetPath = dependAssetInfo.AssetPath;
                    if (result.Contains(assetPath) == false)
                        result.Add(assetPath);
                }
            }

            EditorTools.DisplayProgressBar("搜集所有材质球", ++progressValue, collectResult.CollectAssets.Count);
        }

        EditorTools.ClearProgressBar();

        // 返回结果
        return result.ToList();
    }

    private static void CollectVariants(List<string> materials)
    {
        Camera camera = Camera.main;
        if (camera == null)
            throw new System.Exception("Not found main camera.");

        // 设置主相机
        float aspect = camera.aspect;
        int totalMaterials = materials.Count;
        float height = Mathf.Sqrt(totalMaterials / aspect) + 1;
        float width = Mathf.Sqrt(totalMaterials / aspect) * aspect + 1;
        float halfHeight = Mathf.CeilToInt(height / 2f);
        float halfWidth = Mathf.CeilToInt(width / 2f);
        camera.orthographic = true;
        camera.orthographicSize = halfHeight;
        camera.transform.position = new Vector3(0f, 0f, -10f);

        // 创建测试球体
        int xMax = (int)(width - 1);
        int x = 0, y = 0;
        int progressValue = 0;
        for (int i = 0; i < materials.Count; i++)
        {
            var material = materials[i];
            var position = new Vector3(x - halfWidth + 1f, y - halfHeight + 1f, 0f);
            var go = CreateSphere(material, position, i);
            if (go != null)
                _currentTask.AllSpheres.Add(go);
            if (x == xMax)
            {
                x = 0;
                y++;
            }
            else
            {
                x++;
            }

            EditorTools.DisplayProgressBar("照射所有材质球", ++progressValue, materials.Count);
        }

        EditorTools.ClearProgressBar();
    }

    private static GameObject CreateSphere(string assetPath, Vector3 position, int index)
    {
        var material = AssetDatabase.LoadAssetAtPath<Material>(assetPath);
        var shader = material.shader;
        if (shader == null)
            return null;

        var go = GameObject.CreatePrimitive(PrimitiveType.Sphere);
        go.GetComponent<Renderer>().sharedMaterial = material;
        go.transform.position = position;
        go.name = $"Sphere_{index} | {material.name}";
        return go;
    }

    private static void DestroyAllSpheres()
    {
        foreach (var go in _currentTask.AllSpheres)
        {
            GameObject.DestroyImmediate(go);
        }

        _currentTask.AllSpheres.Clear();

        // 尝试释放编辑器加载的资源
        EditorUtility.UnloadUnusedAssetsImmediate(true);
    }

    private static void CreateManifest()
    {
        AssetDatabase.Refresh(ImportAssetOptions.ForceUpdate);

        ShaderVariantCollection svc = AssetDatabase.LoadAssetAtPath<ShaderVariantCollection>(_currentTask.SavePath);
        if (svc != null)
        {
            var wrapper = ShaderVariantCollectionManifest.Extract(svc);
            string jsonData = JsonUtility.ToJson(wrapper, true);
            string savePath = _currentTask.SavePath.Replace(".shadervariants", ".json");
            File.WriteAllText(savePath, jsonData);
        }

        AssetDatabase.Refresh(ImportAssetOptions.ForceUpdate);
    }


    // 新增任务类
    private class CollectTask
    {
        public string SavePath;
        public string PackageName;
        public int ProcessMaxNum;
        public Action CompletedCallback;
        public ESteps Steps = ESteps.None;
        public Stopwatch ElapsedTime;
        public List<string> AllMaterials;
        public List<GameObject> AllSpheres = new List<GameObject>();
    }
}