From 5b8b80cc72dc189acc273055e0c6b8c2121454f3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20Srbeck=C3=BD?= <dsrbecky@gmail.com>
Date: Sun, 27 Jan 2008 23:13:37 +0000
Subject: [PATCH] Restore loop iterator - move assignment from the end of loop.

---
 src/Transforms/Ast/RestoreLoop.cs | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/Transforms/Ast/RestoreLoop.cs b/src/Transforms/Ast/RestoreLoop.cs
index 77f8a8c6e..a9ef14c36 100644
--- a/src/Transforms/Ast/RestoreLoop.cs
+++ b/src/Transforms/Ast/RestoreLoop.cs
@@ -12,6 +12,7 @@ namespace Decompiler.Transforms.Ast
 		{
 			base.VisitForStatement(forStatement, data);
 			
+			// Restore loop condition
 			if (forStatement.Condition.IsNull &&
 				forStatement.EmbeddedStatement.Children.Count >= 3)
 			{
@@ -25,13 +26,28 @@ namespace Decompiler.Transforms.Ast
 				{
 					MyGotoStatement gotoStmt = condition.TrueStatement[0] as MyGotoStatement;
 					if (gotoStmt != null && gotoStmt.NodeLabel == label.NodeLabel) {
-						forStatement.EmbeddedStatement.Children.RemoveAt(0);
-						forStatement.EmbeddedStatement.Children.RemoveAt(0);
+						forStatement.EmbeddedStatement.Children.Remove(condition);
+						forStatement.EmbeddedStatement.Children.Remove(breakStmt);
 						gotoStmt.NodeLabel.ReferenceCount--;
 						forStatement.Condition = condition.Condition;
 					}
 				}
 			}
+			
+			// Restore loop iterator
+			if (forStatement.EmbeddedStatement.Children.Count > 0 &&
+			    forStatement.Iterator.Count == 0)
+			{
+				ExpressionStatement lastStmt = forStatement.EmbeddedStatement.Children[forStatement.EmbeddedStatement.Children.Count - 1] as ExpressionStatement;
+				if (lastStmt != null) {
+					AssignmentExpression assign = lastStmt.Expression as AssignmentExpression;
+					if (assign != null) {
+						forStatement.EmbeddedStatement.Children.Remove(lastStmt);
+						forStatement.Iterator.Add(lastStmt);
+					}
+				}
+			}
+			
 			return null;
 		}
 	}