Implements dead branch removal for switch opcode (#2905)

Fixes https://github.com/dotnet/linker/issues/2888

2 files changed, 213 insertions, 37 deletions

diff --git a/src/linker/Linker.Steps/UnreachableBlocksOptimizer.cs b/src/linker/Linker.Steps/UnreachableBlocksOptimizer.cs
index 179f15d87..2a6e273f7 100644
--- a/src/linker/Linker.Steps/UnreachableBlocksOptimizer.cs
+++ b/src/linker/Linker.Steps/UnreachableBlocksOptimizer.cs
@@ -535,9 +535,15 @@ namespace Mono.Linker.Steps
 			//
 			// Sorted list of body instruction indexes which were
 			// replaced pass-through nop
-			// 
+			//
 			List<int>? conditionInstrsToRemove;
 
+			//
+			// Sorted list of body instruction indexes which were
+			// set to be replaced with different intstruction
+			//
+			List<(int, Instruction)>? conditionInstrsToReplace;
+
 			public BodyReducer (MethodBody body, LinkContext context)
 			{
 				Body = body;
@@ -546,6 +552,7 @@ namespace Mono.Linker.Steps
 				FoldedInstructions = null;
 				mapping = null;
 				conditionInstrsToRemove = null;
+				conditionInstrsToReplace = null;
 				InstructionsReplaced = 0;
 			}
 
@@ -582,6 +589,43 @@ namespace Mono.Linker.Steps
 				FoldedInstructions[index] = newInstruction;
 			}
 
+			void RewriteCondition (int index, Instruction instr, int operand)
+			{
+				switch (instr.OpCode.Code) {
+				case Code.Brfalse:
+				case Code.Brfalse_S:
+					if (operand == 0) {
+						Rewrite (index, Instruction.Create (OpCodes.Br, (Instruction) instr.Operand));
+					} else {
+						RewriteConditionToNop (index);
+					}
+
+					break;
+				case Code.Brtrue:
+				case Code.Brtrue_S:
+					if (operand != 0) {
+						Rewrite (index, Instruction.Create (OpCodes.Br, (Instruction) instr.Operand));
+					} else {
+						RewriteConditionToNop (index);
+					}
+
+					break;
+
+				case Code.Switch:
+					var targets = (Instruction[]) instr.Operand;
+					if (operand <= targets.Length) {
+						// It does not need to be conditional but existing logic in BodySweeper would
+						// need to be updated to deal with 1->2 instruction replacement
+						RewriteConditionTo (index, Instruction.Create (operand == 0 ? OpCodes.Brfalse : OpCodes.Brtrue, targets[operand]));
+						Rewrite (index, Instruction.Create (OpCodes.Br, targets[operand]));
+					} else {
+						RewriteConditionToNop (index);
+					}
+
+					break;
+				}
+			}
+
 			void RewriteConditionToNop (int index)
 			{
 				conditionInstrsToRemove ??= new List<int> ();
@@ -590,6 +634,13 @@ namespace Mono.Linker.Steps
 				RewriteToNop (index);
 			}
 
