Only compute if method returns a constant when needed (#1734)

This changes how the `RemoveUnreachableBlocksStep` computes if a method returns a constant or not. Previously we would go over all methods and detect those which return constants (and store only those in a dictionary). Then we would proceed with branch removal relying on this dictionary.

For a hello world console app this means we ran the detection of constants on almost 50K methods. On the other hand the dictionary stored only about 1.5K records (actually constant methods).

With this change the constant detection only runs on methods we need the result for later on. This means the dictionary stores records for all methods asked about (`null` value represents "We checked, and it's not constant").

For the same hello world console app we now run the detection only on ~5K methods (10x less). On the other hand we have records for all (5K -> 3x increase). So this trades CPU for memory.

The purpose of this change is not the CPU/memory tradeoff, it's to prepare the constant propagation to make it possible to call it on-demand per-method from `MarkStep` which will be needed once we switch over to the current lazy-loading of assemblies.

4 files changed, 87 insertions, 62 deletions

diff --git a/src/linker/Linker.Steps/RemoveUnreachableBlocksStep.cs b/src/linker/Linker.Steps/RemoveUnreachableBlocksStep.cs
index 12b249ea8..b9cd165c0 100644
--- a/src/linker/Linker.Steps/RemoveUnreachableBlocksStep.cs
+++ b/src/linker/Linker.Steps/RemoveUnreachableBlocksStep.cs
@@ -17,6 +17,7 @@ namespace Mono.Linker.Steps
 	public class RemoveUnreachableBlocksStep : BaseStep
 	{
 		Dictionary<MethodDefinition, Instruction> constExprMethods;
+		bool constExprMethodsAdded;
 		MethodDefinition IntPtrSize, UIntPtrSize;
 
 		protected override void Process ()
@@ -24,19 +25,12 @@ namespace Mono.Linker.Steps
 			var assemblies = Context.Annotations.GetAssemblies ().ToArray ();
 
 			constExprMethods = new Dictionary<MethodDefinition, Instruction> ();
-			foreach (var assembly in assemblies) {
-				FindConstantExpressionsMethods (assembly.MainModule.Types);
-			}
-
-			if (constExprMethods.Count == 0)
-				return;
 
-			int constExprMethodsCount;
 			do {
 				//
 				// Body rewriting can produce more methods with constant expression
 				//
-				constExprMethodsCount = constExprMethods.Count;
+				constExprMethodsAdded = false;
 
 				foreach (var assembly in assemblies) {
 					if (Annotations.GetAction (assembly) != AssemblyAction.Link)
@@ -44,54 +38,40 @@ namespace Mono.Linker.Steps
 
 					RewriteBodies (assembly.MainModule.Types);
 				}
-			} while (constExprMethodsCount < constExprMethods.Count);
+			} while (constExprMethodsAdded);
 		}
 
-		void FindConstantExpressionsMethods (Collection<TypeDefinition> types)
+		bool TryGetConstantResultInstructionForMethod (MethodDefinition method, out Instruction constantResultInstruction)
 		{
-			foreach (var type in types) {
-				if (type.IsInterface)
-					continue;
-
-				if (!type.HasMethods)
-					continue;
+			if (constExprMethods.TryGetValue (method, out constantResultInstruction))
+				return constantResultInstruction != null;
 
-				foreach (var method in type.Methods) {
-					if (!method.HasBody)
-						continue;
-
-					if (method.ReturnType.MetadataType == MetadataType.Void)
-						continue;
+			constantResultInstruction = GetConstantResultInstructionForMethod (method);
+			constExprMethods.Add (method, constantResultInstruction);
 
-					switch (Annotations.GetAction (method)) {
-					case MethodAction.ConvertToThrow:
-						continue;
-					case MethodAction.ConvertToStub:
-						var instruction = CodeRewriterStep.CreateConstantResultInstruction (Context, method);
-						if (instruction != null)
-							constExprMethods[method] = instruction;
-
-						continue;
-					}
+			return constantResultInstruction != null;
+		}
 
-					if (method.IsIntrinsic () || method.NoInlining)
-						continue;
+		Instruction GetConstantResultInstructionForMethod (MethodDefinition method)
+		{
+			if (!method.HasBody)
+				return null;
 
-					if (constExprMethods.ContainsKey (method))
-						continue;
+			if (method.ReturnType.MetadataType == MetadataType.Void)
+				return null;
 
-					if (!Context.IsOptimizationEnabled (CodeOptimizations.IPConstantPropagation, method))
-						continue;
+			if (method.IsIntrinsic () || method.NoInlining)
+				return null;
 
-					var analyzer = new ConstantExpressionMethodAnalyzer (method);
-					if (analyzer.Analyze ()) {
-						constExprMethods[method] = analyzer.Result;
-					}
-				}
+			if (!Context.IsOptimizationEnabled (CodeOptimizations.IPConstantPropagation, method))
+				return null;
 
-				if (type.HasNestedTypes)
-					FindConstantExpressionsMethods (type.NestedTypes);
+			var analyzer = new ConstantExpressionMethodAnalyzer (Context, method);
+			if (analyzer.Analyze ()) {
+				return analyzer.Result;
 			}
+
+			return null;
 		}
 
 		void RewriteBodies (Collection<TypeDefinition> types)
@@ -115,7 +95,6 @@ namespace Mono.Linker.Steps
 					case MetadataType.FunctionPointer:
 						continue;
 					}
