Fix quadratic algorithm in CompilerGeneratedState (#3150)

The way this code was supposed to work was that it would scan the compiler-
generated type and all its descendants, record each generated type it found,
then fill in information for all of the found types. The way it actually
worked was that it would scan the descendants, record each generated type, then
try to fill in information *for all generated types found in the program*. This
is quadratic as you start adding types, as you rescan everything you've added
before. The fix is to record just the types from the current pass, and then
add them to the larger bag when everything's complete.

Author: agocke

src/linker/Linker.Dataflow/CompilerGeneratedState.cs

@@ -126,6 +126,7 @@ public static bool TryGetStateMachineType (MethodDefinition method, [NotNullWhen
 
 			var callGraph = new CompilerGeneratedCallGraph ();
 			var userDefinedMethods = new HashSet<MethodDefinition> ();
+			var generatedTypeToTypeArgs = new Dictionary<TypeDefinition, TypeArgumentInfo> ();
 
 			void ProcessMethod (MethodDefinition method)
 			{
@@ -152,14 +153,15 @@ void ProcessMethod (MethodDefinition method)
 								if (referencedMethod == null)
 									continue;
 
+								// Find calls to state machine constructors that occur outside the type
 								if (referencedMethod.IsConstructor &&
 									referencedMethod.DeclaringType is var generatedType &&
 									// Don't consider calls in the same type, like inside a static constructor
 									method.DeclaringType != generatedType &&
 									CompilerGeneratedNames.IsLambdaDisplayClass (generatedType.Name)) {
 									// fill in null for now, attribute providers will be filled in later
-									if (!_generatedTypeToTypeArgumentInfo.TryAdd (generatedType, new TypeArgumentInfo (method, null))) {
-										var alreadyAssociatedMethod = _generatedTypeToTypeArgumentInfo[generatedType].CreatingMethod;
+									if (!generatedTypeToTypeArgs.TryAdd (generatedType, new TypeArgumentInfo (method, null))) {
+										var alreadyAssociatedMethod = generatedTypeToTypeArgs[generatedType].CreatingMethod;
 										_context.LogWarning (new MessageOrigin (method), DiagnosticId.MethodsAreAssociatedWithUserMethod, method.GetDisplayName (), alreadyAssociatedMethod.GetDisplayName (), generatedType.GetDisplayName ());
 									}
 									continue;
@@ -189,7 +191,7 @@ referencedMethod.DeclaringType is var generatedType &&
 									// Don't consider field accesses in the same type, like inside a static constructor
 									method.DeclaringType != generatedType &&
 									CompilerGeneratedNames.IsLambdaDisplayClass (generatedType.Name)) {
-									if (!_generatedTypeToTypeArgumentInfo.TryAdd (generatedType, new TypeArgumentInfo (method, null))) {
+									if (!generatedTypeToTypeArgs.TryAdd (generatedType, new TypeArgumentInfo (method, null))) {
 										// It's expected that there may be multiple methods associated with the same static closure environment.
 										// All of these methods will substitute the same type arguments into the closure environment
 										// (if it is generic). Don't warn.
@@ -214,7 +216,7 @@ referencedMethod.DeclaringType is var generatedType &&
 					}
 					// Already warned above if multiple methods map to the same type
 					// Fill in null for argument providers now, the real providers will be filled in later
-					_generatedTypeToTypeArgumentInfo[stateMachineType] = new TypeArgumentInfo (method, null);
+					generatedTypeToTypeArgs[stateMachineType] = new TypeArgumentInfo (method, null);
 				}
 			}
 
@@ -280,9 +282,17 @@ referencedMethod.DeclaringType is var generatedType &&
 
 			// Now that we have instantiating methods fully filled out, walk the generated types and fill in the attribute
 			// providers
-			foreach (var generatedType in _generatedTypeToTypeArgumentInfo.Keys) {
-				if (HasGenericParameters (generatedType))
-					MapGeneratedTypeTypeParameters (generatedType);
+			foreach (var generatedType in generatedTypeToTypeArgs.Keys) {
+				if (HasGenericParameters (generatedType)) {
+					MapGeneratedTypeTypeParameters (generatedType, generatedTypeToTypeArgs, _context);
+					// Finally, add resolved type arguments to the cache
+					var info = generatedTypeToTypeArgs[generatedType];
+					if (!_generatedTypeToTypeArgumentInfo.TryAdd (generatedType, info)) {
+						var method = info.CreatingMethod;
+						var alreadyAssociatedMethod = _generatedTypeToTypeArgumentInfo[generatedType].CreatingMethod;
+						_context.LogWarning (new MessageOrigin (method), DiagnosticId.MethodsAreAssociatedWithUserMethod, method.GetDisplayName (), alreadyAssociatedMethod.GetDisplayName (), generatedType.GetDisplayName ());
+					}
+				}
 			}
 
 			_cachedTypeToCompilerGeneratedMembers.Add (type, compilerGeneratedCallees);
@@ -301,18 +311,42 @@ static bool HasGenericParameters (TypeDefinition typeDef)
 				return typeDef.GenericParameters.Count > typeDef.DeclaringType.GenericParameters.Count;
 			}
 
-			void MapGeneratedTypeTypeParameters (TypeDefinition generatedType)
+			/// <summary>
+			/// Attempts to reverse the process of the compiler's alpha renaming. So if the original code was
+			/// something like this:
+			/// <code>
+			/// void M&lt;T&gt; () {
+			///     Action a = () => { Console.WriteLine (typeof (T)); };
+			/// }
+			/// </code>
+			/// The compiler will generate a nested class like this:
+			/// <code>
+			/// class &lt;&gt;c__DisplayClass0&lt;T&gt; {
+			///     public void &lt;M&gt;b__0 () {
+			///         Console.WriteLine (typeof (T));
+			///     }
+			/// }
+			/// </code>
+			/// The task of this method is to figure out that the type parameter T in the nested class is the same
+			/// as the type parameter T in the parent method M.
+			/// <paramref name="generatedTypeToTypeArgs"/> acts as a memoization table to avoid recalculating the
+			/// mapping multiple times.
+			/// </summary>
+			static void MapGeneratedTypeTypeParameters (
+				TypeDefinition generatedType,
+				Dictionary<TypeDefinition, TypeArgumentInfo> generatedTypeToTypeArgs,
+				LinkContext context)
 			{
 				Debug.Assert (CompilerGeneratedNames.IsGeneratedType (generatedType.Name));
 
-				var typeInfo = _generatedTypeToTypeArgumentInfo[generatedType];
+				var typeInfo = generatedTypeToTypeArgs[generatedType];
 				if (typeInfo.OriginalAttributes is not null) {
 					return;
 				}
 				var method = typeInfo.CreatingMethod;
 				if (method.Body is { } body) {
 					var typeArgs = new ICustomAttributeProvider[generatedType.GenericParameters.Count];
-					var typeRef = ScanForInit (generatedType, body);
+					var typeRef = ScanForInit (generatedType, body, context);
 					if (typeRef is null) {
 						return;
 					}
@@ -334,9 +368,9 @@ void MapGeneratedTypeTypeParameters (TypeDefinition generatedType)
 								var owningRef = (TypeReference) owner;
 								if (!CompilerGeneratedNames.IsGeneratedType (owningRef.Name)) {
 									userAttrs = param;
-								} else if (_context.TryResolve ((TypeReference) param.Owner) is { } owningType) {
-									MapGeneratedTypeTypeParameters (owningType);
-									if (_generatedTypeToTypeArgumentInfo[owningType].OriginalAttributes is { } owningAttrs) {
+								} else if (context.TryResolve ((TypeReference) param.Owner) is { } owningType) {
+									MapGeneratedTypeTypeParameters (owningType, generatedTypeToTypeArgs, context);
+									if (generatedTypeToTypeArgs[owningType].OriginalAttributes is { } owningAttrs) {
 										userAttrs = owningAttrs[param.Position];
 									} else {
 										Debug.Assert (false, "This should be impossible in valid code");
@@ -348,27 +382,30 @@ void MapGeneratedTypeTypeParameters (TypeDefinition generatedType)
 						typeArgs[i] = userAttrs;
 					}
 
-					_generatedTypeToTypeArgumentInfo[generatedType] = typeInfo with { OriginalAttributes = typeArgs };
+					generatedTypeToTypeArgs[generatedType] = typeInfo with { OriginalAttributes = typeArgs };
 				}
 			}
 
-			GenericInstanceType? ScanForInit (TypeDefinition compilerGeneratedType, MethodBody body)
+			static GenericInstanceType? ScanForInit (
+				TypeDefinition compilerGeneratedType,
+				MethodBody body,
+				LinkContext context)
 			{
-				foreach (var instr in _context.GetMethodIL (body).Instructions) {
+				foreach (var instr in context.GetMethodIL (body).Instructions) {
 					bool handled = false;
 					switch (instr.OpCode.Code) {
 					case Code.Initobj:
 					case Code.Newobj: {
 							if (instr.Operand is MethodReference { DeclaringType: GenericInstanceType typeRef }
-								&& compilerGeneratedType == _context.TryResolve (typeRef)) {
+								&& compilerGeneratedType == context.TryResolve (typeRef)) {
 								return typeRef;
 							}
 							handled = true;
 						}
 						break;
 					case Code.Stsfld: {
 							if (instr.Operand is FieldReference { DeclaringType: GenericInstanceType typeRef }
-								&& compilerGeneratedType == _context.TryResolve (typeRef)) {
+								&& compilerGeneratedType == context.TryResolve (typeRef)) {
 								return typeRef;
 							}
 							handled = true;
@@ -381,7 +418,7 @@ void MapGeneratedTypeTypeParameters (TypeDefinition generatedType)
 					if (!handled && instr.OpCode.OperandType is OperandType.InlineMethod) {
 						if (instr.Operand is GenericInstanceMethod gim) {
 							foreach (var tr in gim.GenericArguments) {
-								if (tr is GenericInstanceType git && compilerGeneratedType == _context.TryResolve (git)) {
+								if (tr is GenericInstanceType git && compilerGeneratedType == context.TryResolve (git)) {
 									return git;
 								}
 							}

test/Mono.Linker.Tests/Tests/PerfTestGeneratorForCompilerGeneratedCode.cs

@@ -0,0 +1,47 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System.IO;
+using System.Linq;
+
+/// <summary>
+/// This class generates a test that can be used to test perf of analyzing
+/// compiler-generated code. Run it by copying this file into a console app and
+/// calling <see cref="PerfTestGeneratorForCompilerGeneratedCode.Run"/>. A file
+/// will be generated in the current directory named GeneratedLinkerTests.cs.
+/// Copy this file into another Console app and trim the app to measure the
+/// perf.
+/// </summary>
+static class PerfTestGeneratorForCompilerGeneratedCode
+{
+	const int FuncNumber = 10000;
+	public static void Run ()
+	{
+		using var fstream = File.Create ("GeneratedLinkerTests.cs");
+		using var writer = new StreamWriter (fstream);
+		writer.WriteLine ($$"""
+class C {
+    public static async void Main()
+    {
+        int x = 0;
+        {{string.Join (@"
+        ", Enumerable.Range (0, FuncNumber).Select (i => $"x += await N{i}<int>.M();"))}}
+        Console.WriteLine(x);
+    }
+}
+""");
+		for (int i = 0; i < FuncNumber; i++) {
+			writer.WriteLine ($$"""
+public static class N{{i}}<T>
+{
+    public static async ValueTask<int> M()
+    {
+        Func<int> a = () => 1;
+        await Task.Delay(0);
+        return a();
+    }
+}
+""");
+		}
+	}
+}