Merge pull request #493 from Washi1337/feature/ca-type-arg-cloning

Let MemberCloner keep track of included types by value

src/AsmResolver.DotNet/Cloning/CloneContextAwareReferenceImporter.cs

@@ -1,3 +1,5 @@
+using System;
+
 namespace AsmResolver.DotNet.Cloning
 {
     /// <summary>
@@ -49,11 +51,16 @@ public override IMethodDefOrRef ImportMethod(IMethodDefOrRef method)
         /// <inheritdoc />
         protected override ITypeDefOrRef ImportType(TypeReference type)
         {
-            return type.Namespace == "System"
-                && type.Name == nameof(System.Object)
-                && (type.Scope?.GetAssembly()?.IsCorLib ?? false)
-                ? _context.Module.CorLibTypeFactory.Object.Type
-                : base.ImportType(type);
+            // Special case for System.Object.
+            if (type.IsTypeOf(nameof(System), nameof(Object)) && (type.Scope?.GetAssembly()?.IsCorLib ?? false))
+                return _context.Module.CorLibTypeFactory.Object.Type;
+
+            // Rare case where a type reference could point to one of the included type definitions
+            // (e.g., in custom attributes type arguments, see https://github.com/Washi1337/AsmResolver/issues/482).
+            if (_context.ClonedTypes.TryGetValue(type, out var clonedType))
+                return (ITypeDefOrRef) clonedType;
+
+            return base.ImportType(type);
         }
     }
 }

src/AsmResolver.DotNet/Cloning/MemberCloneContext.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using AsmResolver.DotNet.Signatures;
 
 namespace AsmResolver.DotNet.Cloning
 {
@@ -49,5 +50,17 @@ public CloneContextAwareReferenceImporter Importer
         {
             get;
         } = new Dictionary<IMemberDescriptor, IMemberDescriptor>();
+
+        /// <summary>
+        /// Gets a mapping of original types to their cloned counterparts.
+        /// </summary>
+        /// <remarks>
+        /// This dictionary performs lookups based on value using a <see cref="SignatureComparer"/> instead of object
+        /// identity, and can thus be used to translate type references to included type definitions.
+        /// </remarks>
+        public IDictionary<ITypeDescriptor, ITypeDescriptor> ClonedTypes
+        {
+            get;
+        } = new Dictionary<ITypeDescriptor, ITypeDescriptor>(SignatureComparer.Default);
     }
 }

src/AsmResolver.DotNet/Cloning/MemberCloner.cs

@@ -3,6 +3,7 @@
 using AsmResolver.DotNet.Signatures;
 using AsmResolver.DotNet.Signatures.Marshal;
 using AsmResolver.DotNet.Signatures.Security;
+using AsmResolver.DotNet.Signatures.Types;
 using AsmResolver.PE.DotNet.Metadata.Tables;
 
 namespace AsmResolver.DotNet.Cloning
@@ -320,6 +321,7 @@ private static void CreateTypeStub(MemberCloneContext context, TypeDefinition ty
 
             var typeStub = new TypeDefinition(type.Namespace, type.Name, type.Attributes);
             context.ClonedMembers.Add(type, typeStub);
+            context.ClonedTypes.Add(type, typeStub);
         }
 
         private void DeepCopyMembers(MemberCloneContext context)
@@ -442,11 +444,19 @@ private static CustomAttributeArgument CloneCustomAttributeArgument(MemberCloneC
 
             // Copy all elements.
             for (int i = 0; i < argument.Elements.Count; i++)
-                clonedArgument.Elements.Add(argument.Elements[i]);
+                clonedArgument.Elements.Add(CloneElement(context, argument.Elements[i]));
 
             return clonedArgument;
         }
 
