生成代理类

MarvelTiter · Aug 26, 2024 · 437c20a · 437c20a
1 parent 3c8800a
commit 437c20a
Show file tree

Hide file tree

Showing 9 changed files with 288 additions and 76 deletions.
diff --git a/src/AutoAopProxy.Roslyn/AutoAopProxyClassGenerator.cs b/src/AutoAopProxy.Roslyn/AutoAopProxyClassGenerator.cs
@@ -61,7 +61,6 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
             }
             #endregion
 
-
             #region 获取所有的AspectHandler
             var allHandlers = allInterfaces.SelectMany(x =>
             {
@@ -84,38 +83,139 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
                 return handlers.Where(i => i != null).Cast<INamedTypeSymbol>();
             }).Distinct(EqualityComparer<INamedTypeSymbol>.Default).ToArray();
             #endregion
+
             var file = CreateGeneratedProxyClassFile(targetSymbol, allInterfaces, allHandlers);
-            context.AddSource(file);
+            if (file != null)
+            {
+                //var ss = file.ToString();
+                context.AddSource(file);
+            }
         });
     }
 
     private static CodeFile? CreateGeneratedProxyClassFile(INamedTypeSymbol classSymbol, INamedTypeSymbol[] interfaces, INamedTypeSymbol[] allHandlers)
     {
-        List<MethodBuilder> methods = [];
+        // 代理字段和aspect handler 字段
+        List<Node> members = [
+            FieldBuilder.Default.MemberType(classSymbol.ToDisplayString()).FieldName("proxy")
+            , .. allHandlers.Select(n =>FieldBuilder.Default.MemberType(n.ToDisplayString()).FieldName(n.MetadataName))
+            ];
         // 构造函数
-
+        List<Statement> ctorbody = [
+            "this.proxy = proxy"
+            , ..allHandlers.Select(n => $"this.{n.MetadataName} = {n.MetadataName}")
+            ];
+        var ctor = ConstructorBuilder.Default.MethodName($"{classSymbol.FormatClassName()}GeneratedProxy")
+            .AddParameter([$"{classSymbol.ToDisplayString()} proxy", .. allHandlers.Select(n => $"{n.ToDisplayString()} {n.MetadataName}")]).AddBody([.. ctorbody]);
+        members.Add(ctor);
         // 接口方法
         foreach (var iface in interfaces)
         {
-            methods.AddRange(CreateProxyMethod(iface));
+            members.AddRange(CreateProxyMethod(iface, classSymbol));
         }
 
-        var ctor = ConstructorBuilder.Default.MethodName($"{classSymbol.FormatClassName()}GeneratedProxy");
-
-        var proxyClass = ClassBuilder.Default.ClassName($"{classSymbol.FormatClassName()}GeneratedProxy");
+        var proxyClass = ClassBuilder.Default
+            .ClassName($"{classSymbol.FormatClassName()}GeneratedProxy")
+            .AddMembers([.. members])
+            .Interface([.. interfaces.Select(i => i.ToDisplayString())])
+            .AddGeneratedCodeAttribute(typeof(AutoAopProxyClassGenerator));
 
         return CodeFile.New($"{classSymbol.FormatFileName()}GeneratedProxyClass.g.cs")
+            .AddUsings("using AutoAopProxyGenerator;")
             .AddMembers(NamespaceBuilder.Default.Namespace(classSymbol.ContainingNamespace.ToDisplayString()).AddMembers(proxyClass));
     }
 
