TypeProvider Optimizations (#322)

* Optimalizing TypeProvider.cs

* Update TypeProvider.cs

* bug fix

* PR comments

* Update TypeProvider.cs

* Update TypeProvider.cs

* Update TypeProvider.cs

src/Draco.Compiler/Api/Compilation.cs

@@ -133,6 +133,11 @@ public static Compilation Create(
     /// </summary>
     internal WellKnownTypes WellKnownTypes { get; }
 
+    /// <summary>
+    /// The type provider used for metadata references.
+    /// </summary>
+    internal TypeProvider TypeProvider { get; }
+
     /// <summary>
     /// Intrinsicly defined symbols for the compilation.
     /// </summary>
@@ -153,6 +158,7 @@ private Compilation(
         ModuleSymbol? sourceModule = null,
         DeclarationTable? declarationTable = null,
         WellKnownTypes? wellKnownTypes = null,
+        TypeProvider? typeProvider = null,
         IntrinsicSymbols? intrinsicSymbols = null,
         BinderCache? binderCache = null)
     {
@@ -166,6 +172,7 @@ private Compilation(
         this.sourceModule = sourceModule;
         this.declarationTable = declarationTable;
         this.WellKnownTypes = wellKnownTypes ?? new WellKnownTypes(this);
+        this.TypeProvider = typeProvider ?? new TypeProvider(this);
         this.IntrinsicSymbols = intrinsicSymbols ?? new IntrinsicSymbols(this);
         this.binderCache = binderCache ?? new BinderCache(this);
     }
@@ -215,6 +222,10 @@ public Compilation UpdateSyntaxTree(SyntaxTree? oldTree, SyntaxTree? newTree)
             // Or we keep it as long as metadata refs don't change?
             // Just a cache
             wellKnownTypes: this.WellKnownTypes,
+            // TODO: We might want to change the compilation of type provider?
+            // Or we keep it as long as metadata refs don't change?
+            // Just a cache
+            typeProvider: this.TypeProvider,
             // TODO: We might want to change the compilation of intrinsic-symbols?
             // Or we keep it as long as metadata refs don't change?
             // Just a cache

src/Draco.Compiler/Internal/Symbols/Metadata/MetadataFieldSymbol.cs

@@ -77,12 +77,8 @@ public MetadataFieldSymbol(Symbol containingSymbol, FieldDefinition fieldDefinit
         this.fieldDefinition = fieldDefinition;
     }
 
-    private TypeSymbol BuildType()
-    {
-        // Decode signature
-        var decoder = new TypeProvider(this.Assembly.Compilation);
-        return this.fieldDefinition.DecodeSignature(decoder, this);
-    }
+    private TypeSymbol BuildType() =>
+        this.fieldDefinition.DecodeSignature(this.Assembly.Compilation.TypeProvider, this);
 
     private object? BuildDefaultValue()
     {

src/Draco.Compiler/Internal/Symbols/Metadata/MetadataMethodSymbol.cs

@@ -139,8 +139,7 @@ private ImmutableArray<TypeParameterSymbol> BuildGenericParameters()
     private void BuildSignature()
     {
         // Decode signature
-        var decoder = new TypeProvider(this.Assembly.Compilation);
-        var signature = this.methodDefinition.DecodeSignature(decoder, this);
+        var signature = this.methodDefinition.DecodeSignature(this.Assembly.Compilation.TypeProvider, this);
 
         // Build parameters
         var parameters = ImmutableArray.CreateBuilder<ParameterSymbol>();
@@ -194,7 +193,7 @@ private void BuildOverride()
     {
         var definition = this.MetadataReader.GetMethodDefinition(methodDef);
         var name = this.MetadataReader.GetString(definition.Name);
-        var provider = new TypeProvider(this.Assembly.Compilation);
+        var provider = this.Assembly.Compilation.TypeProvider;
         var signature = definition.DecodeSignature(provider, this);
         var containingType = provider.GetTypeFromDefinition(this.MetadataReader, definition.GetDeclaringType(), 0);
         return GetFunctionWithSignature(containingType, name, signature);
@@ -204,7 +203,7 @@ private void BuildOverride()
     {
         var reference = this.MetadataReader.GetMemberReference(methodRef);
         var name = this.MetadataReader.GetString(reference.Name);
-        var provider = new TypeProvider(this.Assembly.Compilation);
+        var provider = this.Assembly.Compilation.TypeProvider;
         var signature = reference.DecodeMethodSignature(provider, this);
         var containingType = provider.GetTypeFromReference(this.MetadataReader, (TypeReferenceHandle)reference.Parent, 0);
         return GetFunctionWithSignature(containingType, name, signature);

src/Draco.Compiler/Internal/Symbols/Metadata/MetadataSymbol.cs

@@ -58,7 +58,7 @@ public static IEnumerable<Symbol> ToSymbol(
         foreach (var attributeHandle in typeDefinition.GetCustomAttributes())
         {
             var attribute = reader.GetCustomAttribute(attributeHandle);
-            var typeProvider = new TypeProvider(compilation!);
+            var typeProvider = compilation.TypeProvider;
             switch (attribute.Constructor.Kind)
             {
             case HandleKind.MethodDefinition:

src/Draco.Compiler/Internal/Symbols/Metadata/MetadataTypeSymbol.cs

@@ -95,7 +95,7 @@ private ImmutableArray<TypeParameterSymbol> BuildGenericParameters()
     private ImmutableArray<TypeSymbol> BuildBaseTypes()
     {
         var builder = ImmutableArray.CreateBuilder<TypeSymbol>();
-        var typeProvider = new TypeProvider(this.Assembly.Compilation);
+        var typeProvider = this.Assembly.Compilation.TypeProvider;
         if (!this.typeDefinition.BaseType.IsNil)
         {
             builder.Add(GetTypeFromMetadata(this.typeDefinition.BaseType));

src/Draco.Compiler/Internal/Symbols/Metadata/TypeProvider.cs

@@ -1,7 +1,9 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Linq;
+using System.Reflection;
 using System.Reflection.Metadata;
+using System.Reflection.PortableExecutable;
 using Draco.Compiler.Api;
 using Draco.Compiler.Internal.Symbols.Synthetized;
 
@@ -12,13 +14,16 @@ namespace Draco.Compiler.Internal.Symbols.Metadata;
 /// </summary>
 internal sealed class TypeProvider : ISignatureTypeProvider<TypeSymbol, Symbol>, ICustomAttributeTypeProvider<TypeSymbol>
 {
+    private readonly record struct CacheKey(MetadataReader Reader, EntityHandle Handle);
+
     // TODO: We return a special error type for now to swallow errors
     private static TypeSymbol UnknownType { get; } = new PrimitiveTypeSymbol("<unknown>", false);
 
     private WellKnownTypes WellKnownTypes => this.compilation.WellKnownTypes;
     private IntrinsicSymbols IntrinsicSymbols => this.compilation.IntrinsicSymbols;
 
     private readonly Compilation compilation;
+    private readonly Dictionary<CacheKey, TypeSymbol> cache = new();
 
     public TypeProvider(Compilation compilation)
     {
@@ -36,6 +41,7 @@ public TypeSymbol GetGenericInstantiation(TypeSymbol genericType, ImmutableArray
         if (ReferenceEquals(genericType, UnknownType)) return UnknownType;
         return genericType.GenericInstantiate(genericType.ContainingSymbol, typeArguments);
     }
+
     public TypeSymbol GetGenericMethodParameter(Symbol genericContext, int index)
     {
         var methodAncestor = genericContext.AncestorChain
@@ -46,6 +52,7 @@ public TypeSymbol GetGenericMethodParameter(Symbol genericContext, int index)
             ? methodAncestor.GenericParameters[index]
             : methodAncestor.GenericDefinition!.GenericParameters[index];
     }
+
     public TypeSymbol GetGenericTypeParameter(Symbol genericContext, int index)
     {
         var typeAncestor = genericContext.AncestorChain
@@ -56,6 +63,7 @@ public TypeSymbol GetGenericTypeParameter(Symbol genericContext, int index)
             ? typeAncestor.GenericParameters[index]
             : typeAncestor.GenericDefinition!.GenericParameters[index];
     }
+
     public TypeSymbol GetModifiedType(TypeSymbol modifier, TypeSymbol unmodifiedType, bool isRequired) => UnknownType;
     public TypeSymbol GetPinnedType(TypeSymbol elementType) => UnknownType;
     public TypeSymbol GetPointerType(TypeSymbol elementType) => UnknownType;
@@ -84,7 +92,19 @@ public TypeSymbol GetGenericTypeParameter(Symbol genericContext, int index)
 
         _ => UnknownType,
     };
+
     public TypeSymbol GetTypeFromDefinition(MetadataReader reader, TypeDefinitionHandle handle, byte rawTypeKind)
+    {
+        var key = new CacheKey(reader, handle);
+        if (!this.cache.TryGetValue(key, out var type))
+        {
+            type = this.BuildTypeFromDefinition(reader, handle, rawTypeKind);
+            this.cache.Add(key, type);
+        }
+        return type;
+    }
+
+    private TypeSymbol BuildTypeFromDefinition(MetadataReader reader, TypeDefinitionHandle handle, byte rawTypeKind)
     {
         var definition = reader.GetTypeDefinition(handle);
         if (definition.IsNested)
@@ -96,12 +116,11 @@ public TypeSymbol GetTypeFromDefinition(MetadataReader reader, TypeDefinitionHan
             // Search for this type by name and generic argument count
             var nestedName = reader.GetString(definition.Name);
             var nestedGenericArgc = definition.GetGenericParameters().Count;
-            var nestedSymbol = declaringSymbol
+            return declaringSymbol
                 .DefinedMembers
                 .OfType<TypeSymbol>()
                 .Where(t => t.Name == nestedName && t.GenericParameters.Length == nestedGenericArgc)
                 .Single();
-            return nestedSymbol;
         }
 
         var assemblyName = reader
@@ -116,7 +135,19 @@ public TypeSymbol GetTypeFromDefinition(MetadataReader reader, TypeDefinitionHan
 
         return this.WellKnownTypes.GetTypeFromAssembly(assemblyName, path);
     }
+
     public TypeSymbol GetTypeFromReference(MetadataReader reader, TypeReferenceHandle handle, byte rawTypeKind)
+    {
+        var key = new CacheKey(reader, handle);
+        if (!this.cache.TryGetValue(key, out var type))
+        {
+            type = this.BuildTypeFromReference(reader, handle, rawTypeKind);
+            this.cache.Add(key, type);
+        }
+        return type;
+    }
+
+    private TypeSymbol BuildTypeFromReference(MetadataReader reader, TypeReferenceHandle handle, byte rawTypeKind)
     {
         var parts = new List<string>();
         var reference = reader.GetTypeReference(handle);
@@ -136,6 +167,8 @@ public TypeSymbol GetTypeFromReference(MetadataReader reader, TypeReferenceHandl
         var assembly = this.compilation.MetadataAssemblies.Values.Single(x => x.AssemblyName.FullName == assemblyName.FullName);
         return assembly.RootNamespace.Lookup(parts.ToImmutableArray()).OfType<TypeSymbol>().Single();
     }
+
+    // TODO: Should we cache this as well? doesn't seem to have any effect
     public TypeSymbol GetTypeFromSpecification(MetadataReader reader, Symbol genericContext, TypeSpecificationHandle handle, byte rawTypeKind)
     {
         var specification = reader.GetTypeSpecification(handle);