Merge pull request #351 from Washi1337/feature/raw-ca-blobs

Lazy custom attribute signatures

src/AsmResolver.DotNet/CustomAttribute.cs

@@ -25,6 +25,17 @@ protected CustomAttribute(MetadataToken token)
             _signature = new LazyVariable<CustomAttributeSignature?>(GetSignature);
         }
 
+        /// <summary>
+        /// Creates a new custom attribute.
+        /// </summary>
+        /// <param name="constructor">The constructor of the attribute to call.</param>
+        public CustomAttribute(ICustomAttributeType? constructor)
+            : this(new MetadataToken(TableIndex.CustomAttribute, 0))
+        {
+            Constructor = constructor;
+            Signature = new CustomAttributeSignature();
+        }
+
         /// <summary>
         /// Creates a new custom attribute.
         /// </summary>

src/AsmResolver.DotNet/Serialized/SerializedCustomAttribute.cs

@@ -65,10 +65,7 @@ public SerializedCustomAttribute(ModuleReaderContext context, MetadataToken toke
                     $"Invalid signature blob index in custom attribute {MetadataToken}.");
             }
 
-            return CustomAttributeSignature.FromReader(
-                new BlobReadContext(_context),
-                Constructor,
-                ref reader);
+            return CustomAttributeSignature.FromReader(new BlobReadContext(_context), Constructor, reader);
         }
     }
 }

src/AsmResolver.DotNet/Signatures/CustomAttributeSignature.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using AsmResolver.DotNet.Signatures.Types;
 using AsmResolver.IO;
