Refactor TransformDisplayClassUsage

ICSharpCode.Decompiler.Tests/TestCases/Pretty/LocalFunctions.cs

@@ -310,6 +310,17 @@ namespace LocalFunctions
 			}
 		}
 
+		public void WriteCapturedParameter(int i)
+		{
+			ParamWrite();
+			Console.WriteLine(i);
+
+			void ParamWrite()
+			{
+				i++;
+			}
+		}
+
 		//public static void LocalFunctionInUsing()
 		//{
 		//	using (MemoryStream memoryStream = new MemoryStream()) {

ICSharpCode.Decompiler/IL/Transforms/TransformDisplayClassUsage.cs

@@ -37,26 +37,19 @@ namespace ICSharpCode.Decompiler.IL.Transforms
 			public ILInstruction Initializer;
 			public ILVariable Variable;
 			public ITypeDefinition Definition;
-			public Dictionary<IField, DisplayClassVariable> Variables;
+			public Dictionary<IField, ILVariable> Variables;
 			public BlockContainer CaptureScope;
 			public ILFunction DeclaringFunction;
 		}
 
-		struct DisplayClassVariable
-		{
-			public ILVariable Variable;
-			public ILInstruction Value;
-		}
-
 		ILTransformContext context;
-		ILFunction currentFunction;
 		readonly Dictionary<ILVariable, DisplayClass> displayClasses = new Dictionary<ILVariable, DisplayClass>();
 		readonly List<ILInstruction> instructionsToRemove = new List<ILInstruction>();
 
 		public void Run(ILFunction function, ILTransformContext context)
 		{
 			try {
-				if (this.context != null || this.currentFunction != null)
+				if (this.context != null)
 					throw new InvalidOperationException("Reentrancy in " + nameof(TransformDisplayClassUsage));
 				this.context = context;
 				var decompilationContext = new SimpleTypeResolveContext(context.Function.Method);
@@ -73,12 +66,8 @@ namespace ICSharpCode.Decompiler.IL.Transforms
 							AddOrUpdateDisplayClass(f, v, ((ByReferenceType)p.Type).ElementType.GetDefinition(), f.Body, localFunctionClosureParameter: true);
 						}
 					}
-					foreach (var displayClass in displayClasses.Values.OrderByDescending(d => d.Initializer.StartILOffset).ToArray()) {
-						context.Step($"Transform references to " + displayClass.Variable, displayClass.Initializer);
-						this.currentFunction = f;
-						VisitILFunction(f);
-					}
 				}
+				VisitILFunction(function);
 				if (instructionsToRemove.Count > 0) {
 					context.Step($"Remove instructions", function);
 					foreach (var store in instructionsToRemove) {
@@ -91,7 +80,6 @@ namespace ICSharpCode.Decompiler.IL.Transforms
 				instructionsToRemove.Clear();
 				displayClasses.Clear();
 				this.context = null;
-				this.currentFunction = null;
 			}
 		}
 
@@ -106,7 +94,7 @@ namespace ICSharpCode.Decompiler.IL.Transforms
 					Initializer = inst,
 					Variable = v,
 					Definition = closureType,
-					Variables = new Dictionary<IField, DisplayClassVariable>(),
+					Variables = new Dictionary<IField, ILVariable>(),
 					CaptureScope = (isMono && IsMonoNestedCaptureScope(closureType)) || localFunctionClosureParameter ? null : v.CaptureScope,
 					DeclaringFunction = localFunctionClosureParameter ? f.DeclarationScope.Ancestors.OfType<ILFunction>().First() : f
 				});
@@ -151,11 +139,6 @@ namespace ICSharpCode.Decompiler.IL.Transforms
 			return false;
 		}
 
