Browse Source

Add LightAsyncValueTaskMethodBuilder and related types for zero allocation async support

Akeit0 7 months ago
parent
commit
7452141a41

+ 10 - 0
src/Lua/Internal/CompilerServices/AsyncMethodBuilderAttribute.cs

@@ -0,0 +1,10 @@
+#pragma warning disable CS1591 // Missing XML comment for publicly visible type or member
+#pragma warning disable CS0436
+
+namespace System.Runtime.CompilerServices
+{
+    internal sealed class AsyncMethodBuilderAttribute(Type builderType) : Attribute
+    {
+        public Type BuilderType { get; } = builderType;
+    }
+}

+ 257 - 0
src/Lua/Internal/CompilerServices/LightAsyncTaskMethodBuilder.cs

@@ -0,0 +1,257 @@
+#pragma warning disable CS1591
+/*
+
+LightAsyncValueTaskMethodBuilder is  based on UniTask
+https://github.com/Cysharp/UniTask
+
+MIT License
+
+Copyright (c) 2019 Cysharp, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+*/
+using System.Diagnostics;
+using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+using System.Security;
+
+namespace Lua.Internal.CompilerServices;
+
+[StructLayout(LayoutKind.Auto)]
+internal struct LightAsyncValueTaskMethodBuilder
+{
+    IStateMachineRunnerPromise? runnerPromise;
+    Exception? ex;
+
+    // 1. Static Create method.
+    [DebuggerHidden]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static LightAsyncValueTaskMethodBuilder Create()
+    {
+        return default;
+    }
+
+    // 2. TaskLike Task property.
+    public ValueTask Task
+    {
+        [DebuggerHidden]
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        get
+        {
+            if (runnerPromise != null)
+            {
+                return runnerPromise.Task;
+            }
+
+            if (ex != null)
+            {
+                return new(System.Threading.Tasks.Task.FromException(ex));
+            }
+
+            return new(System.Threading.Tasks.Task.CompletedTask);
+        }
+    }
+
+    // 3. SetException
+    [DebuggerHidden]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void SetException(Exception exception)
+    {
+        if (runnerPromise == null)
+        {
+            ex = exception;
+        }
+        else
+        {
+            runnerPromise.SetException(exception);
+        }
+    }
+
+    // 4. SetResult
+    [DebuggerHidden]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void SetResult()
+    {
+        runnerPromise?.SetResult();
+    }
+
+    // 5. AwaitOnCompleted
+    [DebuggerHidden]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine)
+        where TAwaiter : INotifyCompletion
+        where TStateMachine : IAsyncStateMachine
+    {
+        if (runnerPromise == null)
+        {
+            LightAsyncValueTask<TStateMachine>.SetStateMachine(ref stateMachine, ref runnerPromise);
+        }
+
+        awaiter.OnCompleted(runnerPromise!.MoveNext);
+    }
+
+    // 6. AwaitUnsafeOnCompleted
+    [DebuggerHidden]
+    [SecuritySafeCritical]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine)
+        where TAwaiter : ICriticalNotifyCompletion
+        where TStateMachine : IAsyncStateMachine
+    {
+        if (runnerPromise == null)
+        {
+            LightAsyncValueTask<TStateMachine>.SetStateMachine(ref stateMachine, ref runnerPromise);
+        }
+
+        awaiter.UnsafeOnCompleted(runnerPromise!.MoveNext);
+    }
+
+    // 7. Start
+    [DebuggerHidden]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void Start<TStateMachine>(ref TStateMachine stateMachine)
+        where TStateMachine : IAsyncStateMachine
+    {
+        stateMachine.MoveNext();
+    }
+
+    // 8. SetStateMachine
+    [DebuggerHidden]
+    public void SetStateMachine(IAsyncStateMachine stateMachine)
+    {
+        // don't use boxed stateMachine.
+    }
+}
+
+[StructLayout(LayoutKind.Auto)]
+internal struct LightAsyncValueTaskMethodBuilder<T>
+{
+    IStateMachineRunnerPromise<T>? runnerPromise;
+    Exception? ex;
+    T result;
+
+    // 1. Static Create method.
+    [DebuggerHidden]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static LightAsyncValueTaskMethodBuilder<T> Create()
+    {
+        return default;
+    }
+
+    // 2. TaskLike Task property.
+    public ValueTask<T> Task
+    {
+        [DebuggerHidden]
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        get
+        {
+            if (runnerPromise != null)
+            {
+                return runnerPromise.Task;
+            }
+
+            if (ex != null)
+            {
+                return new(System.Threading.Tasks.Task.FromException<T>(ex));
+            }
+
+            {
+                return new(System.Threading.Tasks.Task.FromResult(result));
+            }
+        }
+    }
+
+    // 3. SetException
+    [DebuggerHidden]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void SetException(Exception exception)
+    {
+        if (runnerPromise == null)
+        {
+            ex = exception;
+        }
+        else
+        {
+            runnerPromise.SetException(exception);
+        }
+    }
+
+    // 4. SetResult
+    [DebuggerHidden]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void SetResult(T result)
+    {
+        if (runnerPromise == null)
+        {
+            this.result = result;
+        }
+        else
+        {
+            runnerPromise.SetResult(result);
+        }
+    }
+
+    // 5. AwaitOnCompleted
+    [DebuggerHidden]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine)
+        where TAwaiter : INotifyCompletion
+        where TStateMachine : IAsyncStateMachine
+    {
+        if (runnerPromise == null)
+        {
+            LightAsyncValueTask<TStateMachine, T>.SetStateMachine(ref stateMachine, ref runnerPromise);
+        }
+
+        awaiter.OnCompleted(runnerPromise!.MoveNext);
+    }
+
+    // 6. AwaitUnsafeOnCompleted
+    [DebuggerHidden]
+    [SecuritySafeCritical]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine)
+        where TAwaiter : ICriticalNotifyCompletion
+        where TStateMachine : IAsyncStateMachine
+    {
+        if (runnerPromise == null)
+        {
+            LightAsyncValueTask<TStateMachine, T>.SetStateMachine(ref stateMachine, ref runnerPromise);
+        }
+
+        awaiter.UnsafeOnCompleted(runnerPromise!.MoveNext);
+    }
+
+    // 7. Start
+    [DebuggerHidden]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void Start<TStateMachine>(ref TStateMachine stateMachine)
+        where TStateMachine : IAsyncStateMachine
+    {
+        stateMachine.MoveNext();
+    }
+
+    // 8. SetStateMachine
+    [DebuggerHidden]
+    public void SetStateMachine(IAsyncStateMachine stateMachine)
+    {
+        // don't use boxed stateMachine.
+    }
+}