-    private static IEnumerable<MethodBuilder> CreateProxyMethod(INamedTypeSymbol iface)
+    private static IEnumerable<MethodBuilder> CreateProxyMethod(INamedTypeSymbol iface, INamedTypeSymbol classSymbol)
     {
-        var handlers = iface.GetAttributes(AspectHandler);
+        var handlers = iface.GetAttributes(AspectHandler).Select(a => a.GetNamedValue("AspectType")).OfType<INamedTypeSymbol>().ToArray();
 
         foreach (var m in iface.GetMethods())
         {
+            var method = m.IsGenericMethod ? m.ConstructedFrom : m;
+            var returnType = method.ReturnType.GetGenericTypes().FirstOrDefault() ?? method.ReturnType;
+            var isAsync = method.IsAsync || method.ReturnType.ToDisplayString().StartsWith("System.Threading.Tasks.Task");
+            var builder = MethodBuilder.Default
+                .MethodName(method.Name)
+                .Async(isAsync)
+                .Generic([.. method.GetTypeParameters()])
+                .AddParameter([.. method.Parameters.Select(p => $"{p.Type.ToDisplayString()} {p.Name}")])
+                .ReturnType(method.ReturnType.ToDisplayString())
+                .AddGeneratedCodeAttribute(typeof(AutoAopProxyClassGenerator));
+            List<Statement> statements = [];
+            var hasReturn = !method.ReturnsVoid && returnType.ToDisplayString() != "System.Threading.Tasks.Task";
+            if (hasReturn)
+            {
+                statements.Add($"{returnType.ToDisplayString()} returnValue = default;");
+            }
+            var done = LocalFunction.Default
+                .MethodName("Done")
+                .AddParameters("ProxyContext ctx")
+                .Async(isAsync)
+                .Return("System.Threading.Tasks.Task")
+                .AddBody([.. CreateLocalFunctionBody(method, builder.ConstructedMethodName, isAsync, hasReturn)]);
+
+            statements.Add(done);
+            statements.Add("var builder = AsyncPipelineBuilder<ProxyContext>.Create(Done)");
+            foreach (var handler in handlers)
+            {
+                statements.Add($"builder.Use({handler.MetadataName}.Invoke)");
+            }
+            statements.Add("var job = builder.Build()");
+            var ptypes = method.Parameters.Length > 0 ? $"[{string.Join(", ", method.Parameters.Select(p => $"typeof({p.Type.ToDisplayString()})"))}]" : "Type.EmptyTypes";
+            statements.Add($"var context = ContextHelper.GetOrCreate(typeof({iface.ToDisplayString()}), typeof({classSymbol.ToDisplayString()}), nameof({method.Name}), {ptypes})");
+            if (hasReturn)
+            {
+                statements.Add("context.HasReturnValue = true");
+            }
+            statements.Add($"context.Parameters = new object?[] {{{string.Join(", ", method.Parameters.Select(p => p.Name))}}};");
+            if (isAsync)
+            {
+                statements.Add("await job.Invoke(context)");
+            }
+            else
+            {
+                statements.Add("job.Invoke(context).GetAwaiter().GetResult()");
+            }
+            if (hasReturn)
+                statements.Add("return returnValue");
 
+            builder.AddBody([.. statements]);
+
+            yield return builder;
+        }
+    }
+
+    private static IEnumerable<Statement> CreateLocalFunctionBody(IMethodSymbol method, string proxyName, bool isAsync, bool hasReturn)
+    {
+        if (isAsync)
+        {
+            if (hasReturn)
+            {
+                yield return $"returnValue = await proxy.{proxyName}({string.Join(", ", method.Parameters.Select(p => p.Name))})";
+                yield return "ctx.Executed = true";
+                yield return "ctx.ReturnValue = returnValue";
+            }
+            else
+            {
+                yield return $"await proxy.{proxyName}({string.Join(", ", method.Parameters.Select(p => p.Name))})";
+                yield return "ctx.Executed = true";
+            }
+        }
+        else
+        {
+            if (hasReturn)
+            {
+                yield return $"returnValue = proxy.{proxyName}({string.Join(", ", method.Parameters.Select(p => p.Name))})";
+                yield return "ctx.Executed = true";
+                yield return "ctx.ReturnValue = returnValue";
+                yield return "return global::System.Threading.Tasks.Task.CompletedTask";
+            }
+            else
+            {
+                yield return $"proxy.{proxyName}({string.Join(", ", method.Parameters.Select(p => p.Name))})";
+                yield return "ctx.Executed = true";
+                yield return "return global::System.Threading.Tasks.Task.CompletedTask";
+            }
         }
-        yield break;
     }
 }
diff --git a/src/AutoAopProxyGenerator/IAspectHandler.cs b/src/AutoAopProxyGenerator/IAspectHandler.cs
@@ -1,4 +1,6 @@
 using System;
+using System.Collections.Concurrent;
+using System.Linq;
 using System.Threading.Tasks;
 
 namespace AutoAopProxyGenerator;
@@ -8,3 +10,25 @@ public interface IAspectHandler
     Task Invoke(ProxyContext context, Func<Task> process);
 }
 