-		bool IsOuterClosureReference(IField field)
-		{
-			return displayClasses.Values.Any(disp => disp.Definition == field.DeclaringTypeDefinition);
-		}
-
 		bool IsMonoNestedCaptureScope(ITypeDefinition closureType)
 		{
 			var decompilationContext = new SimpleTypeResolveContext(context.Function.Method);
@@ -181,17 +164,14 @@ namespace ICSharpCode.Decompiler.IL.Transforms
 				Initializer = nestedFunction.Body,
 				Variable = thisVariable,
 				Definition = thisVariable.Type.GetDefinition(),
-				Variables = new Dictionary<IField, DisplayClassVariable>(),
+				Variables = new Dictionary<IField, ILVariable>(),
 				CaptureScope = (BlockContainer)nestedFunction.Body
 			};
 			displayClasses.Add(thisVariable, displayClass);
 			foreach (var stateMachineVariable in nestedFunction.Variables) {
 				if (stateMachineVariable.StateMachineField == null || displayClass.Variables.ContainsKey(stateMachineVariable.StateMachineField))
 					continue;
-				displayClass.Variables.Add(stateMachineVariable.StateMachineField, new DisplayClassVariable {
-					Variable = stateMachineVariable,
-					Value = new LdLoc(stateMachineVariable)
-				});
+				displayClass.Variables.Add(stateMachineVariable.StateMachineField, stateMachineVariable);
 			}
 			if (!currentFunction.Method.IsStatic && FindThisField(out var thisField)) {
 				var thisVar = currentFunction.Variables
@@ -200,7 +180,7 @@ namespace ICSharpCode.Decompiler.IL.Transforms
 					thisVar = new ILVariable(VariableKind.Parameter, decompilationContext.CurrentTypeDefinition, -1) { Name = "this" };
 					currentFunction.Variables.Add(thisVar);
 				}
-				displayClass.Variables.Add(thisField, new DisplayClassVariable { Variable = thisVar, Value = new LdLoc(thisVar) });
+				displayClass.Variables.Add(thisField, thisVar);
 			}
 			return true;
 
@@ -263,11 +243,17 @@ namespace ICSharpCode.Decompiler.IL.Transforms
 			return true;
 		}
 
-		bool IsDisplayClassLoad(ILInstruction target, out ILVariable variable)
+		ILFunction currentFunction;
+
+		protected internal override void VisitILFunction(ILFunction function)
 		{
-			if (target.MatchLdLoc(out variable) || target.MatchLdLoca(out variable))
-				return true;
-			return false;
+			var oldFunction = this.currentFunction;
+			try {
+				this.currentFunction = function;
+				base.VisitILFunction(function);
+			} finally {
+				this.currentFunction = oldFunction;
+			}
 		}
 
 		protected override void Default(ILInstruction inst)
