com.alicizax.unity.framework/Runtime/MemoryPool/MemoryPoolRegistry.cs

using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;

namespace AlicizaX
{

    public static class MemoryPoolRegistry
    {
        internal sealed class MemoryPoolHandle
        {
            public delegate IMemory AcquireHandler();
            public delegate void ReleaseHandler(IMemory memory);
            public delegate void ClearHandler();
            public delegate void IntHandler(int value);
            public delegate void GetInfoHandler(ref MemoryPoolInfo info);

            public readonly AcquireHandler Acquire;
            public readonly ReleaseHandler Release;
            public readonly ClearHandler Clear;
            public readonly IntHandler Prewarm;
            public readonly GetInfoHandler GetInfo;
            public readonly IntHandler Tick;
            public readonly IntHandler Shrink;

            public MemoryPoolHandle(
                AcquireHandler acquire,
                ReleaseHandler release,
                ClearHandler clear,
                IntHandler prewarm,
                GetInfoHandler getInfo,
                IntHandler tick,
                IntHandler shrink)
            {
                Acquire = acquire;
                Release = release;
                Clear = clear;
                Prewarm = prewarm;
                GetInfo = getInfo;
                Tick = tick;
                Shrink = shrink;
            }
        }

        private static readonly Dictionary<Type, MemoryPoolHandle> s_Handles
            = new Dictionary<Type, MemoryPoolHandle>(64);

        private static MemoryPoolHandle.IntHandler[] s_TickArray = Array.Empty<MemoryPoolHandle.IntHandler>();
        private static int s_TickCount;
        private static bool s_TickArrayDirty;

        public static int Count => s_Handles.Count;

        internal static void Register(Type type, MemoryPoolHandle handle)
        {
            s_Handles[type] = handle;
            s_TickArrayDirty = true;
        }

        /// <summary>
        /// 非泛型 Acquire，用于只有 Type 没有泛型参数的场景。
        /// </summary>
        public static IMemory Acquire(Type type)
        {
            if (type == null)
                throw new ArgumentNullException(nameof(type));

            if (s_Handles.TryGetValue(type, out var handle))
                return handle.Acquire();

            EnsureRegistered(type);

            if (s_Handles.TryGetValue(type, out handle))
                return handle.Acquire();

            throw new Exception($"MemoryPool: Type '{type.FullName}' is not a valid IMemory type.");
        }

        public static void Release(IMemory memory)
        {
            if (memory == null)
                throw new ArgumentNullException(nameof(memory));

            Type type = memory.GetType();
            if (s_Handles.TryGetValue(type, out var handle))
            {
                handle.Release(memory);
                return;
            }

            EnsureRegistered(type);

            if (s_Handles.TryGetValue(type, out handle))
            {
                handle.Release(memory);
                return;
            }

            throw new Exception($"MemoryPool: Type '{type.FullName}' is not a valid IMemory type.");
        }

        public static int GetAllInfos(MemoryPoolInfo[] infos)
        {
            if (infos == null)
                throw new ArgumentNullException(nameof(infos));

            int count = s_Handles.Count;
            if (infos.Length < count)
                throw new ArgumentException("Target buffer is too small.", nameof(infos));

            int i = 0;
            foreach (var kv in s_Handles)
            {
                kv.Value.GetInfo(ref infos[i]);
                i++;
            }

            return count;
        }

        public static MemoryPoolInfo[] GetAllInfos()
        {
            var infos = new MemoryPoolInfo[s_Handles.Count];
            GetAllInfos(infos);
            return infos;
        }

        public static void ClearAll()
        {
            foreach (var kv in s_Handles)
                kv.Value.Clear();
        }

        public static void Prewarm(Type type, int count)
        {
            if (type == null)
                throw new ArgumentNullException(nameof(type));

            if (s_Handles.TryGetValue(type, out var handle))
            {
                handle.Prewarm(count);
                return;
            }

            EnsureRegistered(type);

            if (s_Handles.TryGetValue(type, out handle))
            {
                handle.Prewarm(count);
                return;
            }

            throw new Exception($"MemoryPool: Type '{type.FullName}' is not a valid IMemory type.");
        }

        public static void ClearType(Type type)
        {
            if (type == null)
                throw new ArgumentNullException(nameof(type));

            if (!s_Handles.TryGetValue(type, out var handle))
            {
                EnsureRegistered(type);
                if (!s_Handles.TryGetValue(type, out handle))
                    throw new Exception($"MemoryPool: Type '{type.FullName}' is not a valid IMemory type.");
            }

            handle.Clear();
        }

        public static void RemoveFromType(Type type, int count)
        {
            if (type == null)
                throw new ArgumentNullException(nameof(type));

            if (s_Handles.TryGetValue(type, out var handle))
            {
                MemoryPoolInfo info = default;
                handle.GetInfo(ref info);
                int unused = info.UnusedCount;
                handle.Shrink(unused - count);
                return;
            }

            EnsureRegistered(type);
            if (s_Handles.TryGetValue(type, out handle))
            {
                MemoryPoolInfo info = default;
                handle.GetInfo(ref info);
                int unused = info.UnusedCount;
                handle.Shrink(unused - count);
                return;
            }

            throw new Exception($"MemoryPool: Type '{type.FullName}' is not a valid IMemory type.");
        }


        public static void TickAll(int frameCount)
        {
            if (s_TickArrayDirty)
                RebuildTickArray();

            for (int i = 0; i < s_TickCount; i++)
                s_TickArray[i](frameCount);
        }

        private static void RebuildTickArray()
        {
            s_TickCount = s_Handles.Count;
            if (s_TickArray.Length < s_TickCount)
                s_TickArray = new MemoryPoolHandle.IntHandler[s_TickCount];

            int i = 0;
            foreach (var kv in s_Handles)
                s_TickArray[i++] = kv.Value.Tick;

            s_TickArrayDirty = false;
        }

        private static void EnsureRegistered(Type type)
        {
            if (type == null)
                throw new ArgumentNullException(nameof(type));

            if (!typeof(IMemory).IsAssignableFrom(type))
                throw new Exception($"MemoryPool: Type '{type.FullName}' is not a valid IMemory type.");

            RuntimeHelpers.RunClassConstructor(
                typeof(MemoryPool<>).MakeGenericType(type).TypeHandle);
        }
    }
}