+ 3 - 6
src/Lua/LuaCoroutine.cs

@@ -85,9 +85,8 @@ public sealed class LuaCoroutine : LuaThread, IValueTaskSource<LuaCoroutine.Yiel
         return ResumeAsyncCore(context.Thread.Stack, context.ArgumentCount, context.ReturnFrameBase, context.Thread, cancellationToken);
     }
 
-#if NET6_0_OR_GREATER
-    [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))]
-#endif
+
+    [AsyncMethodBuilder(typeof(LightAsyncValueTaskMethodBuilder<>))]
     async ValueTask<int> ResumeAsyncCore(LuaStack stack, int argCount, int returnBase, LuaThread? baseThread, CancellationToken cancellationToken = default)
     {
         if (baseThread != null)
@@ -229,9 +228,7 @@ public sealed class LuaCoroutine : LuaThread, IValueTaskSource<LuaCoroutine.Yiel
         return YieldAsyncCore(context.Thread.Stack, context.ArgumentCount, context.ReturnFrameBase, context.Thread, cancellationToken);
     }
 
-#if NET6_0_OR_GREATER
-    [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))]
-#endif
+    [AsyncMethodBuilder(typeof(LightAsyncValueTaskMethodBuilder<>))]
     async ValueTask<int> YieldAsyncCore(LuaStack stack, int argCount, int returnBase, LuaThread? baseThread, CancellationToken cancellationToken = default)
     {
         if (Volatile.Read(ref status) != (byte)LuaThreadStatus.Running)