+public class ContextHelper
+{
+    private static readonly ConcurrentDictionary<string, ProxyContext> caches = [];
+    public static ProxyContext GetOrCreate(Type @interface, Type impl, string methodName, Type[] types)
+    {
+        var key = $"{@interface.FullName}_{impl.FullName}_{methodName}_{string.Join("_", types.Select(t => t.Name))}";
+        return caches.GetOrAdd(key, (k) =>
+          {
+              var interfacesMethods = @interface.GetMethod(methodName, types);
+              var implMethod = impl.GetMethod(methodName, types);
+              return new ProxyContext()
+              {
+                  ServiceType = @interface,
+                  ImplementType = impl,
+                  ServiceMethod = interfacesMethods,
+                  ImplementMethod = implMethod
+              };
+          });
+
+    }
+}
+
diff --git a/src/AutoAopProxyGenerator/Models/ProxyContext.cs b/src/AutoAopProxyGenerator/Models/ProxyContext.cs
@@ -5,7 +5,7 @@
 
 namespace AutoAopProxyGenerator
 {
-    public class ProxyContext
+    public record ProxyContext
     {
         public ProxyContext(IServiceProvider services)
         {

diff --git a/src/Generators.Shared/Builder/MethodBuilder.cs b/src/Generators.Shared/Builder/MethodBuilder.cs
@@ -16,7 +16,7 @@ public MethodBuilder()
     public bool IsAsync { get; set; }
     string Async => IsAsync ? " async " : " ";
     public string? ReturnType { get; set; } = "void";
-
+    public string ConstructedMethodName => $"{Name}{Types}";
     public override string ToString()
     {
         if (IsLambdaBody)
@@ -209,10 +209,10 @@ public override string ToString()
     {
         return
 $$"""
-{{Indent}}{{Async}}{{Name}}({{string.Join(", ", Parameters)}})
-{
+{{Indent}}{{Async}}{{ReturnType}} {{Name}}({{string.Join(", ", Parameters)}})
+{{Indent}}{
 {{string.Join("\n", Body.Select(s => $"    {s}"))}}
-}
+{{Indent}}}
 """;
     }
 }
diff --git a/src/Generators.Shared/Builder/TypeBuilder.cs b/src/Generators.Shared/Builder/TypeBuilder.cs
@@ -36,7 +36,7 @@ protected string TypeConstraints
         get
         {
             if (TypeArguments.Count == 0) return "";
-            return $"\n{string.Join("\n", TypeArguments.Select(ta => $"{Indent}    where {ta.Name} : {string.Join(", ", ta.Constraints)}"))}";
+            return $"\n{string.Join("\n", TypeArguments.Where(ta => ta.Constraints.Length > 0).Select(ta => $"{Indent}    where {ta.Name} : {string.Join(", ", ta.Constraints)}"))}";
         }
     }
 }

diff --git a/src/Generators.Shared/Extensions/MethodBuilderExtensions.cs b/src/Generators.Shared/Extensions/MethodBuilderExtensions.cs
@@ -30,9 +30,9 @@ public static MethodBuilder ReturnType(this MethodBuilder builder, string return
         builder.ReturnType = returnType;
         return builder;
     }
-    public static MethodBuilder Async(this MethodBuilder builder)
+    public static MethodBuilder Async(this MethodBuilder builder, bool isAsync = true)
     {
-        builder.IsAsync = true;
+        builder.IsAsync = isAsync;
         return builder;
     }
 
@@ -143,6 +143,12 @@ public static LocalFunction MethodName(this LocalFunction localFunction, string
         return localFunction;
     }
 
+    public static LocalFunction Async(this LocalFunction localFunction, bool isAsync = true)
+    {
+        localFunction.IsAsync = isAsync;
+        return localFunction;
+    }
+
     public static LocalFunction Return(this LocalFunction localFunction, string returnType)
     {
         localFunction.ReturnType = returnType;

diff --git a/src/TestProject1/AopGeneratorTest/LogAop.cs b/src/TestProject1/AopGeneratorTest/LogAop.cs
@@ -7,11 +7,11 @@
 
 namespace TestProject1.AopGeneratorTest
 {
-    internal class LogAop : IAspectHandler
+    public class LogAop : IAspectHandler
     {
         public Task Invoke(ProxyContext context, Func<Task> process)
         {
-            throw new NotImplementedException();
+            return process.Invoke();
         }
     }
 }
diff --git a/src/TestProject1/AopGeneratorTest/ProxyTest.cs b/src/TestProject1/AopGeneratorTest/ProxyTest.cs
@@ -0,0 +1,24 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace TestProject1.AopGeneratorTest
+{
+    [TestClass]
+    public class ProxyTest
+    {
+        [TestMethod]
+        public void TestReturn()
+        {
+
+        }
+
+        [TestMethod]
+        public Task TestReturnAsync()
+        {
+            return Task.FromResult(0);
+        }
+    }
+}
-Original file line number
+Diff line change
@@ Expand Up / @@ -36,7 +36,7 @@ protected string TypeConstraints @@
             get
             {
                 if (TypeArguments.Count == 0) return "";
-                return $"\n{string.Join("\n", TypeArguments.Select(ta => $"{Indent}    where {ta.Name} : {string.Join(", ", ta.Constraints)}"))}";
+                return $"\n{string.Join("\n", TypeArguments.Where(ta => ta.Constraints.Length > 0).Select(ta => $"{Indent}    where {ta.Name} : {string.Join(", ", ta.Constraints)}"))}";
             }
         }
     }
@@ Expand Down @@