From bd31a0e5f055c2192d99670b88fdb2ba94da1360 Mon Sep 17 00:00:00 2001
From: unknown <Meepo_Tian@tcossha.com>
Date: Mon, 3 Jul 2023 09:40:53 +0800
Subject: [PATCH] fix generatetype assembly ref issue

---
 .../CodeGeneration/CodeGenerationContext.cs    | 18 ++++++++++++++++++
 .../CodeGeneration/InterfaceProxyGenerator.cs  |  4 ++--
 .../VirtualMethodProxyGenerator.cs             |  4 ++--
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/Interception/Interception/CodeGeneration/CodeGenerationContext.cs b/src/Interception/Interception/CodeGeneration/CodeGenerationContext.cs
index 95dea7d..c328b7f 100644
--- a/src/Interception/Interception/CodeGeneration/CodeGenerationContext.cs
+++ b/src/Interception/Interception/CodeGeneration/CodeGenerationContext.cs
@@ -25,6 +25,24 @@ public sealed class CodeGenerationContext
             typeof(CodeGenerationContext).Assembly
         };
 
+        /// <summary>
+        /// Add reference, support generic type
+        /// </summary>
+        /// <param name="type"></param>
+        /// <returns></returns>
+        public ISet<Assembly> AddReference(Type type)
+        {
+            if (type.IsGenericType)
+            {
+                var types = type.GetGenericArguments();
+                foreach (var t in types)
+                   return AddReference(t);
+            }
+            else
+                this.References.Add(type.Assembly);
+            return References;
+        }
+
         /// <summary>Gets the indent level.</summary>
         /// <value>The indent level.</value>
         public int IndentLevel { get; private set; }
diff --git a/src/Interception/Interception/CodeGeneration/InterfaceProxyGenerator.cs b/src/Interception/Interception/CodeGeneration/InterfaceProxyGenerator.cs
index d47f53a..bf5cc51 100644
--- a/src/Interception/Interception/CodeGeneration/InterfaceProxyGenerator.cs
+++ b/src/Interception/Interception/CodeGeneration/InterfaceProxyGenerator.cs
@@ -50,10 +50,10 @@ public override bool TryGenerate(ServiceDescriptor serviceDescriptor, CodeGenera
 
             foreach (var method in methods)
             {
-                codeGenerationContext.References.Add(method.ReturnType.Assembly);
+                codeGenerationContext.AddReference(method.ReturnType);
                 foreach (var parameter in method.GetParameters())
                 {
-                    codeGenerationContext.References.Add(parameter.ParameterType.Assembly);
+                    codeGenerationContext.AddReference(parameter.ParameterType);
                 }
             }
 
diff --git a/src/Interception/Interception/CodeGeneration/VirtualMethodProxyGenerator.cs b/src/Interception/Interception/CodeGeneration/VirtualMethodProxyGenerator.cs
index 079a573..d9ea5ae 100644
--- a/src/Interception/Interception/CodeGeneration/VirtualMethodProxyGenerator.cs
+++ b/src/Interception/Interception/CodeGeneration/VirtualMethodProxyGenerator.cs
@@ -203,10 +203,10 @@ public override bool TryGenerate(ServiceDescriptor serviceDescriptor, CodeGenera
 
             foreach (var method in interceptableMethods)
             {                
-                codeGenerationContext.References.Add(method.ReturnType.Assembly);
+                codeGenerationContext.AddReference(method.ReturnType);
                 foreach (var parameter in method.GetParameters())
                 {
-                    codeGenerationContext.References.Add(parameter.ParameterType.Assembly);
+                    codeGenerationContext.AddReference(parameter.ParameterType);
                 }
             }