@@ -280,105 +266,102 @@ namespace ICSharpCode.Decompiler.IL.Transforms
 		protected internal override void VisitStLoc(StLoc inst)
 		{
 			base.VisitStLoc(inst);
-			// Sometimes display class references are copied into other local variables.
-			// We remove the assignment and store the relationship between the display class and the variable in the
-			// displayClasses dictionary.
-			if (inst.Value.MatchLdLoc(out var closureVariable) && displayClasses.TryGetValue(closureVariable, out var displayClass)) {
-				displayClasses[inst.Variable] = displayClass;
-				instructionsToRemove.Add(inst);
-			} else if (inst.Variable.Kind == VariableKind.Local && inst.Variable.IsSingleDefinition && inst.Variable.LoadCount == 0 && inst.Value is StLoc) {
+
+			if (inst.Variable.Kind == VariableKind.Local && inst.Variable.IsSingleDefinition && inst.Variable.LoadCount == 0 && inst.Value is StLoc) {
+				context.Step($"Remove unused variable assignment {inst.Variable.Name}", inst);
 				inst.ReplaceWith(inst.Value);
 			}
 		}
 
 		protected internal override void VisitStObj(StObj inst)
 		{
-			base.VisitStObj(inst);
-			// This instruction has been marked deletable, do not transform it further
-			if (instructionsToRemove.Contains(inst))
-				return;
-			// The target of the store instruction must be a field reference
-			if (!inst.Target.MatchLdFlda(out ILInstruction target, out IField field))
-				return;
-			// Get display class info
-			if (!IsDisplayClassLoad(target, out var displayClassLoad) || !displayClasses.TryGetValue(displayClassLoad, out var displayClass))
-				return;
-			field = (IField)field.MemberDefinition;
-			if (displayClass.Variables.TryGetValue(field, out DisplayClassVariable info)) {
-				// If the display class field was previously initialized, we use a simple assignment.
-				inst.ReplaceWith(new StLoc(info.Variable, inst.Value).WithILRange(inst));
-			} else {
-				// This is an uninitialized variable:
-				ILInstruction value;
-				if (inst.Value.MatchLdLoc(out var v) && v.Kind == VariableKind.Parameter && currentFunction == v.Function) {
-					// Special case for parameters: remove copies of parameter values.
+			inst.Value.AcceptVisitor(this);
+			if (IsParameterAssignment(inst, out var displayClass, out var field, out var parameter)) {
+				context.Step($"Detected parameter assignment {parameter.Name}", inst);
+				displayClass.Variables.Add(field, parameter);
+				instructionsToRemove.Add(inst);
+			} else if (IsDisplayClassAssignment(inst, out displayClass, out field, out var variable)) {
+				context.Step($"Detected display-class assignment {variable.Name}", inst);
+				displayClass.Variables.Add(field, variable);
 				instructionsToRemove.Add(inst);
-					value = inst.Value;
 			} else {
-					context.Step($"Introduce captured variable for {field.FullName}", inst);
-					Debug.Assert(displayClass.Definition == field.DeclaringTypeDefinition);
-					// Introduce a fresh variable for the display class field.
-					if (displayClass.IsMono && displayClass.CaptureScope == null && !IsOuterClosureReference(field)) {
-						displayClass.CaptureScope = BlockContainer.FindClosestContainer(inst);
-					}
-					v = displayClass.DeclaringFunction.RegisterVariable(VariableKind.Local, field.Type, field.Name);
-					v.HasInitialValue = true;
-					v.CaptureScope = displayClass.CaptureScope;
-					inst.ReplaceWith(new StLoc(v, inst.Value).WithILRange(inst));
-					value = new LdLoc(v);
-				}
-				displayClass.Variables.Add(field, new DisplayClassVariable { Value = value, Variable = v });
+				inst.Target.AcceptVisitor(this);
+				EarlyExpressionTransforms.StObjToStLoc(inst, context);
 			}
 		}
 
 		protected internal override void VisitLdObj(LdObj inst)
 		{
 			base.VisitLdObj(inst);
-			// The target of the store instruction must be a field reference
-			if (!inst.Target.MatchLdFlda(out var target, out IField field))
-				return;
-			// Get display class info
-			if (!IsDisplayClassLoad(target, out var displayClassLoad) || !displayClasses.TryGetValue(displayClassLoad, out var displayClass))
-				return;
-			// Get display class variable info
-			if (!displayClass.Variables.TryGetValue((IField)field.MemberDefinition, out DisplayClassVariable info))
-				return;
-			// Replace usage of display class field with the variable.
-			var replacement = info.Value.Clone();
-			replacement.SetILRange(inst);
-			inst.ReplaceWith(replacement);
+			EarlyExpressionTransforms.LdObjToLdLoc(inst, context);
+		}
+
+		private bool IsDisplayClassLoad(ILInstruction target, out ILVariable variable)
+		{
+			// We cannot use MatchLdLocRef here because local functions use ref parameters
+			if (target.MatchLdLoc(out variable) || target.MatchLdLoca(out variable))
+				return true;
+			return false;
+		}
+
+		private bool IsDisplayClassAssignment(StObj inst, out DisplayClass displayClass, out IField field, out ILVariable variable)
+		{
+			variable = null;
+			if (!IsDisplayClassFieldAccess(inst.Target, out var displayClassVar, out displayClass, out field))
+				return false;
+			if (!(inst.Value.MatchLdLoc(out var v) && displayClasses.ContainsKey(v)))
+				return false;
+			if (displayClassVar.Function != currentFunction)
+				return false;
+			variable = v;
+			return true;
+		}
+
+		private bool IsParameterAssignment(StObj inst, out DisplayClass displayClass, out IField field, out ILVariable parameter)
+		{
+			parameter = null;
+			if (!IsDisplayClassFieldAccess(inst.Target, out var displayClassVar, out displayClass, out field))
+				return false;
+			if (!(inst.Value.MatchLdLoc(out var v) && v.Kind == VariableKind.Parameter && v.Function == currentFunction))
+				return false;
+			if (displayClass.Variables.ContainsKey(field))
+				return false;
+			if (displayClassVar.Function != currentFunction)
+				return false;
+			parameter = v;
+			return true;
+		}
+
+		private bool IsDisplayClassFieldAccess(ILInstruction inst, out ILVariable displayClassVar, out DisplayClass displayClass, out IField field)
+		{
+			displayClass = null;
+			displayClassVar = null;
+			field = null;
+			if (!(inst is LdFlda ldflda))
+				return false;
+			field = (IField)ldflda.Field.MemberDefinition;
+			return IsDisplayClassLoad(ldflda.Target, out displayClassVar)
+				&& displayClasses.TryGetValue(displayClassVar, out displayClass);
 		}
 
 		protected internal override void VisitLdFlda(LdFlda inst)
 		{
 			base.VisitLdFlda(inst);
-			// TODO : Figure out why this was added in https://github.com/icsharpcode/ILSpy/pull/1303
-			if (inst.Target.MatchLdThis() && inst.Field.Name == "$this"
-				&& inst.Field.MemberDefinition.ReflectionName.Contains("c__Iterator")) {
-				//Debug.Assert(false, "This should not be executed!");
-				var variable = currentFunction.Variables.First((f) => f.Index == -1);
-				inst.ReplaceWith(new LdLoca(variable).WithILRange(inst));
-			}
-			// Skip stfld/ldfld
-			if (inst.Parent is LdObj || inst.Parent is StObj)
-				return;
 			// Get display class info
-			if (!IsDisplayClassLoad(inst.Target, out var displayClassLoad) || !displayClasses.TryGetValue(displayClassLoad, out var displayClass))
+			if (!IsDisplayClassFieldAccess(inst, out _, out DisplayClass displayClass, out IField field))
 				return;
-			var field = (IField)inst.Field.MemberDefinition;
-			if (!displayClass.Variables.TryGetValue(field, out DisplayClassVariable info)) {
+			if (!displayClass.Variables.TryGetValue(field, out var v)) {
 				context.Step($"Introduce captured variable for {field.FullName}", inst);
 				// Introduce a fresh variable for the display class field.
 				Debug.Assert(displayClass.Definition == field.DeclaringTypeDefinition);
-				var v = displayClass.DeclaringFunction.RegisterVariable(VariableKind.Local, field.Type, field.Name);
+				v = displayClass.DeclaringFunction.RegisterVariable(VariableKind.Local, field.Type, field.Name);
 				v.HasInitialValue = true;
 				v.CaptureScope = displayClass.CaptureScope;
 				inst.ReplaceWith(new LdLoca(v).WithILRange(inst));
-				displayClass.Variables.Add(field, new DisplayClassVariable { Value = new LdLoc(v), Variable = v });
-			} else if (info.Value is LdLoc l) {
-				inst.ReplaceWith(new LdLoca(l.Variable).WithILRange(inst));
+				displayClass.Variables.Add(field, v);
 			} else {
-				Debug.Fail("LdFlda pattern not supported!");
+				context.Step($"Reuse captured variable {v.Name} for {field.FullName}", inst);
+				inst.ReplaceWith(new LdLoca(v).WithILRange(inst));
 			}
 		}
 	}