From cc1e79472c46c83162067f9bdccbe2b2d6ce613e Mon Sep 17 00:00:00 2001
From: Washi <washi1337@protonmail.com>
Date: Mon, 12 Jun 2023 15:18:32 +0200
Subject: [PATCH] BUGFIX: Handle type specifications in GetImpliedMemoryLayout

---
 .../Memory/TypeMemoryLayoutDetector.cs        |  6 +++
 .../Memory/SequentialStructLayoutTest.cs      | 45 +++++++++++--------
 .../Memory/SequentialTestStructs.cs           | 23 ++++++++--
 .../Memory/StructLayoutTestBase.cs            | 10 ++---
 4 files changed, 55 insertions(+), 29 deletions(-)

diff --git a/src/AsmResolver.DotNet/Memory/TypeMemoryLayoutDetector.cs b/src/AsmResolver.DotNet/Memory/TypeMemoryLayoutDetector.cs
index 2ae1e422c..588c9dae9 100644
--- a/src/AsmResolver.DotNet/Memory/TypeMemoryLayoutDetector.cs
+++ b/src/AsmResolver.DotNet/Memory/TypeMemoryLayoutDetector.cs
@@ -150,10 +150,16 @@ public TypeMemoryLayout VisitTypeDefOrRef(ITypeDefOrRef type)
             {
                 TableIndex.TypeRef => VisitTypeReference((TypeReference) type),
                 TableIndex.TypeDef => VisitTypeDefinition((TypeDefinition) type),
+                TableIndex.TypeSpec => VisitTypeSpecification((TypeSpecification) type),
                 _ => throw new ArgumentException("Invalid type.")
             };
         }
 
+        private TypeMemoryLayout VisitTypeSpecification(TypeSpecification type)
+        {
+            return type.Signature!.AcceptVisitor(this);
+        }
+
         private TypeMemoryLayout VisitTypeReference(TypeReference type) =>
             VisitTypeDefinition(type.Resolve() ?? throw new ArgumentException(
                 $"Could not resolve type {type.SafeToString()}."));
diff --git a/test/AsmResolver.DotNet.Tests/Memory/SequentialStructLayoutTest.cs b/test/AsmResolver.DotNet.Tests/Memory/SequentialStructLayoutTest.cs
index 7d15df344..dd2e84ac2 100644
--- a/test/AsmResolver.DotNet.Tests/Memory/SequentialStructLayoutTest.cs
+++ b/test/AsmResolver.DotNet.Tests/Memory/SequentialStructLayoutTest.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Runtime.CompilerServices;
 using AsmResolver.DotNet.Memory;
 using AsmResolver.PE.DotNet.Metadata.Tables.Rows;
@@ -51,62 +52,68 @@ public void ReferenceTypedCorLibTypeShouldReturnElementSize(ElementType elementT
         [Fact]
         public void MultipleFieldsSequentialStructDefaultPack() =>
             VerifySize<SequentialTestStructs.MultipleFieldsSequentialStructDefaultPack>();
-        
+
         [Fact]
         public void MultipleFieldsSequentialStructPack1() =>
             VerifySize<SequentialTestStructs.MultipleFieldsSequentialStructPack1>();
-        
+
         [Fact]
         public void LargeAndSmallFieldSequentialDefaultPack() =>
             VerifySize<SequentialTestStructs.LargeAndSmallFieldSequentialDefaultPack>();
-        
+
         [Fact]
         public void NestedStruct1() => VerifySize<SequentialTestStructs.NestedStruct1>();
-        
+
         [Fact]
         public void NestedStruct2() => VerifySize<SequentialTestStructs.NestedStruct2>();
-        
+
         [Fact]
         public void NestedStructWithEnclosingPack1() => VerifySize<SequentialTestStructs.NestedStructWithEnclosingPack1>();
-        
+
         [Fact]
         public void NestedStructWithNestedPack1() => VerifySize<SequentialTestStructs.NestedStructWithNestedPack1>();
-        
+
         [Fact]
         public void NestedStructInNestedStruct() => VerifySize<SequentialTestStructs.NestedStructInNestedStruct>();
-        
+
         [Fact]
-        public void ThreeLevelsNestingSequentialStructDefaultPack() => 
+        public void ThreeLevelsNestingSequentialStructDefaultPack() =>
             VerifySize<SequentialTestStructs.ThreeLevelsNestingSequentialStructDefaultPack>();
-        
+
         [Fact]
-        public void ThreeLevelsNestingSequentialStructPack1() => 
+        public void ThreeLevelsNestingSequentialStructPack1() =>
             VerifySize<SequentialTestStructs.ThreeLevelsNestingSequentialWithNestedStructPack1>();
-        
+
         [Fact]
         public void ExplicitlySizedEmptyStruct() => VerifySize<SequentialTestStructs.ExplicitlySizedEmptyStruct>();
-        
+
         [Fact]
         public void ExplicitlySizedSingleField() => VerifySize<SequentialTestStructs.ExplicitlySizedSingleField>();
-        
+
         [Fact]
         public void ExplicitlySizedSmallerExplicitSizeThanActualSize() =>
             VerifySize<SequentialTestStructs.ExplicitlySizedSmallerExplicitSizeThanActualSize>();
-        
+
         [Fact]
         public void StructWithPrimitiveFieldSmallerThanPack() =>
             VerifySize<SequentialTestStructs.StructWithPrimitiveFieldSmallerThanPack>();
-        
+
         [Fact]
         public void StructWithStructFieldSmallerThanPack() =>
             VerifySize<SequentialTestStructs.StructWithStructFieldSmallerThanPack>();
-        
+
         [Fact]
         public void PackLargerThanLargestField() =>
             VerifySize<SequentialTestStructs.PackLargerThanLargestField>();
-        
+
         [Fact]
         public void PackLargerThanLargestFieldWithImplicitAlignment() =>
             VerifySize<SequentialTestStructs.PackLargerThanLargestFieldWithImplicitAlignment>();
+
+        [Fact]
+        public void GenericStruct() => VerifySize<SequentialTestStructs.GenericStruct<int, byte>>();
+
+        [Fact]
+        public void GenericNestedStruct() => VerifySize<SequentialTestStructs.GenericStruct<int, byte>.NestedStruct>();
     }
