Avoid allocation of traversed tokens when possible.

Washi1337 · Nov 27, 2022 · c0283c7 · c0283c7
1 parent bebe5d6
commit c0283c7
Show file tree

Hide file tree

Showing 22 changed files with 87 additions and 67 deletions.
diff --git a/src/AsmResolver.DotNet.Dynamic/DynamicTypeSignatureReader.cs b/src/AsmResolver.DotNet.Dynamic/DynamicTypeSignatureReader.cs
@@ -41,7 +41,7 @@ static DynamicTypeSignatureResolver()
         public static bool IsSupported => GetTypeFromHandleUnsafeMethod is not null;
 
         /// <inheritdoc />
-        public override TypeSignature ResolveRuntimeType(in BlobReadContext context, nint address)
+        public override TypeSignature ResolveRuntimeType(ref BlobReadContext context, nint address)
         {
             if (!IsSupported)
                 throw new PlatformNotSupportedException("The current platform does not support the translation of raw type handles to System.Type instances.");

diff --git a/src/AsmResolver.DotNet/Serialized/SerializedFieldDefinition.cs b/src/AsmResolver.DotNet/Serialized/SerializedFieldDefinition.cs
@@ -47,7 +47,8 @@ public SerializedFieldDefinition(ModuleReaderContext context, MetadataToken toke
                     $"Invalid signature blob index in field {MetadataToken.ToString()}.");
             }
 
-            return FieldSignature.FromReader(new BlobReadContext(_context), ref reader);
+            var context = new BlobReadContext(_context);
+            return FieldSignature.FromReader(ref context, ref reader);
         }
 
         /// <inheritdoc />

diff --git a/src/AsmResolver.DotNet/Serialized/SerializedMemberReference.cs b/src/AsmResolver.DotNet/Serialized/SerializedMemberReference.cs
@@ -60,7 +60,8 @@ public SerializedMemberReference(
                     $"Invalid signature blob index in member reference {MetadataToken.ToString()}.");
             }
 
-            return CallingConventionSignature.FromReader(new BlobReadContext(_context), ref reader, true);
+            var context = new BlobReadContext(_context);
+            return CallingConventionSignature.FromReader(ref context, ref reader, true);
         }
 
         /// <inheritdoc />

diff --git a/src/AsmResolver.DotNet/Serialized/SerializedMethodDefinition.cs b/src/AsmResolver.DotNet/Serialized/SerializedMethodDefinition.cs
@@ -52,7 +52,8 @@ public SerializedMethodDefinition(ModuleReaderContext context, MetadataToken tok
                     $"Invalid signature blob index in method {MetadataToken.ToString()}.");
             }
 
-            return MethodSignature.FromReader(new BlobReadContext(_context), ref reader);
+            var context = new BlobReadContext(_context);
+            return MethodSignature.FromReader(ref context, ref reader);
         }
 
         /// <inheritdoc />

diff --git a/src/AsmResolver.DotNet/Serialized/SerializedMethodSpecification.cs b/src/AsmResolver.DotNet/Serialized/SerializedMethodSpecification.cs
@@ -52,7 +52,8 @@ public SerializedMethodSpecification(ModuleReaderContext context, MetadataToken
                     $"Invalid instantiation blob index in method specification {MetadataToken.ToString()}.");
             }
 
-            return GenericInstanceMethodSignature.FromReader(new BlobReadContext(_context), ref reader);
+            var context = new BlobReadContext(_context);
+            return GenericInstanceMethodSignature.FromReader(ref context, ref reader);
         }
 
         /// <inheritdoc />

diff --git a/src/AsmResolver.DotNet/Serialized/SerializedPropertyDefinition.cs b/src/AsmResolver.DotNet/Serialized/SerializedPropertyDefinition.cs
@@ -48,7 +48,8 @@ public SerializedPropertyDefinition(ModuleReaderContext context, MetadataToken t
                     $"Invalid signature blob index in property {MetadataToken.ToString()}.");
             }
 
-            return PropertySignature.FromReader(new BlobReadContext(_context), ref reader);
+            var context = new BlobReadContext(_context);
+            return PropertySignature.FromReader(ref context, ref reader);
         }
 
         /// <inheritdoc />

diff --git a/src/AsmResolver.DotNet/Serialized/SerializedStandAloneSignature.cs b/src/AsmResolver.DotNet/Serialized/SerializedStandAloneSignature.cs
@@ -43,7 +43,8 @@ public SerializedStandAloneSignature(
                     $"Invalid signature blob index in stand-alone signature {MetadataToken.ToString()}.");
             }
 