+        private static object? CloneElement(MemberCloneContext context, object? element)
+        {
+            if (element is TypeSignature type)
+                return context.Importer.ImportTypeSignature(type);
+
+            return element;
+        }
+
         private static ImplementationMap? CloneImplementationMap(MemberCloneContext context, ImplementationMap? map)
         {
             if (map is null)

src/AsmResolver.DotNet/Signatures/Types/CorLibTypeFactory.cs

@@ -11,16 +11,6 @@ namespace AsmResolver.DotNet.Signatures.Types
     /// </summary>
     public class CorLibTypeFactory
     {
-        /// <summary>
-        /// Creates a new type factory that references mscorlib 4.0.0.0.
-        /// </summary>
-        /// <returns>The factory.</returns>
-        public static CorLibTypeFactory CreateMscorlib40TypeFactory(ModuleDefinition module)
-        {
-            var importer = new ReferenceImporter(module);
-            return new CorLibTypeFactory(importer.ImportScope(KnownCorLibs.MsCorLib_v4_0_0_0));
-        }
-
         private CorLibTypeSignature? _void;
         private CorLibTypeSignature? _boolean;
         private CorLibTypeSignature? _char;
@@ -146,6 +136,16 @@ public IResolutionScope CorLibScope
         /// Gets the element type signature for <see cref="System.Object"/>.
         /// </summary>
         public CorLibTypeSignature Object => GetOrCreateCorLibTypeSignature(ref _object, ElementType.Object, nameof(Object));
+
+        /// <summary>
+        /// Creates a new type factory that references mscorlib 4.0.0.0.
+        /// </summary>
+        /// <returns>The factory.</returns>
+        public static CorLibTypeFactory CreateMscorlib40TypeFactory(ModuleDefinition module)
+        {
+            var importer = new ReferenceImporter(module);
+            return new CorLibTypeFactory(importer.ImportScope(KnownCorLibs.MsCorLib_v4_0_0_0));
+        }
 
         /// <summary>
         /// Transforms the provided type descriptor to a common object runtime type signature.

test/AsmResolver.DotNet.Tests/Cloning/MetadataClonerTest.cs

@@ -4,6 +4,8 @@
 using System.Reflection;
 using AsmResolver.DotNet.Cloning;
 using AsmResolver.DotNet.Signatures;
+using AsmResolver.DotNet.Signatures.Types;
+using AsmResolver.DotNet.TestCases.CustomAttributes;
 using AsmResolver.DotNet.TestCases.Events;
 using AsmResolver.DotNet.TestCases.Fields;
 using AsmResolver.DotNet.TestCases.Generics;
@@ -408,5 +410,31 @@ public void CloneAndInjectAndAssignToken()
             Assert.All(result.ClonedTopLevelTypes, t => Assert.Contains(t, targetModule.TopLevelTypes));
             Assert.All(result.ClonedMembers, m => Assert.NotEqual(0u, ((IMetadataMember) m).MetadataToken.Rid));
         }
+
+        [Fact]
+        public void CloneIncludedTypeArgument()
+        {
+            // https://github.com/Washi1337/AsmResolver/issues/482
+
+            var sourceModule = ModuleDefinition.FromFile(typeof(CustomAttributesTestClass).Assembly.Location);
+            var targetModule = PrepareTempModule();
+
+            var type = sourceModule.LookupMember<TypeDefinition>(typeof(TestEnum).MetadataToken);
+            var method = sourceModule.LookupMember<MethodDefinition>(typeof(CustomAttributesTestClass)
+                .GetMethod(nameof(CustomAttributesTestClass.FIxedLocalTypeArgument))!
+                .MetadataToken);
+
+            var result = new MemberCloner(targetModule)
+                .Include(type)
+                .Include(method)
+                .AddListener(new InjectTypeClonerListener(targetModule))
+                .Clone();
+
+            var newType = result.GetClonedMember(type);
+            var newMethod = result.GetClonedMember(method);
+
+            var newArgument = Assert.IsAssignableFrom<ITypeDescriptor>(newMethod.CustomAttributes[0].Signature!.FixedArguments[0].Element);
+            Assert.Equal(newType, newArgument, _signatureComparer);
+        }
     }
 }

test/TestBinaries/DotNet/AsmResolver.DotNet.TestCases.CustomAttributes/CustomAttributesTestClass.cs

@@ -54,6 +54,11 @@ public void FixedComplexTypeArgument()
         {
         }
 
+        [TestCase(typeof(TestEnum))]
+        public void FIxedLocalTypeArgument()
+        {
+        }
+
         [TestCase(2, "Fixed arg", TestEnum.Value3)]
         public void FixedMultipleArguments()
         {