ILSpy/ICSharpCode.Decompiler/IL/ControlFlow/YieldReturnDecompiler.cs

// Copyright (c) 2011 AlphaSierraPapa for the SharpDevelop Team
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy of this
// software and associated documentation files (the "Software"), to deal in the Software
// without restriction, including without limitation the rights to use, copy, modify, merge,
// publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons
// to whom the Software is furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in all copies or
// substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
// INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
// PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
// FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
// OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

using ICSharpCode.Decompiler.CSharp;
using ICSharpCode.Decompiler.IL.Transforms;
using ICSharpCode.Decompiler.TypeSystem;
using ICSharpCode.Decompiler.Util;
using Mono.Cecil;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace ICSharpCode.Decompiler.IL.ControlFlow
{
	class YieldReturnDecompiler : IILTransform
	{
		// For a description on the code generated by the C# compiler for yield return:
		// http://csharpindepth.com/Articles/Chapter6/IteratorBlockImplementation.aspx

		// The idea here is:
		// - Figure out whether the current method is instanciating an enumerator
		// - Figure out which of the fields is the state field
		// - Construct an exception table based on states. This allows us to determine, for each state, what the parent try block is.

		// See http://community.sharpdevelop.net/blogs/danielgrunwald/archive/2011/03/06/ilspy-yield-return.aspx
		// for a description of this step.

		ILTransformContext context;

		/// <summary>The type that contains the function being decompiled.</summary>
		TypeDefinition currentType;

		/// <summary>The compiler-generated enumerator class.</summary>
		/// <remarks>Set in MatchEnumeratorCreationPattern()</remarks>
		TypeDefinition enumeratorType;

		/// <summary>The constructor of the compiler-generated enumerator class.</summary>
		/// <remarks>Set in MatchEnumeratorCreationPattern()</remarks>
		MethodDefinition enumeratorCtor;

		/// <summary>The dispose method of the compiler-generated enumerator class.</summary>
		/// <remarks>Set in ConstructExceptionTable()</remarks>
		MethodDefinition disposeMethod;

		/// <summary>The field in the compiler-generated class holding the current state of the state machine</summary>
		/// <remarks>Set in AnalyzeCtor()</remarks>
		IField stateField;

		/// <summary>The backing field of the 'Current' property in the compiler-generated class</summary>
		/// <remarks>Set in AnalyzeCurrentProperty()</remarks>
		IField currentField;

		/// <summary>Maps the fields of the compiler-generated class to the original parameters.</summary>
		/// <remarks>Set in MatchEnumeratorCreationPattern() and ResolveIEnumerableIEnumeratorFieldMapping()</remarks>
		readonly Dictionary<IField, ILVariable> fieldToParameterMap = new Dictionary<IField, ILVariable>();

		/// <summary>This dictionary stores the information extracted from the Dispose() method:
		/// for each "Finally Method", it stores the set of states for which the method is being called.</summary>
		/// <remarks>Set in ConstructExceptionTable()</remarks>
		Dictionary<IMethod, LongSet> finallyMethodToStateRange;

		/// <summary>
		/// List of blocks that change the iterator state on block entry.
		/// </summary>
		readonly List<(int state, Block block)> stateChanges = new List<(int state, Block block)>();

		#region Run() method
		public void Run(ILFunction function, ILTransformContext context)
		{
			if (!context.Settings.YieldReturn)
				return; // abort if enumerator decompilation is disabled
			this.context = context;
			this.currentType = function.Method.DeclaringType;
			this.enumeratorType = null;
			this.enumeratorCtor = null;
			this.stateField = null;
			this.currentField = null;
			this.fieldToParameterMap.Clear();
			this.finallyMethodToStateRange = null;
			this.stateChanges.Clear();
			if (!MatchEnumeratorCreationPattern(function))
				return;
			BlockContainer newBody;
			try {
				AnalyzeCtor();
				AnalyzeCurrentProperty();
				ResolveIEnumerableIEnumeratorFieldMapping();
				ConstructExceptionTable();
				newBody = AnalyzeMoveNext();
			} catch (SymbolicAnalysisFailedException) {
				return;
			}

			context.Step("Replacing body with MoveNext() body", function);
			function.IsIterator = true;
			function.Body = newBody;
			// register any locals used in newBody
			function.Variables.AddRange(newBody.Descendants.OfType<IInstructionWithVariableOperand>().Select(inst => inst.Variable).Distinct());
			function.CheckInvariant(ILPhase.Normal);

			PrintFinallyMethodStateRanges(newBody);

			context.Step("Delete unreachable blocks", function);
			// Note: because this only deletes blocks outright, the 'stateChanges' entries remain valid
			// (though some may point to now-deleted blocks)
			newBody.SortBlocks(deleteUnreachableBlocks: true);

			context.Step("Reconstruct try-finally blocks", function);
			ReconstructTryFinallyBlocks(newBody);

			context.Step("Translate fields to local accesses", function);
			TranslateFieldsToLocalAccess(function, function, fieldToParameterMap);

			// Re-run control flow simplification over the newly constructed set of gotos,
			// and inlining because TranslateFieldsToLocalAccess() might have opened up new inlining opportunities.
			function.RunTransforms(CSharpDecompiler.EarlyILTransforms(), context);
		}
		#endregion

		#region Match the enumerator creation pattern
		bool MatchEnumeratorCreationPattern(ILFunction function)
		{
			Block body = SingleBlock(function.Body);
			if (body == null || body.Instructions.Count == 0) {
				return false;
			}

			ILInstruction newObj;
			if (body.Instructions.Count == 1) {
				// No parameters passed to enumerator (not even 'this'):
				// ret(newobj(...))
				if (body.Instructions[0].MatchReturn(out newObj))
					return MatchEnumeratorCreationNewObj(newObj);
				else
					return false;
			}

			// If there's parameters passed to the helper class, the class instance is first
			// stored in a variable, then the parameters are copied over, then the instance is returned.

			// stloc(var_1, newobj(..))
			if (!body.Instructions[0].MatchStLoc(out var var1, out newObj))
				return false;
			if (!MatchEnumeratorCreationNewObj(newObj))
				return false;

			int i;
			for (i = 1; i < body.Instructions.Count; i++) {
				// stfld(..., ldloc(var_1), ldloc(parameter))
				if (!body.Instructions[i].MatchStFld(out var ldloc, out var storedField, out var loadParameter))
					break;
				if (ldloc.MatchLdLoc(var1)
					&& loadParameter.MatchLdLoc(out var parameter)
					&& parameter.Kind == VariableKind.Parameter) {
					fieldToParameterMap[(IField)storedField.MemberDefinition] = parameter;
				} else {
					return false;
				}
			}

			// In debug builds, the compiler may copy the var1 into another variable (var2) before returning it.
			if (i < body.Instructions.Count
				&& body.Instructions[i].MatchStLoc(out var var2, out var ldlocForStloc2)
				&& ldlocForStloc2.MatchLdLoc(var1)) {
				// stloc(var_2, ldloc(var_1))
				i++;
			} else {
				// in release builds, var1 is returned directly
				var2 = var1;
			}
			if (i < body.Instructions.Count
				&& body.Instructions[i].MatchReturn(out var retVal)
				&& retVal.MatchLdLoc(var2)) {
				// ret(ldloc(var_2))
				return true;
			} else {
				return false;
			}
		}

		/// <summary>
		/// Matches the body of a method as a single basic block.
		/// </summary>
		static Block SingleBlock(ILInstruction body)
		{
			var block = body as Block;
			if (body is BlockContainer blockContainer && blockContainer.Blocks.Count == 1) {
				block = blockContainer.Blocks.Single() as Block;
			}
			return block;
		}

		/// <summary>
		/// Matches the newobj instruction that creates an instance of the compiler-generated enumerator helper class.
		/// </summary>
		bool MatchEnumeratorCreationNewObj(ILInstruction inst)
		{
			// newobj(CurrentType/...::.ctor, ldc.i4(-2))
			if (!(inst is NewObj newObj))
				return false;
			if (newObj.Arguments.Count != 1)
				return false;
			if (!newObj.Arguments[0].MatchLdcI4(out int initialState))
				return false;
			if (!(initialState == -2 || initialState == 0))
				return false;
			enumeratorCtor = context.TypeSystem.GetCecil(newObj.Method) as MethodDefinition;
			enumeratorType = enumeratorCtor?.DeclaringType;
			return enumeratorType?.DeclaringType == currentType
				&& IsCompilerGeneratorEnumerator(enumeratorType);
		}

		public static bool IsCompilerGeneratorEnumerator(TypeDefinition type)
		{
			if (!(type?.DeclaringType != null && type.IsCompilerGenerated()))
				return false;
			foreach (var i in type.Interfaces) {
				var tr = i.InterfaceType;
				if (tr.Namespace == "System.Collections" && tr.Name == "IEnumerator")
					return true;
			}
			return false;
		}
		#endregion

		#region Figure out what the 'state' field is (analysis of .ctor())
		/// <summary>
		/// Looks at the enumerator's ctor and figures out which of the fields holds the state.
		/// </summary>
		void AnalyzeCtor()
		{
			Block body = SingleBlock(CreateILAst(enumeratorCtor).Body);
			if (body == null)
				throw new SymbolicAnalysisFailedException("Missing enumeratorCtor.Body");
			foreach (var inst in body.Instructions) {
				if (inst.MatchStFld(out var target, out var field, out var value)
					&& target.MatchLdThis()
					&& value.MatchLdLoc(out var arg)
					&& arg.Kind == VariableKind.Parameter && arg.Index == 0) {
					stateField = (IField)field.MemberDefinition;
				}
			}
			if (stateField == null)
				throw new SymbolicAnalysisFailedException("Could not find stateField");
		}

		/// <summary>
		/// Creates ILAst for the specified method, optimized up to before the 'YieldReturn' step.
		/// </summary>
		ILFunction CreateILAst(MethodDefinition method)
		{
			if (method == null || !method.HasBody)
				throw new SymbolicAnalysisFailedException();

			var il = new ILReader(context.TypeSystem).ReadIL(method.Body, context.CancellationToken);
			il.RunTransforms(CSharpDecompiler.EarlyILTransforms(), new ILTransformContext {
				Settings = context.Settings,
				CancellationToken = context.CancellationToken,
				TypeSystem = context.TypeSystem
			});
			return il;
		}
		#endregion

		#region Figure out what the 'current' field is (analysis of get_Current())
		/// <summary>
		/// Looks at the enumerator's get_Current method and figures out which of the fields holds the current value.
		/// </summary>
		void AnalyzeCurrentProperty()
		{
			MethodDefinition getCurrentMethod = enumeratorType.Methods.FirstOrDefault(
				m => m.Name.StartsWith("System.Collections.Generic.IEnumerator", StringComparison.Ordinal)
				&& m.Name.EndsWith(".get_Current", StringComparison.Ordinal));
			Block body = SingleBlock(CreateILAst(getCurrentMethod).Body);
			if (body == null)
				throw new SymbolicAnalysisFailedException();
			if (body.Instructions.Count == 1) {
				// release builds directly return the current field
				// ret(ldfld F(ldloc(this)))
				if (body.Instructions[0].MatchReturn(out var retVal)
					&& retVal.MatchLdFld(out var target, out var field)
					&& target.MatchLdThis()) {
					currentField = (IField)field.MemberDefinition;
				}
			} else if (body.Instructions.Count == 2) {
				// debug builds store the return value in a temporary
				// stloc V = ldfld F(ldloc(this))
				// ret(ldloc V)
				if (body.Instructions[0].MatchStLoc(out var v, out var ldfld)
					&& ldfld.MatchLdFld(out var target, out var field)
					&& target.MatchLdThis()
					&& body.Instructions[1].MatchReturn(out var retVal)
					&& retVal.MatchLdLoc(v)) {
					currentField = (IField)field.MemberDefinition;
				}
			}
			if (currentField == null)
				throw new SymbolicAnalysisFailedException("Could not find currentField");
		}
		#endregion

		#region Figure out the mapping of IEnumerable fields to IEnumerator fields  (analysis of GetEnumerator())
		void ResolveIEnumerableIEnumeratorFieldMapping()
		{
			MethodDefinition getEnumeratorMethod = enumeratorType.Methods.FirstOrDefault(
				m => m.Name.StartsWith("System.Collections.Generic.IEnumerable", StringComparison.Ordinal)
				&& m.Name.EndsWith(".GetEnumerator", StringComparison.Ordinal));
			if (getEnumeratorMethod == null)
				return; // no mappings (maybe it's just an IEnumerator implementation?)
			var function = CreateILAst(getEnumeratorMethod);
			foreach (var block in function.Descendants.OfType<Block>()) {
				foreach (var inst in block.Instructions) {
					// storeTarget.storeField = this.loadField;
					if (inst.MatchStFld(out var storeTarget, out var storeField, out var storeValue)
						&& storeValue.MatchLdFld(out var loadTarget, out var loadField)
						&& loadTarget.MatchLdThis()) {
						storeField = (IField)storeField.MemberDefinition;
						loadField = (IField)loadField.MemberDefinition;
						if (fieldToParameterMap.TryGetValue(loadField, out var mappedParameter))
							fieldToParameterMap[storeField] = mappedParameter;
					}
				}
			}
		}
		#endregion

		#region Construction of the exception table (analysis of Dispose())
		// We construct the exception table by analyzing the enumerator's Dispose() method.

		void ConstructExceptionTable()
		{
			disposeMethod = enumeratorType.Methods.FirstOrDefault(m => m.Name == "System.IDisposable.Dispose");
			var function = CreateILAst(disposeMethod);

			var rangeAnalysis = new StateRangeAnalysis(StateRangeAnalysisMode.IteratorDispose, stateField);
			rangeAnalysis.AssignStateRanges(function.Body, LongSet.Universe);
			finallyMethodToStateRange = rangeAnalysis.finallyMethodToStateRange;
		}
		
		[Conditional("DEBUG")]
		void PrintFinallyMethodStateRanges(BlockContainer bc)
		{
			foreach (var (method, stateRange) in finallyMethodToStateRange) {
				bc.Blocks[0].Instructions.Insert(0, new Nop {
					Comment = method.Name + " in " + stateRange
				});
			}
		}
		#endregion

		#region Analyze MoveNext() and generate new body
		BlockContainer AnalyzeMoveNext()
		{
			MethodDefinition moveNextMethod = enumeratorType.Methods.FirstOrDefault(m => m.Name == "MoveNext");
			ILFunction moveNextFunction = CreateILAst(moveNextMethod);

			var body = (BlockContainer)moveNextFunction.Body;
			if (body.Blocks.Count == 1 && body.Blocks[0].Instructions.Count == 1 && body.Blocks[0].Instructions[0] is TryFault tryFault) {
				body = (BlockContainer)tryFault.TryBlock;
				var faultBlockContainer = tryFault.FaultBlock as BlockContainer;
				if (faultBlockContainer?.Blocks.Count != 1)
					throw new SymbolicAnalysisFailedException("Unexpected number of blocks in MoveNext() fault block");
				var faultBlock = faultBlockContainer.Blocks.Single();
				if (!(faultBlock.Instructions.Count == 2
					&& faultBlock.Instructions[0] is Call call
					&& context.TypeSystem.GetCecil(call.Method) == disposeMethod
					&& call.Arguments.Count == 1
					&& call.Arguments[0].MatchLdThis()
					&& faultBlock.Instructions[1].MatchLeave(faultBlockContainer))) {
					throw new SymbolicAnalysisFailedException("Unexpected fault block contents in MoveNext()");
				}
			}

			// Note: body may contain try-catch or try-finally statements that have nested block containers,
			// but those cannot contain any yield statements.
			// So for reconstructing the control flow, we only need at the blocks directly within body.

			var rangeAnalysis = new StateRangeAnalysis(StateRangeAnalysisMode.IteratorMoveNext, stateField);
			rangeAnalysis.AssignStateRanges(body, LongSet.Universe);

			var newBody = ConvertBody(body, rangeAnalysis.GetBlockStateSetMapping(body));
			moveNextFunction.Variables.Clear();
			// release references from old moveNextFunction to instructions that were moved over to newBody
			moveNextFunction.ReleaseRef();
			return newBody;
		}

		/// <summary>
		/// Convert the old body (of MoveNext function) to the new body (of decompiled iterator method).
		/// 
		/// * Replace the sequence
		///       this.currentField = expr;
		///       this.state = N;
		///       return true;
		///     with:
		///       yield return expr;
		///       goto blockForState(N);
		///  * Replace the sequence:
		///       this._finally2();
		///       this._finally1();
		///       return false;
		///     with:
		///       yield break;
		///  * Reconstruct try-finally blocks from
		///      (on enter) this.state = N;
		///      (on exit)  this._finallyX();
		/// </summary>
		private BlockContainer ConvertBody(BlockContainer oldBody, IEnumerable<(Block, LongSet)> blockStateSets)
		{
			BlockContainer newBody = new BlockContainer();
			// create all new blocks so that they can be referenced by gotos
			for (int blockIndex = 0; blockIndex < oldBody.Blocks.Count; blockIndex++) {
				newBody.Blocks.Add(new Block { ILRange = oldBody.Blocks[blockIndex].ILRange });
			}
			// convert contents of blocks
			var ssaDefs = new Dictionary<ILVariable, ILInstruction>();
			
			for (int i = 0; i < oldBody.Blocks.Count; i++) {
				var oldBlock = oldBody.Blocks[i];
				var newBlock = newBody.Blocks[i];
				foreach (var oldInst in oldBlock.Instructions) {
					if (oldInst.MatchStFld(out var target, out var field, out var value) && target.MatchLdThis()) {
						if (field.MemberDefinition.Equals(stateField)) {
							if (value.MatchLdcI4(out int newState)) {
								// On state change, break up the block (if necessary):
								if (newBlock.Instructions.Count > 0) {
									var newBlock2 = new Block();
									newBlock2.ILRange = new Interval(oldInst.ILRange.Start, oldInst.ILRange.Start);
									newBody.Blocks.Add(newBlock2);
									newBlock.Instructions.Add(new Branch(newBlock2));
									newBlock = newBlock2;
								}
#if DEBUG
								newBlock.Instructions.Add(new Nop { Comment = "iterator._state = " + newState });
#endif
								stateChanges.Add((newState, newBlock));
							} else {
								newBlock.Instructions.Add(new InvalidExpression("Assigned non-constant to iterator.state field") {
									ILRange = oldInst.ILRange
								});
							}
							continue; // don't copy over this instruction, but continue with the basic block
						} else if (field.MemberDefinition.Equals(currentField)) {
							// create yield return
							newBlock.Instructions.Add(new YieldReturn(value) { ILRange = oldInst.ILRange });
							ConvertBranchAfterYieldReturn(newBlock, oldBlock, oldInst.ChildIndex);
							break; // we're done with this basic block
						}
					} else if (oldInst.MatchReturn(out value)) {
						if (value.MatchLdLoc(out var v)) {
							ssaDefs.TryGetValue(v, out value);
						}
						if (value.MatchLdcI4(0)) {
							// yield break
							newBlock.Instructions.Add(new Leave(newBody) { ILRange = oldInst.ILRange });
						} else {
							newBlock.Instructions.Add(new InvalidBranch("Unexpected return in MoveNext()"));
						}
						break; // we're done with this basic block
					} else if (oldInst.MatchStLoc(out var v, out value) && v.IsSingleDefinition) {
						ssaDefs.Add(v, value);
					}
					// copy over the instruction to the new block
					UpdateBranchTargets(oldInst);
					newBlock.Instructions.Add(oldInst);
				}
			}

			// Insert new artificial block as entry point, and jump to state 0.
			// This causes the method to start directly at the first user code,
			// and the whole compiler-generated state-dispatching logic becomes unreachable code
			// and gets deleted.
			newBody.Blocks.Insert(0, new Block {
				Instructions = { MakeGoTo(0) }
			});
			return newBody;

			void ConvertBranchAfterYieldReturn(Block newBlock, Block oldBlock, int i)
			{
				if (!(oldBlock.Instructions[i + 1].MatchStFld(out var target, out var field, out var value)
					&& target.MatchLdThis()
					&& field.MemberDefinition == stateField
					&& value.MatchLdcI4(out int newState))) {
					newBlock.Instructions.Add(new InvalidBranch("Unable to find new state assignment for yield return"));
					return;
				}
				if (!(oldBlock.Instructions[i + 2].MatchReturn(out var retVal)
					&& retVal.MatchLdcI4(1))) {
					newBlock.Instructions.Add(new InvalidBranch("Unable to find 'return true' for yield return"));
					return;
				}
				newBlock.Instructions.Add(MakeGoTo(newState));
			}

			ILInstruction MakeGoTo(int v)
			{
				Block targetBlock = null;
				foreach (var (block, stateSet) in blockStateSets) {
					if (stateSet.Contains(v))
						targetBlock = block;
				}
				if (targetBlock != null)
					return new Branch(newBody.Blocks[targetBlock.ChildIndex]);
				else
					return new InvalidBranch("Could not find block for state " + v);
			}

			void UpdateBranchTargets(ILInstruction inst)
			{
				switch (inst) {
					case Branch branch:
						if (branch.TargetContainer == oldBody) {
							branch.TargetBlock = newBody.Blocks[branch.TargetBlock.ChildIndex];
						}
						break;
					case Leave leave:
						if (leave.TargetContainer == oldBody) {
							leave.TargetContainer = newBody;
						}
						break;
				}
				foreach (var child in inst.Children) {
					UpdateBranchTargets(child);
				}
			}
		}
		#endregion

		#region TranslateFieldsToLocalAccess
		/// <summary>
		/// Translates all field accesses in `function` to local variable accesses.
		/// </summary>
		internal static void TranslateFieldsToLocalAccess(ILFunction function, ILInstruction inst, Dictionary<IField, ILVariable> fieldToVariableMap)
		{
			if (inst is LdFlda ldflda && ldflda.Target.MatchLdThis()) {
				var fieldDef = (IField)ldflda.Field.MemberDefinition;
				if (!fieldToVariableMap.TryGetValue(fieldDef, out var v)) {
					string name = null;
					if (!string.IsNullOrEmpty(fieldDef.Name) && fieldDef.Name[0] == '<') {
						int pos = fieldDef.Name.IndexOf('>');
						if (pos > 1)
							name = fieldDef.Name.Substring(1, pos - 1);
					}
					v = function.RegisterVariable(VariableKind.Local, ldflda.Field.ReturnType, name);
					fieldToVariableMap.Add(fieldDef, v);
				}
				inst.ReplaceWith(new LdLoca(v));
			} else if (inst.MatchLdThis()) {
				inst.ReplaceWith(new InvalidExpression("iterator") { ExpectedResultType = inst.ResultType });
			} else {
				foreach (var child in inst.Children) {
					TranslateFieldsToLocalAccess(function, child, fieldToVariableMap);
				}
			}
		}
		#endregion

		#region Reconstruct try-finally blocks

		private void ReconstructTryFinallyBlocks(BlockContainer newBody)
		{
			// TODO
		}

		#endregion
	}
}