-            return CallingConventionSignature.FromReader(new BlobReadContext(_context), ref reader);
+            var context = new BlobReadContext(_context);
+            return CallingConventionSignature.FromReader(ref context, ref reader);
         }
 
         /// <inheritdoc />

diff --git a/src/AsmResolver.DotNet/Serialized/SerializedTypeSpecification.cs b/src/AsmResolver.DotNet/Serialized/SerializedTypeSpecification.cs
@@ -42,8 +42,8 @@ public SerializedTypeSpecification(ModuleReaderContext context, MetadataToken to
             }
 
             var context = new BlobReadContext(_context);
-            context.TraversedTokens.Add(MetadataToken);
-            return TypeSignature.FromReader(context, ref reader);
+            context.StepInToken(MetadataToken);
+            return TypeSignature.FromReader(ref context, ref reader);
         }
 
         /// <inheritdoc />

diff --git a/src/AsmResolver.DotNet/Signatures/BlobReadContext.cs b/src/AsmResolver.DotNet/Signatures/BlobReadContext.cs
@@ -10,14 +10,16 @@ namespace AsmResolver.DotNet.Signatures
     /// Provides a context in which a metadata blob parser exists in. This includes the original module reader context
     /// as well as a mechanism to protect against infinite recursion.
     /// </summary>
-    public readonly struct BlobReadContext
+    public struct BlobReadContext
     {
+        private Stack<MetadataToken>? _traversedTokens;
+
         /// <summary>
         /// Creates a new instance of the <see cref="BlobReadContext"/> structure.
         /// </summary>
         /// <param name="readerContext">The original read context.</param>
         public BlobReadContext(ModuleReaderContext readerContext)
-            : this(readerContext, PhysicalTypeSignatureResolver.Instance, Enumerable.Empty<MetadataToken>())
+            : this(readerContext, PhysicalTypeSignatureResolver.Instance)
         {
         }
 
@@ -27,8 +29,10 @@ public BlobReadContext(ModuleReaderContext readerContext)
         /// <param name="readerContext">The original read context.</param>
         /// <param name="resolver">The object responsible for resolving raw type metadata tokens and addresses.</param>
         public BlobReadContext(ModuleReaderContext readerContext, ITypeSignatureResolver resolver)
-            : this(readerContext, resolver, Enumerable.Empty<MetadataToken>())
         {
+            ReaderContext = readerContext;
+            TypeSignatureResolver = resolver;
+            _traversedTokens = null;
         }
 
         /// <summary>
@@ -41,7 +45,7 @@ public BlobReadContext(ModuleReaderContext readerContext, ITypeSignatureResolver
         {
             ReaderContext = readerContext;
             TypeSignatureResolver = resolver;
-            TraversedTokens = new HashSet<MetadataToken>(traversedTokens);
+            _traversedTokens = new Stack<MetadataToken>(traversedTokens);
         }
 
         /// <summary>
@@ -60,12 +64,15 @@ public ITypeSignatureResolver TypeSignatureResolver
             get;
         }
 
-        /// <summary>
-        /// Gets a collection of metadata tokens that were traversed during the parsing of metadata.
-        /// </summary>
-        public ISet<MetadataToken> TraversedTokens
+        public bool StepInToken(MetadataToken token)
         {
-            get;
+            _traversedTokens ??= new Stack<MetadataToken>();
+
+            if (_traversedTokens.Contains(token))
+                return false;
+
+            _traversedTokens.Push(token);
+            return true;
         }
     }
 }