@@ -12,55 +13,26 @@ namespace AsmResolver.DotNet.Signatures
     /// </summary>
     public class CustomAttributeSignature : ExtendableBlobSignature
     {
-        private readonly List<CustomAttributeArgument> _fixedArguments;
-        private readonly List<CustomAttributeNamedArgument> _namedArguments;
-        private const ushort CustomAttributeSignaturePrologue = 0x0001;
-
         /// <summary>
-        /// Reads a single custom attribute signature from the input stream.
+        /// The header value of every custom attribute signature.
         /// </summary>
-        /// <param name="context">The blob reader context.</param>
-        /// <param name="ctor">The constructor that was called.</param>
-        /// <param name="reader">The input stream.</param>
-        /// <returns>The signature.</returns>
-        /// <exception cref="FormatException">Occurs when the input stream does not point to a valid signature.</exception>
-        public static CustomAttributeSignature? FromReader(in BlobReadContext context, ICustomAttributeType ctor, ref BinaryStreamReader reader)
-        {
-            ushort prologue = reader.ReadUInt16();
-            if (prologue != CustomAttributeSignaturePrologue)
-            {
-                context.ReaderContext.BadImage("Input stream does not point to a valid custom attribute signature.");
-                return null;
-            }
-
-            var result = new CustomAttributeSignature();
-
-            // Read fixed arguments.
-            var parameterTypes = ctor.Signature?.ParameterTypes ?? Array.Empty<TypeSignature>();
-            result._fixedArguments.Capacity = parameterTypes.Count;
-            for (int i = 0; i < parameterTypes.Count; i++)
-            {
-                var argument = CustomAttributeArgument.FromReader(context, parameterTypes[i], ref reader);
-                result._fixedArguments.Add(argument);
-            }
-
-            // Read named arguments.
-            ushort namedArgumentCount = reader.ReadUInt16();
-            result._namedArguments.Capacity = namedArgumentCount;
-            for (int i = 0; i < namedArgumentCount; i++)
-            {
-                var argument = CustomAttributeNamedArgument.FromReader(context, ref reader);
-                result._namedArguments.Add(argument);
-            }
+        protected const ushort CustomAttributeSignaturePrologue = 0x0001;
 
-            return result;
-        }
+        private List<CustomAttributeArgument>? _fixedArguments;
+        private List<CustomAttributeNamedArgument>? _namedArguments;
 
         /// <summary>
         /// Creates a new empty custom attribute signature.
         /// </summary>
         public CustomAttributeSignature()
-            : this(Enumerable.Empty<CustomAttributeArgument>(), Enumerable.Empty<CustomAttributeNamedArgument>())
+        {
+        }
+
+        /// <summary>
+        /// Creates a new custom attribute signature with the provided fixed arguments.
+        /// </summary>
+        public CustomAttributeSignature(params CustomAttributeArgument[] fixedArguments)
+            : this(fixedArguments, Enumerable.Empty<CustomAttributeNamedArgument>())
         {
         }
 
@@ -84,12 +56,86 @@ public CustomAttributeSignature(IEnumerable<CustomAttributeArgument> fixedArgume
         /// <summary>
         /// Gets a collection of fixed arguments that are passed onto the constructor of the attribute.
         /// </summary>
-        public IList<CustomAttributeArgument> FixedArguments => _fixedArguments;
+        public IList<CustomAttributeArgument> FixedArguments
+        {
+            get
+            {
+                EnsureIsInitialized();
+                return _fixedArguments;
+            }
+        }
 
         /// <summary>
         /// Gets a collection of values that are assigned to fields and/or members of the attribute class.
         /// </summary>
-        public IList<CustomAttributeNamedArgument> NamedArguments => _namedArguments;
+        public IList<CustomAttributeNamedArgument> NamedArguments
+        {
+            get
+            {
+                EnsureIsInitialized();
+                return _namedArguments;
+            }
+        }
+
+        /// <summary>
+        /// Reads a single custom attribute signature from the input stream.
+        /// </summary>
+        /// <param name="context">The blob reader context.</param>
+        /// <param name="ctor">The constructor that was called.</param>
+        /// <param name="reader">The input stream.</param>
+        /// <returns>The signature.</returns>
+        /// <exception cref="FormatException">Occurs when the input stream does not point to a valid signature.</exception>
+        public static CustomAttributeSignature FromReader(
+            in BlobReadContext context,
+            ICustomAttributeType ctor,
+            in BinaryStreamReader reader)
+        {
+            var argumentTypes = ctor.Signature?.ParameterTypes ?? Array.Empty<TypeSignature>();
+            return new SerializedCustomAttributeSignature(context, argumentTypes, reader);
+        }
+
+        /// <summary>
+        /// Gets a value indicating whether the <see cref="FixedArguments"/> and <see cref="NamedArguments"/> collections
+        /// are initialized or not.
+        /// </summary>
+        [MemberNotNullWhen(true, nameof(_fixedArguments))]
+        [MemberNotNullWhen(true, nameof(_namedArguments))]
+        protected bool IsInitialized => _fixedArguments is not null && _namedArguments is not null;
+
+        /// <summary>
+        /// Ensures that the <see cref="FixedArguments"/> and <see cref="NamedArguments"/> are initialized.
+        /// </summary>
+        [MemberNotNull(nameof(_fixedArguments))]
+        [MemberNotNull(nameof(_namedArguments))]
+        protected void EnsureIsInitialized()
+        {
+            if (IsInitialized)
+                return;
+
+            var fixedArguments = new List<CustomAttributeArgument>();
+            var namedArguments = new List<CustomAttributeNamedArgument>();
+
+            Initialize(fixedArguments, namedArguments);
+
+            lock (this)
+            {
+                if (IsInitialized)
+                    return;
+                _fixedArguments = fixedArguments;
+                _namedArguments = namedArguments;
+            }
+        }
+
+        /// <summary>
+        /// Initializes the argument collections of the signature.
+        /// </summary>
+        /// <param name="fixedArguments">The collection that will receive the fixed arguments.</param>
+        /// <param name="namedArguments">The collection that will receive the named arguments.</param>
+        protected virtual void Initialize(
+            IList<CustomAttributeArgument> fixedArguments,
+            IList<CustomAttributeNamedArgument> namedArguments)
+        {
+        }
 
         /// <inheritdoc />
         public override string ToString()
@@ -110,4 +156,5 @@ protected override void WriteContents(BlobSerializationContext context)
                 NamedArguments[i].Write(context);
         }
     }
+
 }

src/AsmResolver.DotNet/Signatures/SerializedCustomAttributeSignature.cs