+			void RewriteConditionTo (int index, Instruction instruction)
+			{
+				conditionInstrsToReplace ??= new List<(int, Instruction)> ();
+
+				conditionInstrsToReplace.Add ((index, instruction));
+			}
+
 			public void RewriteToNop (int index, int stackDepth)
 			{
 				if (FoldedInstructions == null)
@@ -718,7 +769,7 @@ namespace Mono.Linker.Steps
 				var bodySweeper = new BodySweeper (Body, reachableInstrs, unreachableEH, context);
 				bodySweeper.Initialize ();
 
-				bodySweeper.Process (conditionInstrsToRemove, out var nopInstructions);
+				bodySweeper.Process (conditionInstrsToRemove, conditionInstrsToReplace, out var nopInstructions);
 				InstructionsReplaced = bodySweeper.InstructionsReplaced;
 				if (InstructionsReplaced == 0)
 					return false;
@@ -859,10 +910,6 @@ namespace Mono.Linker.Steps
 					var opcode = instr.OpCode;
 
 					if (opcode.FlowControl == FlowControl.Cond_Branch) {
-						// No support for removing branches from switch instruction
-						if (opcode == OpCodes.Switch)
-							continue;
-
 						if (opcode.StackBehaviourPop == StackBehaviour.Pop1_pop1) {
 							if (!GetOperandsConstantValues (i, out left, out right))
 								continue;
@@ -894,12 +941,7 @@ namespace Mono.Linker.Steps
 										continue;
 
 									RewriteToNop (i - 1);
-
-									if (IsConstantBranch (opcode, opint)) {
-										Rewrite (i, Instruction.Create (OpCodes.Br, (Instruction) instr.Operand));
-									} else {
-										RewriteConditionToNop (i);
-									}
+									RewriteCondition (i, instr, opint);
 
 									changed = true;
 									continue;
@@ -924,12 +966,27 @@ namespace Mono.Linker.Steps
 								RewriteToNop (i - 3);
 								RewriteToNop (i - 2);
 								RewriteToNop (i - 1);
+								RewriteCondition (i, instr, opint2);
 
-								if (IsConstantBranch (opcode, opint2)) {
-									Rewrite (i, Instruction.Create (OpCodes.Br, (Instruction) instr.Operand));
-								} else {
-									RewriteConditionToNop (i);
-								}
+								changed = true;
+								continue;
+							}
+
+							// Pattern for non-zero based switch with constant input
+							if (i >= 5 && opcode == OpCodes.Switch && GetConstantValue (FoldedInstructions[i - 5], out operand) && operand is int opint3 && IsPairedStlocLdloc (FoldedInstructions[i - 4], FoldedInstructions[i - 3])) {
+								if (IsJumpTargetRange (i - 4, i))
+									continue;
+
+								if (!GetConstantValue (FoldedInstructions[i - 2], out operand) || operand is not int offset)
+									continue;
+
+								if (FoldedInstructions[i - 1].OpCode != OpCodes.Sub)
+									continue;
+
+								RewriteToNop (i - 5);
+								RewriteToNop (i - 4);
+								RewriteToNop (i - 3);
+								RewriteCondition (i, instr, opint3 - offset);
 
 								changed = true;
 								continue;
@@ -1137,20 +1194,6 @@ namespace Mono.Linker.Steps
 				return false;
 			}
 
-			static bool IsConstantBranch (OpCode opCode, int operand)
-			{
-				switch (opCode.Code) {
-				case Code.Brfalse:
-				case Code.Brfalse_S:
-					return operand == 0;
-				case Code.Brtrue:
-				case Code.Brtrue_S:
-					return operand != 0;
-				}
-
-				throw new NotImplementedException (opCode.ToString ());
-			}
-
 			bool IsJumpTargetRange (int firstInstr, int lastInstr)
 			{
 				Debug.Assert (FoldedInstructions != null);
@@ -1219,15 +1262,34 @@ namespace Mono.Linker.Steps
 				ilprocessor = body.GetLinkerILProcessor ();
 			}
 
-			public void Process (List<int>? conditionInstrsToRemove, out List<Instruction>? sentinelNops)
+			public void Process (List<int>? conditionInstrsToRemove, List<(int, Instruction)>? conditionInstrsToReplace, out List<Instruction>? sentinelNops)
 			{
 				List<VariableDefinition>? removedVariablesReferences = null;
+				var instrs = Instructions;
+
+				//
+				// Process list of conditional instructions that were set to be replaced and not removed
+				//
+				if (conditionInstrsToReplace != null) {
+					foreach (var pair in conditionInstrsToReplace) {
+						var instr = instrs[pair.Item1];
+						switch (instr.OpCode.StackBehaviourPop) {
+						case StackBehaviour.Popi:
+							ILProcessor.Replace (pair.Item1, pair.Item2);
+							InstructionsReplaced++;
+							break;
+						default:
+							Debug.Fail ("not supported");
+							break;
+						}
+					}
+
+				}
 
 				//
 				// Initial pass which replaces unreachable instructions with nops or
 				// ret to keep the body verifiable
 				//
-				var instrs = Instructions;
 				for (int i = 0; i < instrs.Count; ++i) {
 					if (reachable[i])
 						continue;
diff --git a/test/Mono.Linker.Tests.Cases/UnreachableBlock/ComplexConditionsOptimized.cs b/test/Mono.Linker.Tests.Cases/UnreachableBlock/ComplexConditionsOptimized.cs
index a59e0cdc3..1161e7576 100644
--- a/test/Mono.Linker.Tests.Cases/UnreachableBlock/ComplexConditionsOptimized.cs
+++ b/test/Mono.Linker.Tests.Cases/UnreachableBlock/ComplexConditionsOptimized.cs
@@ -14,6 +14,9 @@ namespace Mono.Linker.Tests.Cases.UnreachableBlock
 		public static void Main ()
 		{
 			TestSwitch.Test ();
+			TestSwitch.TestOffset ();
+			TestSwitch2.TestFallThrough ();
+			TestSwitchZero.Test ();
 		}
 
 		[Kept]
@@ -24,7 +27,14 @@ namespace Mono.Linker.Tests.Cases.UnreachableBlock
 			}
 
 			[Kept]
-			[ExpectBodyModified]
+			[ExpectedInstructionSequence (new[] {
+				"ldc.i4.2",
+				"stloc.0",
+				"ldloc.0",
+				"brtrue il_8",
+				"call System.Void Mono.Linker.Tests.Cases.UnreachableBlock.ComplexConditionsOptimized/TestSwitch::Reached()",
+				"ret",
+			})]
 			public static void Test ()
 			{
 				switch (KnownInteger) {
@@ -37,13 +47,117 @@ namespace Mono.Linker.Tests.Cases.UnreachableBlock
 				case 2:
 					Reached ();
 					break;
-				default: throw new ApplicationException ();
+				default:
+					throw new ApplicationException ();
 				}
 			}
 
-			// https://github.com/dotnet/linker/issues/2888
-			// Should be removed
 			[Kept]
+			[ExpectedInstructionSequence (new[] {
+				"ldc.i4.2",
+				"stloc.0",
+				"ldloc.0",
+				"ldc.i4.1",
+				"sub",
+				"brtrue il_10",
+				"call System.Void Mono.Linker.Tests.Cases.UnreachableBlock.ComplexConditionsOptimized/TestSwitch::Reached()",
+				"ret",
+				"call System.Void Mono.Linker.Tests.Cases.UnreachableBlock.ComplexConditionsOptimized/TestSwitch::Reached()",
+				"br.s il_a",
+			})]
+			public static void TestOffset ()
+			{
+				switch (KnownInteger) {
+				case 1:
+					Reached ();
+					break;
+				case 2:
+					Reached ();
+					goto case 1;
+				case 3:
+					Unreached ();
+					break;
+				default:
+					throw new ApplicationException ();
+				}
+			}
+
+			static void Unreached () { }
+
+			[Kept]
+			static void Reached () { }
+		}
+
+		[Kept]
+		class TestSwitch2
+		{
+			static int KnownInteger {
+				get => 9;
+			}
+
+			[Kept]
+			[ExpectedInstructionSequence (new[] {
+				"ldc.i4.s 0x9",
+				"stloc.0",
+				"ldloc.0",
+				"pop",
+				"br.s il_7",
+				"newobj System.Void System.ApplicationException::.ctor()",
+				"throw",
+			})]
+			public static void TestFallThrough ()
+			{
+				switch (KnownInteger) {
+				case 0:
+					Unreached ();
+					break;
+				case 1:
+					Unreached ();
+					break;
+				case 2:
+					Unreached ();
+					break;
+				default:
+					throw new ApplicationException ();
+				}
+			}
+
+			static void Unreached () { }
+		}
+
+		[Kept]
+		class TestSwitchZero
+		{
+			static int KnownInteger {
+				get => 0;
+			}
+
+			[Kept]
+			[ExpectedInstructionSequence (new[] {
+				"ldc.i4.0",
+				"stloc.0",
+				"ldloc.0",
+				"brfalse il_8",
+				"call System.Void Mono.Linker.Tests.Cases.UnreachableBlock.ComplexConditionsOptimized/TestSwitchZero::Reached()",
+				"ret",
+			})]
+			public static void Test ()
+			{
+				switch (KnownInteger) {
+				case 0:
+					Reached ();
+					break;
+				case 1:
+					Unreached ();
+					break;
+				case 2:
+					Unreached ();
+					break;
+				default:
+					throw new ApplicationException ();
+				}
+			}
+
 			static void Unreached () { }
 
 			[Kept]