-}
\ No newline at end of file
+}
diff --git a/test/AsmResolver.DotNet.Tests/Memory/SequentialTestStructs.cs b/test/AsmResolver.DotNet.Tests/Memory/SequentialTestStructs.cs
index a4b75b9f3..b24d05764 100644
--- a/test/AsmResolver.DotNet.Tests/Memory/SequentialTestStructs.cs
+++ b/test/AsmResolver.DotNet.Tests/Memory/SequentialTestStructs.cs
@@ -9,7 +9,7 @@ public static class SequentialTestStructs
         [StructLayout(LayoutKind.Sequential)]
         public struct EmptyStruct
         {
-            
+
         }
 
         [StructLayout(LayoutKind.Sequential)]
@@ -134,7 +134,7 @@ public struct FixedSizeStruct133
         public struct PackLargerThanLargestField
         {
             public FixedSizeStruct133 Field1;
-            
+
             public FixedSizeStruct133 Field2;
         }
 
@@ -148,8 +148,23 @@ public struct FixedSizeStruct133WithField
         public struct PackLargerThanLargestFieldWithImplicitAlignment
         {
             public FixedSizeStruct133WithField Field1;
-            
+
             public FixedSizeStruct133WithField Field2;
         }
+
+        [StructLayout(LayoutKind.Sequential, Pack = 1)]
+        public struct GenericStruct<T1, T2>
+        {
+            public T1 Field1;
+            public T2 Field2;
+
+            [StructLayout(LayoutKind.Sequential, Pack = 1)]
+            public struct NestedStruct
+            {
+                public T1 Field1;
+                public T2 Field2;
+            }
+        }
+
     }
-}
\ No newline at end of file
+}
diff --git a/test/AsmResolver.DotNet.Tests/Memory/StructLayoutTestBase.cs b/test/AsmResolver.DotNet.Tests/Memory/StructLayoutTestBase.cs
index 26a4c5f4e..9f885581a 100644
--- a/test/AsmResolver.DotNet.Tests/Memory/StructLayoutTestBase.cs
+++ b/test/AsmResolver.DotNet.Tests/Memory/StructLayoutTestBase.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Linq;
 using System.Runtime.CompilerServices;
 using AsmResolver.DotNet.Memory;
 using Xunit;
@@ -17,11 +18,8 @@ protected ModuleDefinition Module
             get;
         }
 
-        protected TypeDefinition FindTestType(Type type)
-        {
-            return (TypeDefinition) Module.LookupMember(type.MetadataToken);
-        }     
-        
+        private ITypeDescriptor FindTestType(Type type) => Module.DefaultImporter.ImportType(type);
+
         protected void VerifySize<T>()
         {
             var type = FindTestType(typeof(T));
@@ -30,4 +28,4 @@ protected void VerifySize<T>()
         }
 
     }
-}
\ No newline at end of file
+}