ShrinkModFramework/Runtime/Context/ShrinkModContextHost.cs

using System;
using System.Collections.Generic;
using System.Linq;
using Cysharp.Threading.Tasks;
using ShrinkContext;

namespace ShrinkModFramework
{
    public sealed class ShrinkModTransactionException : Exception
    {
        public ShrinkModTransactionException(string message, Exception applyError, Exception restoreError = null)
            : base(message, restoreError == null ? applyError : new AggregateException(applyError, restoreError))
        {
            ApplyError = applyError;
            RestoreError = restoreError;
        }

        public Exception ApplyError { get; }
        public Exception RestoreError { get; }
        public bool PreviousCompositionRestored => RestoreError == null;
    }

    /// <summary>
    /// 模组的 Cordis 组合宿主。配置应用是事务：新组合任一 fiber 失败时，重新协调到旧组件源；
    /// Mono/IL2CPP 下旧程序集仍驻留，但旧/新模组实例及其已追踪效应会被正确卸载或恢复。
    /// </summary>
    public sealed class ShrinkModContextHost
    {
        private readonly Dictionary<string, ShrinkModHandle> _mods = new(StringComparer.Ordinal);
        private readonly Dictionary<ShrinkModComponentSource, string> _catalogNames = new();
        private readonly List<ShrinkModComponentSource> _currentSources = new();
        private readonly ShrinkComponentCatalog _catalog = new();
        private long _generation;
        private int _catalogGeneration;
        private bool _applying;
        private IReadOnlyList<ShrinkModComponentSource> _buildingSources = Array.Empty<ShrinkModComponentSource>();

        internal IEnumerable<string> GetPredecessors(string modId) =>
            ShrinkModDependencyGraph.Predecessors(modId, _buildingSources.Select(source => source.Info));

        public ShrinkModContextHost(bool enableHarmonyPatching = true, bool verboseLogging = false)
        {
            EnableHarmonyPatching = enableHarmonyPatching;
            VerboseLogging = verboseLogging;
            Runtime = new ShrinkContextRuntime();
            Loader = new ShrinkContextLoader(Runtime, _catalog);
            RegistryManager = new ShrinkModRegistryManager();
        }

        public event Action<ShrinkModHandle> OnModReady;
        public event Action<IReadOnlyDictionary<string, ShrinkModHandle>> OnAllModsReady;

        public bool EnableHarmonyPatching { get; }
        public bool VerboseLogging { get; }
        public ShrinkContextRuntime Runtime { get; }
        public ShrinkContextLoader Loader { get; }
        public IReadOnlyDictionary<string, ShrinkModHandle> Mods => _mods;
        public IReadOnlyList<ShrinkModComponentSource> CurrentSources => _currentSources.ToArray();
        internal ShrinkModRegistryManager RegistryManager { get; }

        public IReadOnlyShrinkModRegistry<T> GetOrCreateRegistry<T>(string name) =>
            RegistryManager.GetOrCreateRegistry<T>(name);

        public static string GetModKey(string modId)
        {
            if (string.IsNullOrWhiteSpace(modId))
                throw new ArgumentException("ModId must not be empty.", nameof(modId));
            return "shrink.mod/" + modId.Trim();
        }

        public UniTask ApplyAsync(IReadOnlyList<ShrinkModComponentSource> desiredSources) => ApplyAsync(desiredSources, null);

        internal async UniTask ApplyAsync(IReadOnlyList<ShrinkModComponentSource> desiredSources, Action beforeRestore)
        {
            if (desiredSources == null)
                throw new ArgumentNullException(nameof(desiredSources));
            if (_applying)
                throw new InvalidOperationException("A mod composition transaction is already running.");

            var desired = ReuseUnchangedSources(ValidateAndOrder(desiredSources));
            var previous = _currentSources.ToArray();
            var previousById = previous.ToDictionary(source => source.Info.ModId, StringComparer.Ordinal);
            _applying = true;
            try
            {
                await Loader.ApplyAsync(BuildEntries(desired));
                EnsureActive(desired);

                _currentSources.Clear();
                _currentSources.AddRange(desired);
                foreach (var source in desired)
                {
                    if (!previousById.TryGetValue(source.Info.ModId, out var oldSource) ||
                        !ReferenceEquals(source, oldSource))
                    {
                        OnModReady?.Invoke(_mods[source.Info.ModId]);
                    }
                }
                OnAllModsReady?.Invoke(Mods);
            }
            catch (Exception applyError)
            {
                Exception restoreError = null;
                _currentSources.Clear();
                _currentSources.AddRange(previous);
                try
                {
                    beforeRestore?.Invoke();
                    await Loader.ApplyAsync(BuildEntries(previous));
                    EnsureActive(previous);
                }
                catch (Exception ex)
                {
                    restoreError = ex;
                }

                throw new ShrinkModTransactionException(
                    restoreError == null
                        ? "Mod composition failed; the previous composition was restored."
                        : "Mod composition failed and restoring the previous composition also failed.",
                    applyError,
                    restoreError);
            }
            finally
            {
                _applying = false;
            }
        }