diff --git a/src/AsmResolver.DotNet/Signatures/CallingConventionSignature.cs b/src/AsmResolver.DotNet/Signatures/CallingConventionSignature.cs
@@ -20,11 +20,11 @@ public abstract class CallingConventionSignature : ExtendableBlobSignature, IImp
         /// put into the <see cref="ExtendableBlobSignature.ExtraData"/> property.</param>
         /// <returns>The read signature.</returns>
         public static CallingConventionSignature? FromReader(
-            in BlobReadContext context,
+            ref BlobReadContext context,
             ref BinaryStreamReader reader,
             bool readToEnd = true)
         {
-            var signature = ReadSignature(context, ref reader);
+            var signature = ReadSignature(ref context, ref reader);
             if (readToEnd)
             {
                 byte[] extraData = reader.ReadToEnd();
@@ -35,7 +35,7 @@ public abstract class CallingConventionSignature : ExtendableBlobSignature, IImp
             return signature;
         }
 
-        private static CallingConventionSignature? ReadSignature(in BlobReadContext context, ref BinaryStreamReader reader)
+        private static CallingConventionSignature? ReadSignature(ref BlobReadContext context, ref BinaryStreamReader reader)
         {
             byte flag = reader.ReadByte();
             reader.Offset--;
@@ -50,19 +50,19 @@ public abstract class CallingConventionSignature : ExtendableBlobSignature, IImp
                 case CallingConventionAttributes.ThisCall:
                 case CallingConventionAttributes.VarArg:
                 case CallingConventionAttributes.Unmanaged:
-                    return MethodSignature.FromReader(context, ref reader);
+                    return MethodSignature.FromReader(ref context, ref reader);
 
                 case CallingConventionAttributes.Property:
-                    return PropertySignature.FromReader(context, ref reader);
+                    return PropertySignature.FromReader(ref context, ref reader);
 
                 case CallingConventionAttributes.Local:
-                    return LocalVariablesSignature.FromReader(context, ref reader);
+                    return LocalVariablesSignature.FromReader(ref context, ref reader);
 
                 case CallingConventionAttributes.GenericInstance:
-                    return GenericInstanceMethodSignature.FromReader(context, ref reader);
+                    return GenericInstanceMethodSignature.FromReader(ref context, ref reader);
 
                 case CallingConventionAttributes.Field:
-                    return FieldSignature.FromReader(context, ref reader);
+                    return FieldSignature.FromReader(ref context, ref reader);
             }
 
             throw new NotSupportedException($"Invalid or unsupported calling convention signature header {flag:X2}.");

diff --git a/src/AsmResolver.DotNet/Signatures/FieldSignature.cs b/src/AsmResolver.DotNet/Signatures/FieldSignature.cs
@@ -37,11 +37,11 @@ public static FieldSignature CreateInstance(TypeSignature fieldType)
         /// <param name="context">The blob reader context.</param>
         /// <param name="reader">The blob input stream.</param>
         /// <returns>The field signature.</returns>
-        public static FieldSignature FromReader(in BlobReadContext context, ref BinaryStreamReader reader)
+        public static FieldSignature FromReader(ref BlobReadContext context, ref BinaryStreamReader reader)
         {
             return new(
                 (CallingConventionAttributes) reader.ReadByte(),
-                TypeSignature.FromReader(context, ref reader));
+                TypeSignature.FromReader(ref context, ref reader));
         }
 
         /// <summary>

