From dfeb39f6ee77147931b4f794e3837d520766a17c Mon Sep 17 00:00:00 2001
From: Siegfried Pammer <siegfriedpammer@gmail.com>
Date: Sun, 24 Sep 2017 00:06:07 +0200
Subject: [PATCH] TransformForeach: handle optional return statement after
 loop.

---
 .../CSharp/StatementBuilder.cs                | 44 ++++++++++++++++---
 .../IL/Transforms/UsingTransform.cs           | 16 +------
 2 files changed, 40 insertions(+), 20 deletions(-)

diff --git a/ICSharpCode.Decompiler/CSharp/StatementBuilder.cs b/ICSharpCode.Decompiler/CSharp/StatementBuilder.cs
index bbd773705..101ff7e04 100644
--- a/ICSharpCode.Decompiler/CSharp/StatementBuilder.cs
+++ b/ICSharpCode.Decompiler/CSharp/StatementBuilder.cs
@@ -271,14 +271,14 @@ namespace ICSharpCode.Decompiler.CSharp
 			};
 		}
 
-		ForeachStatement TransformToForeach(UsingInstruction inst, out TranslatedExpression resource)
+		Statement TransformToForeach(UsingInstruction inst, out TranslatedExpression resource)
 		{
 			resource = exprBuilder.Translate(inst.ResourceExpression);
 			var m = getEnumeratorPattern.Match(resource.Expression);
 			if (!(inst.Body is BlockContainer container) || !m.Success)
 				return null;
 			var enumeratorVar = m.Get<IdentifierExpression>("enumerator").Single().GetILVariable();
-			var loop = DetectedLoop.DetectLoop(container);
+			var loop = DetectedLoop.DetectLoop(UnwrapNestedContainerIfPossible(container, out var optionalReturnAfterLoop));
 			if (loop.Kind != LoopKind.While || !(loop.Body is Block body))
 				return null;
 			var condition = exprBuilder.TranslateCondition(loop.Conditions.Single());
@@ -319,15 +319,43 @@ namespace ICSharpCode.Decompiler.CSharp
 				itemVariable.Kind = VariableKind.ForeachLocal;
 				itemVariable.Name = AssignVariableNames.GenerateVariableName(currentFunction, collectionExpr.Annotation<ILInstruction>(), "item", itemVariable);
 			}
-			var whileLoop = (WhileStatement)ConvertAsBlock(inst.Body).Single();
+			var whileLoop = (WhileStatement)ConvertAsBlock(inst.Body).First();
 			BlockStatement foreachBody = (BlockStatement)whileLoop.EmbeddedStatement.Detach();
 			foreachBody.Statements.First().Detach();
-			return new ForeachStatement {
+			var foreachStmt = new ForeachStatement {
 				VariableType = settings.AnonymousTypes && itemVariable.Type.ContainsAnonymousType() ? new SimpleType("var") : exprBuilder.ConvertType(itemVariable.Type),
 				VariableName = itemVariable.Name,
 				InExpression = collectionExpr.Detach(),
 				EmbeddedStatement = foreachBody
 			};
+			if (optionalReturnAfterLoop != null) {
+				return new BlockStatement {
+					Statements = {
+						foreachStmt,
+						optionalReturnAfterLoop.AcceptVisitor(this)
+					}
+				};
+			}
+			return foreachStmt;
+		}
+
+		BlockContainer UnwrapNestedContainerIfPossible(BlockContainer container, out Leave optionalReturnInst)
+		{
+			optionalReturnInst = null;
+			if (container.Blocks.Count != 1)
+				return container;
+			var nestedBlock = container.Blocks[0];
+			if (nestedBlock.Instructions.Count != 2 ||
+				!(nestedBlock.Instructions[0] is BlockContainer nestedContainer) ||
+				!(nestedBlock.Instructions[1] is Leave leave))
+				return container;
+			if (leave.MatchLeave(container))
+				return nestedContainer;
+			if (leave.IsLeavingFunction) {
+				optionalReturnInst = leave;
+				return nestedContainer;
+			}
+			return container;
 		}
 
 		bool BodyHasSingleGetCurrent(Block body, ILVariable enumerator, ILInstruction moveNextUsage, out CallInstruction singleGetter, out bool needsUninlining, out ILVariable existingVariable)
@@ -504,7 +532,13 @@ namespace ICSharpCode.Decompiler.CSharp
 						// skip the final 'leave' instruction and just fall out of the BlockStatement
 						continue;
 					}
-					blockStatement.Add(Convert(inst));
+					var stmt = Convert(inst);
+					if (stmt is BlockStatement b) {
+						foreach (var nested in b.Statements)
+							blockStatement.Add(nested.Detach());
+					} else {
+						blockStatement.Add(stmt);
+					}
 				}
 				if (block.FinalInstruction.OpCode != OpCode.Nop) {
 					blockStatement.Add(Convert(block.FinalInstruction));
diff --git a/ICSharpCode.Decompiler/IL/Transforms/UsingTransform.cs b/ICSharpCode.Decompiler/IL/Transforms/UsingTransform.cs
index 67710e2eb..2f932184f 100644
--- a/ICSharpCode.Decompiler/IL/Transforms/UsingTransform.cs
+++ b/ICSharpCode.Decompiler/IL/Transforms/UsingTransform.cs
@@ -92,24 +92,10 @@ namespace ICSharpCode.Decompiler.IL.Transforms
 			}
 			context.Step("UsingTransform", body);
 			block.Instructions.RemoveAt(i);
-			block.Instructions[i - 1] = new UsingInstruction(resourceExpression, UnwrapNestedContainerIfPossible(body.TryBlock));
+			block.Instructions[i - 1] = new UsingInstruction(resourceExpression, body.TryBlock);
 			return true;
 		}
 
-		ILInstruction UnwrapNestedContainerIfPossible(ILInstruction tryBlock)
-		{
-			if (!(tryBlock is BlockContainer container))
-				return tryBlock;
-			if (container.Blocks.Count != 1)
-				return tryBlock;
-			var nestedBlock = container.Blocks[0];
-			if (nestedBlock.Instructions.Count != 2 ||
-				!(nestedBlock.Instructions[0] is BlockContainer nestedContainer) ||
-				!nestedBlock.Instructions[1].MatchLeave(container))
-				return tryBlock;
-			return nestedContainer;
-		}
-
 		bool MatchDisposeBlock(BlockContainer container, ILVariable objVar, bool usingNull)
 		{
 			var entryPoint = container.EntryPoint;