@@ -0,0 +1,73 @@
+using System.Linq;
+using System.Collections.Generic;
+using AsmResolver.DotNet.Signatures.Types;
+using AsmResolver.IO;
+
+namespace AsmResolver.DotNet.Signatures
+{
+    /// <summary>
+    /// Provides a lazy initialized implementation of the <see cref="CustomAttributeSignature"/> class.
+    /// </summary>
+    public class SerializedCustomAttributeSignature : CustomAttributeSignature
+    {
+        private readonly BlobReadContext _context;
+        private readonly TypeSignature[] _fixedArgTypes;
+        private readonly BinaryStreamReader _reader;
+
+        /// <summary>
+        /// Initializes a new lazy custom attribute signature from an input blob stream reader.
+        /// </summary>
+        /// <param name="context">The blob reading context the signature is situated in.</param>
+        /// <param name="fixedArgTypes">The types of all fixed arguments.</param>
+        /// <param name="reader">The input blob reader.</param>
+        public SerializedCustomAttributeSignature(
+            in BlobReadContext context,
+            IEnumerable<TypeSignature> fixedArgTypes,
+            in BinaryStreamReader reader)
+        {
+            _context = context;
+            _fixedArgTypes = fixedArgTypes.ToArray();
+            _reader = reader;
+        }
+
+        /// <inheritdoc />
+        protected override void Initialize(
+            IList<CustomAttributeArgument> fixedArguments,
+            IList<CustomAttributeNamedArgument> namedArguments)
+        {
+            var reader = _reader.Fork();
+
+            // Verify magic header.
+            ushort prologue = reader.ReadUInt16();
+            if (prologue != CustomAttributeSignaturePrologue)
+                _context.ReaderContext.BadImage("Input stream does not point to a valid custom attribute signature.");
+
+            // Read fixed arguments.
+            for (int i = 0; i < _fixedArgTypes.Length; i++)
+                fixedArguments.Add(CustomAttributeArgument.FromReader(_context, _fixedArgTypes[i], ref reader));
+
+            // Read named arguments.
+            ushort namedArgumentCount = reader.ReadUInt16();
+            for (int i = 0; i < namedArgumentCount; i++)
+                namedArguments.Add(CustomAttributeNamedArgument.FromReader(_context, ref reader));
+        }
+
+        /// <inheritdoc />
+        protected override void WriteContents(BlobSerializationContext context)
+        {
+            // If the arguments are not initialized yet, it means nobody accessed the fixed or named arguments of the
+            // signature. In such a case, we can safely assume nothing has changed to the signature.
+            //
+            // Since custom attribute signatures reference types by their fully qualified name, they do not contain
+            // any metadata tokens or type indices. Thus, regardless of whether the assembly changed or not, we can
+            // always just use the raw blob signature without risking breaking any references into the assembly.
+            // This can save a lot of processing time, given the fact that many custom attribute arguments are Enum
+            // typed arguments, which would require an expensive type resolution to determine the size of an element.
+
+            if (IsInitialized)
+                base.WriteContents(context);
+            else
+                _reader.Fork().WriteToOutput(context.Writer);
+        }
+    }
+}

test/AsmResolver.Benchmarks/ModuleReadWriteBenchmark.cs

@@ -1,6 +1,8 @@
 using System;
 using System.IO;
+using System.Linq;
 using AsmResolver.DotNet;
+using AsmResolver.IO;
 using BenchmarkDotNet.Attributes;
 using static AsmResolver.Benchmarks.Properties.Resources;
 
@@ -12,15 +14,22 @@ public class ModuleReadWriteBenchmark
         private static readonly byte[] HelloWorldApp = HelloWorld;
         private static readonly byte[] CrackMeApp = Test;
         private static readonly byte[] ManyMethods = Utilities.DecompressDeflate(HelloWorld_ManyMethods);
-        private static readonly byte[] CoreLib;
+        private static readonly IInputFile SystemPrivateCoreLib;
+        private static readonly IInputFile SystemRuntime;
+        private static readonly IInputFile SystemPrivateXml;
 
         private readonly MemoryStream _outputStream = new();
 
         static ModuleReadWriteBenchmark()
         {
-            var resolver = new DotNetCoreAssemblyResolver(new Version(3, 1, 0));
-            string path = resolver.Resolve(KnownCorLibs.SystemPrivateCoreLib_v4_0_0_0)!.ManifestModule!.FilePath;
-            CoreLib = File.ReadAllBytes(path);
+            string runtimePath = DotNetCorePathProvider.Default
+                .GetRuntimePathCandidates("Microsoft.NETCore.App", new Version(3, 1, 0))
+                .FirstOrDefault() ?? throw new InvalidOperationException(".NET Core 3.1 is not installed.");
+
+            var fs = new ByteArrayFileService();
+            SystemPrivateCoreLib = fs.OpenFile(Path.Combine(runtimePath, "System.Private.CoreLib.dll"));
+            SystemRuntime = fs.OpenFile(Path.Combine(runtimePath, "System.Runtime.dll"));
+            SystemPrivateXml = fs.OpenFile(Path.Combine(runtimePath, "System.Private.Xml.dll"));
         }
 
         [Benchmark]
@@ -63,9 +72,23 @@ public void ManyMethods_ReadWrite()
         }
 
         [Benchmark]
-        public void CoreLib_ReadWrite()
+        public void SystemPrivateCoreLib_ReadWrite()
+        {
+            var module = ModuleDefinition.FromFile(SystemPrivateCoreLib);
+            module.Write(_outputStream);
+        }
+
+        [Benchmark]
+        public void SystemRuntimeLib_ReadWrite()
+        {
+            var module = ModuleDefinition.FromFile(SystemRuntime);
+            module.Write(_outputStream);
+        }
+
+        [Benchmark]
+        public void SystemPrivateXml_ReadWrite()
         {
-            var module = ModuleDefinition.FromBytes(CoreLib);
+            var module = ModuleDefinition.FromFile(SystemPrivateXml);
             module.Write(_outputStream);
         }
     }