// 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.FlowAnalysis;
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>
		/// For each finally method, stores the target state when entering the finally block,
		/// and the decompiled code of the finally method body.
		/// </summary>
		readonly Dictionary<IMethod, (int? outerState, BlockContainer body)> decompiledFinallyMethods = new Dictionary<IMethod, (int? outerState, BlockContainer body)>();

		/// <summary>
		/// Temporary stores for 'yield break'.
		/// </summary>
		readonly List<StLoc> returnStores = new List<StLoc>();

		#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.decompiledFinallyMethods.Clear();
			this.returnStores.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);

			DecompileFinallyBlocks();
			ReconstructTryFinallyBlocks(newBody);

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

			if (returnStores.Count > 0) {
				context.Step("Remove temporaries", function);
				foreach (var store in returnStores) {
					if (store.Variable.LoadCount == 0 && store.Variable.AddressCount == 0 && store.Parent is Block block) {
						block.Instructions.Remove(store);
					}
				}
			}

			// 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
			
			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:
								// (this allows us to consider each block individually for try-finally reconstruction)
								newBlock = SplitBlock(newBlock, oldInst);
								// We keep the state-changing instruction around (as first instruction of the new block)
								// for reconstructing the try-finallys. 
							} 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 is Call call && call.Arguments.Count == 1 && call.Arguments[0].MatchLdThis()
						&& finallyMethodToStateRange.ContainsKey((IMethod)call.Method.MemberDefinition))
					{
						// Break up the basic block on a call to a finally method
						// (this allows us to consider each block individually for try-finally reconstruction)
						newBlock = SplitBlock(newBlock, oldInst);
					}
					// copy over the instruction to the new block
					newBlock.Instructions.Add(oldInst);
					UpdateBranchTargets(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));
			}

			Block SplitBlock(Block newBlock, ILInstruction oldInst)
			{
				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;
				}
				return newBlock;
			}

			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;
					case Return ret:
						ILInstruction value = ret.Value;
						if (value.MatchLdLoc(out var v) && v.IsSingleDefinition
							&& v.StoreInstructions.SingleOrDefault() is StLoc stloc)
						{
							returnStores.Add(stloc);
							value = stloc.Value;
						}
						if (value.MatchLdcI4(0)) {
							// yield break
							ret.ReplaceWith(new Leave(newBody) { ILRange = ret.ILRange });
						} else {
							ret.ReplaceWith(new InvalidBranch("Unexpected return in MoveNext()") { ILRange = ret.ILRange });
						}
						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 DecompileFinallyBlocks
		void DecompileFinallyBlocks()
		{
			foreach (var method in finallyMethodToStateRange.Keys) {
				var function = CreateILAst((MethodDefinition)context.TypeSystem.GetCecil(method));
				var body = (BlockContainer)function.Body;
				var newState = GetNewState(body.EntryPoint);
				if (newState != null) {
					body.EntryPoint.Instructions.RemoveAt(0);
				}
				function.ReleaseRef(); // make body reusable outside of function
				decompiledFinallyMethods.Add(method, (newState, body));
			}
		}
		#endregion

		#region Reconstruct try-finally blocks

		/// <summary>
		/// Reconstruct try-finally blocks.
		/// * The stateChanges (iterator._state = N;) tell us when to open a try-finally block
		/// * The calls to the finally method tell us when to leave the try block.
		/// 
		/// There might be multiple stateChanges for a given try-finally block, e.g.
		/// both the original entry point, and the target when leaving a nested block.
		/// In proper C# code, the entry point of the try-finally will dominate all other code
		/// in the try-block, so we can use dominance to find the proper entry point.
		/// 
		/// Precondition: the blocks in newBody are topologically sorted.
		/// </summary>
		void ReconstructTryFinallyBlocks(BlockContainer newBody)
		{
			context.Stepper.Step("Reconstuct try-finally blocks");
			var blockState = new int[newBody.Blocks.Count];
			blockState[0] = -1;
			var stateToContainer = new Dictionary<int, BlockContainer>();
			stateToContainer.Add(-1, newBody);
			// First, analyse the newBody: for each block, determine the active state number.
			foreach (var block in newBody.Blocks) {
				int oldState = blockState[block.ChildIndex];
				BlockContainer container; // new container for the block
				if (GetNewState(block) is int newState) {
					// OK, state change
					// Remove the state-changing instruction
					block.Instructions.RemoveAt(0);

					if (!stateToContainer.TryGetValue(newState, out container)) {
						// First time we see this state.
						// This means we just found the entry point of a try block.
						CreateTryBlock(block, newState);
						// CreateTryBlock() wraps the contents of 'block' with a TryFinally.
						// We thus need to put the block (which now contains the whole TryFinally)
						// into the parent container.
						// Assuming a state transition never enters more than one state at once,
						// we can use stateToContainer[oldState] as parent.
						container = stateToContainer[oldState];
					}
				} else {
					// Because newBody is topologically sorted we because we just removed unreachable code,
					// we can assume that blockState[] was already set for this block.
					newState = oldState;
					container = stateToContainer[oldState];
				}
				if (container != newBody) {
					// Move the block into the container.
					container.Blocks.Add(block);
					// Keep the stale reference in newBody.Blocks for now, to avoid
					// changing the ChildIndex of the other blocks while we use it
					// to index the blockState array.
				}
#if DEBUG
				block.Instructions.Insert(0, new Nop { Comment = "state == " + newState });
#endif
				// Propagate newState to successor blocks
				foreach (var branch in block.Descendants.OfType<Branch>()) {
					if (branch.TargetBlock.Parent == newBody) {
						Debug.Assert(blockState[branch.TargetBlock.ChildIndex] == newState || blockState[branch.TargetBlock.ChildIndex] == 0);
						blockState[branch.TargetBlock.ChildIndex] = newState;
					}
				}
			}
			newBody.Blocks.RemoveAll(b => b.Parent != newBody);

			void CreateTryBlock(Block block, int state)
			{
				var finallyMethod = FindFinallyMethod(state);
				Debug.Assert(finallyMethod != null);
				// remove the method so that it doesn't get cause ambiguity when processing nested try-finally blocks
				finallyMethodToStateRange.Remove(finallyMethod);

				var tryBlock = new Block();
				tryBlock.ILRange = block.ILRange;
				tryBlock.Instructions.AddRange(block.Instructions);
				var tryBlockContainer = new BlockContainer();
				tryBlockContainer.Blocks.Add(tryBlock);
				stateToContainer.Add(state, tryBlockContainer);

				ILInstruction finallyBlock;
				if (decompiledFinallyMethods.TryGetValue(finallyMethod, out var decompiledMethod)) {
					finallyBlock = decompiledMethod.body;
				} else {
					finallyBlock = new InvalidBranch("Missing decompiledFinallyMethod");
				}

				block.Instructions.Clear();
				block.Instructions.Add(new TryFinally(tryBlockContainer, finallyBlock));
			}

			IMethod FindFinallyMethod(int state)
			{
				IMethod foundMethod = null;
				foreach (var (method, stateRange) in finallyMethodToStateRange) {
					if (stateRange.Contains(state)) {
						if (foundMethod == null)
							foundMethod = method;
						else
							Debug.Fail("Ambiguous finally method for state " + state);
					}
				}
				return foundMethod;
			}
		}

		// Gets the state that is transitioned to at the start of the block
		int? GetNewState(Block block)
		{
			if (block.Instructions[0].MatchStFld(out var target, out var field, out var value)
				&& target.MatchLdThis()
				&& field.MemberDefinition.Equals(stateField)
				&& value.MatchLdcI4(out int newState))
			{
				return newState;
			} else if (block.Instructions[0] is Call call
				&& call.Arguments.Count == 1 && call.Arguments[0].MatchLdThis()
				&& decompiledFinallyMethods.TryGetValue((IMethod)call.Method.MemberDefinition, out var finallyMethod))
			{
				return finallyMethod.outerState;
			}
			return null;
		}
		#endregion
	}
}