        public async UniTask ShutdownAsync()
        {
            if (_applying)
                throw new InvalidOperationException("Cannot shut down while a mod transaction is running.");

            await Loader.ApplyAsync(Array.Empty<ShrinkLoaderEntry>());
            _currentSources.Clear();
        }

        internal long NextGeneration() => ++_generation;

        internal void RegisterHandle(ShrinkModHandle handle)
        {
            if (!_mods.TryAdd(handle.Info.ModId, handle))
                throw new InvalidOperationException($"Mod handle already active: {handle.Info.ModId}");
        }

        internal void UnregisterHandle(ShrinkModHandle handle)
        {
            if (_mods.TryGetValue(handle.Info.ModId, out var current) && ReferenceEquals(current, handle))
                _mods.Remove(handle.Info.ModId);
        }

        private IReadOnlyList<ShrinkLoaderEntry> BuildEntries(IEnumerable<ShrinkModComponentSource> sources)
        {
            _buildingSources = sources.ToArray();
            var entries = new List<ShrinkLoaderEntry>();
            foreach (var source in sources)
            {
                if (!_catalogNames.TryGetValue(source, out var catalogName))
                {
                    catalogName = $"shrink.mod.source/{source.Info.ModId}/{++_catalogGeneration}";
                    _catalogNames.Add(source, catalogName);
                    _catalog.Register(catalogName, () => new ShrinkModComponent(this, source));
                }

                entries.Add(new ShrinkLoaderEntry(source.Info.ModId, catalogName));
            }
            return entries;
        }

        private void EnsureActive(IEnumerable<ShrinkModComponentSource> sources)
        {
            foreach (var source in sources)
            {
                if (!Loader.TryGetFiber(source.Info.ModId, out var fiber))
                    throw new InvalidOperationException($"Mod fiber was not created: {source.Info.ModId}");
                if (fiber.LastError != null)
                    throw new InvalidOperationException($"Mod {source.Info.ModId} failed during apply.", fiber.LastError);
                if (fiber.State != ShrinkFiberState.Active)
                    throw new InvalidOperationException(
                        $"Mod {source.Info.ModId} did not become active (state={fiber.State}).");
            }
        }

        private static List<ShrinkModComponentSource> ValidateAndOrder(
            IReadOnlyList<ShrinkModComponentSource> sources)
        {
            var map = new Dictionary<string, ShrinkModComponentSource>(StringComparer.Ordinal);
            foreach (var source in sources)
            {
                if (source == null)
                    throw new ArgumentException("Mod sources must not contain null.", nameof(sources));
                if (!map.TryAdd(source.Info.ModId, source))
                    throw new InvalidOperationException($"Duplicate mod source id: {source.Info.ModId}");
            }

            return ShrinkModDependencyGraph.Order(sources.Select(source => source.Info))
                .Select(info => map[info.ModId]).ToList();
        }

        private List<ShrinkModComponentSource> ReuseUnchangedSources(
            IEnumerable<ShrinkModComponentSource> desired)
        {
            var desiredArray = desired.ToArray();
            var currentById = _currentSources.ToDictionary(source => source.Info.ModId, StringComparer.Ordinal);
            var normalized = new List<ShrinkModComponentSource>();
            foreach (var source in desiredArray)
            {
                if (currentById.TryGetValue(source.Info.ModId, out var current) &&
                    string.Equals(current.Revision, source.Revision, StringComparison.Ordinal) &&
                    ShrinkModDependencyGraph.SameMetadata(current.Info, source.Info) &&
                    ShrinkModDependencyGraph.Predecessors(source.Info.ModId, _currentSources.Select(item => item.Info)).OrderBy(id => id)
                        .SequenceEqual(ShrinkModDependencyGraph.Predecessors(source.Info.ModId, desiredArray.Select(item => item.Info)).OrderBy(id => id)))
                {
                    normalized.Add(current);
                }
                else
                {
                    normalized.Add(ReferenceEquals(source, current)
                        ? source.WithMetadata(source.Info, source.Revision) : source);
                }
            }
            return normalized;
        }

    }
}