-
 					RewriteBody (method);
 				}
 
@@ -148,12 +127,15 @@ namespace Mono.Linker.Steps
 			if (method.ReturnType.MetadataType == MetadataType.Void)
 				return;
 
-			//
-			// Re-run the analyzer in case body change rewrote it to constant expression
-			//
-			var analyzer = new ConstantExpressionMethodAnalyzer (method, reducer.FoldedInstructions);
-			if (analyzer.Analyze ()) {
-				constExprMethods[method] = analyzer.Result;
+			if (!constExprMethods.TryGetValue (method, out var constInstruction) || constInstruction == null) {
+				//
+				// Re-run the analyzer in case body change rewrote it to constant expression
+				//
+				var analyzer = new ConstantExpressionMethodAnalyzer (Context, method, reducer.FoldedInstructions);
+				if (analyzer.Analyze ()) {
+					constExprMethods[method] = analyzer.Result;
+					constExprMethodsAdded = true;
+				}
 			}
 		}
 
@@ -174,9 +156,6 @@ namespace Mono.Linker.Steps
 					if (md == null)
 						break;
 
-					if (!constExprMethods.TryGetValue (md, out targetResult))
-						break;
-
 					if (md.CallingConvention == MethodCallingConvention.VarArg)
 						break;
 
@@ -208,8 +187,12 @@ namespace Mono.Linker.Steps
 							break;
 					}
 
+					if (!TryGetConstantResultInstructionForMethod (md, out targetResult))
+						break;
+
 					reducer.Rewrite (i, targetResult);
 					changed = true;
+
 					break;
 
 				case Code.Ldsfld:
@@ -244,7 +227,7 @@ namespace Mono.Linker.Steps
 						sizeOfImpl = (IntPtrSize ??= FindSizeMethod (operand.Resolve ()));
 					}
 