diff --git a/src/AsmResolver.DotNet/Signatures/GenericInstanceMethodSignature.cs b/src/AsmResolver.DotNet/Signatures/GenericInstanceMethodSignature.cs
@@ -11,7 +11,7 @@ namespace AsmResolver.DotNet.Signatures
     public class GenericInstanceMethodSignature : CallingConventionSignature, IGenericArgumentsProvider
     {
         internal static GenericInstanceMethodSignature? FromReader(
-            in BlobReadContext context,
+            ref BlobReadContext context,
             ref BinaryStreamReader reader)
         {
             if (!reader.CanRead(sizeof(byte)))
@@ -30,7 +30,7 @@ public class GenericInstanceMethodSignature : CallingConventionSignature, IGener
 
             var result = new GenericInstanceMethodSignature(attributes, (int) count);
             for (int i = 0; i < count; i++)
-                result.TypeArguments.Add(TypeSignature.FromReader(context, ref reader));
+                result.TypeArguments.Add(TypeSignature.FromReader(ref context, ref reader));
 
             return result;
         }

diff --git a/src/AsmResolver.DotNet/Signatures/LocalVariablesSignature.cs b/src/AsmResolver.DotNet/Signatures/LocalVariablesSignature.cs
@@ -16,7 +16,7 @@ public class LocalVariablesSignature : CallingConventionSignature
         /// <param name="context">The blob reader context.</param>
         /// <param name="reader">The input stream.</param>
         /// <returns>The signature.</returns>
-        public static LocalVariablesSignature FromReader(in BlobReadContext context, ref BinaryStreamReader reader)
+        public static LocalVariablesSignature FromReader(ref BlobReadContext context, ref BinaryStreamReader reader)
         {
             var result = new LocalVariablesSignature();
             result.Attributes = (CallingConventionAttributes) reader.ReadByte();
@@ -28,7 +28,7 @@ public static LocalVariablesSignature FromReader(in BlobReadContext context, ref
             }
 
             for (int i = 0; i < count; i++)
-                result.VariableTypes.Add(TypeSignature.FromReader(context, ref reader));
+                result.VariableTypes.Add(TypeSignature.FromReader(ref context, ref reader));
 
             return result;
         }

diff --git a/src/AsmResolver.DotNet/Signatures/MethodSignature.cs b/src/AsmResolver.DotNet/Signatures/MethodSignature.cs
@@ -18,7 +18,7 @@ public class MethodSignature : MethodSignatureBase
         /// <param name="context">The blob reader context.</param>
         /// <param name="reader">The blob input stream.</param>
         /// <returns>The method signature.</returns>
-        public static MethodSignature FromReader(in BlobReadContext context, ref BinaryStreamReader reader)
+        public static MethodSignature FromReader(ref BlobReadContext context, ref BinaryStreamReader reader)
         {
             var result = new MethodSignature(
                 (CallingConventionAttributes) reader.ReadByte(),
@@ -37,7 +37,7 @@ public static MethodSignature FromReader(in BlobReadContext context, ref BinaryS
                 result.GenericParameterCount = (int) genericParameterCount;
             }
 
-            result.ReadParametersAndReturnType(context, ref reader);
+            result.ReadParametersAndReturnType(ref context, ref reader);
             return result;
         }
 

diff --git a/src/AsmResolver.DotNet/Signatures/MethodSignatureBase.cs b/src/AsmResolver.DotNet/Signatures/MethodSignatureBase.cs
@@ -106,7 +106,7 @@ public override bool IsImportedInModule(ModuleDefinition module)
         /// </summary>
         /// <param name="context">The blob reader context.</param>
         /// <param name="reader">The input stream.</param>
-        protected void ReadParametersAndReturnType(in BlobReadContext context, ref BinaryStreamReader reader)
+        protected void ReadParametersAndReturnType(ref BlobReadContext context, ref BinaryStreamReader reader)
         {
             // Parameter count.
             if (!reader.TryReadCompressedUInt32(out uint parameterCount))
@@ -116,14 +116,14 @@ protected void ReadParametersAndReturnType(in BlobReadContext context, ref Binar
             }
 
             // Return type.
-            ReturnType = TypeSignature.FromReader(context, ref reader);
+            ReturnType = TypeSignature.FromReader(ref context, ref reader);
 
             // Parameter types.
             _parameterTypes.Capacity = (int) parameterCount;
             bool sentinel = false;
             for (int i = 0; i < parameterCount; i++)
             {
-                var parameterType = TypeSignature.FromReader(context, ref reader);
+                var parameterType = TypeSignature.FromReader(ref context, ref reader);
 
                 if (parameterType.ElementType == ElementType.Sentinel)
                 {

diff --git a/src/AsmResolver.DotNet/Signatures/PropertySignature.cs b/src/AsmResolver.DotNet/Signatures/PropertySignature.cs
@@ -17,7 +17,7 @@ public class PropertySignature : MethodSignatureBase
         /// <param name="context">The blob reader context.</param>
         /// <param name="reader">The blob input stream.</param>
         /// <returns>The property signature.</returns>
-        public static PropertySignature? FromReader(in BlobReadContext context, ref BinaryStreamReader reader)
+        public static PropertySignature? FromReader(ref BlobReadContext context, ref BinaryStreamReader reader)
         {
             var attributes = (CallingConventionAttributes) reader.ReadByte();
             if ((attributes & CallingConventionAttributes.Property) == 0)
@@ -31,7 +31,7 @@ public class PropertySignature : MethodSignatureBase
                 context.ReaderContext.ParentModule.CorLibTypeFactory.Object,
                 Enumerable.Empty<TypeSignature>());
 
-            result.ReadParametersAndReturnType(context, ref reader);
+            result.ReadParametersAndReturnType(ref context, ref reader);
             return result;
         }
 

diff --git a/src/AsmResolver.DotNet/Signatures/Types/ArrayTypeSignature.cs b/src/AsmResolver.DotNet/Signatures/Types/ArrayTypeSignature.cs
@@ -69,9 +69,9 @@ public IList<ArrayDimension> Dimensions
             get;
         }
 
-        internal new static ArrayTypeSignature FromReader(in BlobReadContext context, ref BinaryStreamReader reader)
+        internal new static ArrayTypeSignature FromReader(ref BlobReadContext context, ref BinaryStreamReader reader)
         {
-            var signature = new ArrayTypeSignature(TypeSignature.FromReader(context, ref reader));
+            var signature = new ArrayTypeSignature(TypeSignature.FromReader(ref context, ref reader));
 
             // Rank
             if (!reader.TryReadCompressedUInt32(out uint rank))

diff --git a/src/AsmResolver.DotNet/Signatures/Types/GenericInstanceTypeSignature.cs b/src/AsmResolver.DotNet/Signatures/Types/GenericInstanceTypeSignature.cs
@@ -14,9 +14,9 @@ public class GenericInstanceTypeSignature : TypeSignature, IGenericArgumentsProv
         private ITypeDefOrRef _genericType;
         private bool _isValueType;
 
-        internal new static GenericInstanceTypeSignature FromReader(in BlobReadContext context, ref BinaryStreamReader reader)
+        internal new static GenericInstanceTypeSignature FromReader(ref BlobReadContext context, ref BinaryStreamReader reader)
         {
-            var genericType = TypeSignature.FromReader(context, ref reader);
+            var genericType = TypeSignature.FromReader(ref context, ref reader);
             var signature = new GenericInstanceTypeSignature(genericType.ToTypeDefOrRef(), genericType.ElementType == ElementType.ValueType);
 
             if (!reader.TryReadCompressedUInt32(out uint count))
@@ -27,7 +27,7 @@ public class GenericInstanceTypeSignature : TypeSignature, IGenericArgumentsProv
 
             signature._typeArguments.Capacity = (int) count;
             for (int i = 0; i < count; i++)
-                signature._typeArguments.Add(TypeSignature.FromReader(context, ref reader));
+                signature._typeArguments.Add(TypeSignature.FromReader(ref context, ref reader));
 
             return signature;
         }

diff --git a/src/AsmResolver.DotNet/Signatures/Types/ITypeSignatureResolver.cs b/src/AsmResolver.DotNet/Signatures/Types/ITypeSignatureResolver.cs
@@ -13,15 +13,15 @@ public interface ITypeSignatureResolver
         /// <param name="context">The blob reading context the type is situated in.</param>
         /// <param name="token">The token to resolve.</param>
         /// <returns>The type.</returns>
-        ITypeDefOrRef ResolveToken(in BlobReadContext context, MetadataToken token);
+        ITypeDefOrRef ResolveToken(ref BlobReadContext context, MetadataToken token);
 
         /// <summary>
         /// Resolves an address to a runtime method table to a type signature.
         /// </summary>
         /// <param name="context">The blob reading context the type is situated in.</param>
         /// <param name="address">The address to resolve.</param>
         /// <returns>The type.</returns>
-        TypeSignature ResolveRuntimeType(in BlobReadContext context, nint address);
+        TypeSignature ResolveRuntimeType(ref BlobReadContext context, nint address);
     }
 
 }
diff --git a/src/AsmResolver.DotNet/Signatures/Types/PhysicalTypeSignatureResolver.cs b/src/AsmResolver.DotNet/Signatures/Types/PhysicalTypeSignatureResolver.cs
@@ -18,12 +18,12 @@ public static PhysicalTypeSignatureResolver Instance
         } = new();
 
         /// <inheritdoc />
-        public virtual ITypeDefOrRef ResolveToken(in BlobReadContext context, MetadataToken token)
+        public virtual ITypeDefOrRef ResolveToken(ref BlobReadContext context, MetadataToken token)
         {
             switch (token.Table)
             {
                 // Check for infinite recursion.
-                case TableIndex.TypeSpec when !context.TraversedTokens.Add(token):
+                case TableIndex.TypeSpec when !context.StepInToken(token):
                     context.ReaderContext.BadImage("Infinite metadata loop was detected.");
                     return InvalidTypeDefOrRef.Get(InvalidTypeSignatureError.MetadataLoop);
 
@@ -47,7 +47,7 @@ public virtual ITypeDefOrRef ResolveToken(in BlobReadContext context, MetadataTo
         }
 
         /// <inheritdoc />
-        public virtual TypeSignature ResolveRuntimeType(in BlobReadContext context, nint address)
+        public virtual TypeSignature ResolveRuntimeType(ref BlobReadContext context, nint address)
         {
             throw new NotSupportedException(
                 "Encountered an COR_ELEMENT_TYPE_INTERNAL type signature which is not supported by this "