-					if (sizeOfImpl != null && constExprMethods.TryGetValue (sizeOfImpl, out targetResult)) {
+					if (sizeOfImpl != null && TryGetConstantResultInstructionForMethod (sizeOfImpl, out targetResult)) {
 						reducer.Rewrite (i, targetResult);
 						changed = true;
 					}
@@ -1042,14 +1025,16 @@ namespace Mono.Linker.Steps
 
 		struct ConstantExpressionMethodAnalyzer
 		{
+			readonly LinkContext context;
 			readonly MethodDefinition method;
 			readonly Collection<Instruction> instructions;
 
 			Stack<Instruction> stack_instr;
 			Dictionary<int, Instruction> locals;
 
-			public ConstantExpressionMethodAnalyzer (MethodDefinition method)
+			public ConstantExpressionMethodAnalyzer (LinkContext context, MethodDefinition method)
 			{
+				this.context = context;
 				this.method = method;
 				instructions = method.Body.Instructions;
 				stack_instr = null;
@@ -1057,19 +1042,24 @@ namespace Mono.Linker.Steps
 				Result = null;
 			}
 
-			public ConstantExpressionMethodAnalyzer (MethodDefinition method, Collection<Instruction> instructions)
+			public ConstantExpressionMethodAnalyzer (LinkContext context, MethodDefinition method, Collection<Instruction> instructions)
+				: this (context, method)
 			{
-				this.method = method;
 				this.instructions = instructions;
-				stack_instr = null;
-				locals = null;
-				Result = null;
 			}
 
 			public Instruction Result { get; private set; }
 
 			public bool Analyze ()
 			{
+				switch (context.Annotations.GetAction (method)) {
+				case MethodAction.ConvertToThrow:
+					return false;
+				case MethodAction.ConvertToStub:
+					Result = CodeRewriterStep.CreateConstantResultInstruction (context, method);
+					return Result != null;
+				}
+
 				var body = method.Body;
 				if (body.HasExceptionHandlers)
 					return false;
diff --git a/test/Mono.Linker.Tests.Cases/FeatureSettings/FeatureSubstitutions.cs b/test/Mono.Linker.Tests.Cases/FeatureSettings/FeatureSubstitutions.cs
index c531f2110..bb93e90c0 100644
--- a/test/Mono.Linker.Tests.Cases/FeatureSettings/FeatureSubstitutions.cs
+++ b/test/Mono.Linker.Tests.Cases/FeatureSettings/FeatureSubstitutions.cs
@@ -6,6 +6,7 @@ namespace Mono.Linker.Tests.Cases.FeatureSettings
 {
 	[SetupLinkerSubstitutionFile ("FeatureSubstitutions.xml")]
 	[SetupLinkerArgument ("--feature", "OptionalFeature", "false")]
+	[SetupLinkerArgument ("--enable-opt", "ipconstprop")]
 	public class FeatureSubstitutions
 	{
 		[Kept]
diff --git a/test/Mono.Linker.Tests.Cases/UnreachableBlock/BodiesWithSubstitutions.cs b/test/Mono.Linker.Tests.Cases/UnreachableBlock/BodiesWithSubstitutions.cs
index 3b8533deb..d06060dbd 100644
--- a/test/Mono.Linker.Tests.Cases/UnreachableBlock/BodiesWithSubstitutions.cs
+++ b/test/Mono.Linker.Tests.Cases/UnreachableBlock/BodiesWithSubstitutions.cs
@@ -40,6 +40,7 @@ namespace Mono.Linker.Tests.Cases.UnreachableBlock
 
 			ConstantFromNewAssembly.Test ();
 			ConstantSubstitutionsFromNewAssembly.Test ();
+			TestSubstitutionCollision ();
 		}
 
 		[Kept]
@@ -246,5 +247,36 @@ namespace Mono.Linker.Tests.Cases.UnreachableBlock
 					Reached ();
 			}
 		}
+
+		[Kept]
+		static bool CollisionProperty {
+			[Kept]
+			[ExpectBodyModified]
+			get {
+				// Need to call something with constant value to force processing of this method
+				_ = Property;
+				return true;
+			} // Substitution will set this to false
+		}
+
+		// This tests that if there's a method (get_CollisionProperty) which itself is constant
+		// and substitution changes its return value, the branch removal reacts to the substituted value
+		// and not the value from the method's body.
+		// This should ideally never happen, but still.
+		// In the original code this test would be order dependent. Depending if TestSubstitutionsCollision
+		// was processed before CollisionProperty, it would either propagate true or false.
+		[Kept]
+		[ExpectBodyModified]
+		static void TestSubstitutionCollision ()
+		{
+			if (CollisionProperty)
+				Collision_NeverReached ();
+			else
+				Collision_Reached ();
+		}
+
+		[Kept]
+		static void Collision_Reached () { }
+		static void Collision_NeverReached () { }
 	}
 }
diff --git a/test/Mono.Linker.Tests.Cases/UnreachableBlock/BodiesWithSubstitutions.xml b/test/Mono.Linker.Tests.Cases/UnreachableBlock/BodiesWithSubstitutions.xml
index 23fb0d5e9..cad903687 100644
--- a/test/Mono.Linker.Tests.Cases/UnreachableBlock/BodiesWithSubstitutions.xml
+++ b/test/Mono.Linker.Tests.Cases/UnreachableBlock/BodiesWithSubstitutions.xml
@@ -3,6 +3,8 @@
     <type fullname="Mono.Linker.Tests.Cases.UnreachableBlock.BodiesWithSubstitutions">
       <method signature="System.Int32 get_Property()" body="stub" value="3">
       </method>
+      <method signature="System.Boolean get_CollisionProperty()" body="stub" value="false">
+      </method>
     </type>
     <type fullname="Mono.Linker.Tests.Cases.UnreachableBlock.BodiesWithSubstitutions/ClassWithField">
       <field name="SField" value="9"/>