Determine loop condition from CFG

ICSharpCode.Decompiler/Ast/AstMethodBodyBuilder.cs

@@ -103,10 +103,17 @@ namespace Decompiler
 						throw new Exception();
 					}
 				}
-			} else if (node is ILLoop) {
-				yield return new Ast.ForStatement {
-					EmbeddedStatement = TransformBlock(((ILLoop)node).ContentBlock)
+			} else if (node is ILWhileLoop) {
+				ILWhileLoop ilLoop = (ILWhileLoop)node;
+				if (ilLoop.PreLoopLabel != null)
+					yield return TransformNode(ilLoop.PreLoopLabel).Single();
+				WhileStatement whileStmt = new WhileStatement() {
+					Condition = ilLoop.Condition != null ? MakeBranchCondition(ilLoop.Condition) : new PrimitiveExpression(true),
+					EmbeddedStatement = TransformBlock(ilLoop.BodyBlock)
 				};
+				yield return whileStmt;
+				if (ilLoop.PostLoopGoto != null)
+					yield return (Statement)TransformExpression(ilLoop.PostLoopGoto);
 			} else if (node is ILCondition) {
 				ILCondition conditionalNode = (ILCondition)node;
 				if (conditionalNode.FalseBlock.Body.Any()) {

ICSharpCode.Decompiler/Ast/Transforms/RemoveGotos.cs

@@ -7,9 +7,9 @@ namespace Decompiler.Transforms.Ast
 {
 	public class RemoveGotos: DepthFirstAstVisitor<object, object>
 	{
-		Stack<ForStatement> enteredLoops = new Stack<ForStatement>();
+		Stack<Statement> enteredLoops = new Stack<Statement>();
 		
-		ForStatement CurrentLoop {
+		Statement CurrentLoop {
 			get {
 				if (enteredLoops.Count > 0) {
 					return enteredLoops.Peek();
@@ -19,6 +19,20 @@ namespace Decompiler.Transforms.Ast
 			}
 		}
 		
+		Statement CurrentLoopBody {
+			get {
+				if (this.CurrentLoop == null) {
+					return null;
+				} else if (this.CurrentLoop is ForStatement) {
+					return ((ForStatement)this.CurrentLoop).EmbeddedStatement;
+				} else if (this.CurrentLoop is WhileStatement) {
+					return ((WhileStatement)this.CurrentLoop).EmbeddedStatement;
+				} else {
+					return null;
+				}
+			}
+		}
+		
 		public override object VisitForStatement(ForStatement forStatement, object data)
 		{
 			enteredLoops.Push(forStatement);
@@ -27,13 +41,13 @@ namespace Decompiler.Transforms.Ast
 			return null;
 		}
 		
-//		public override object VisitWhileStatement(WhileStatement whileStatement, object data)
-//		{
-//			enteredLoops.Push(whileStatement);
-//			base.VisitWhileStatement(whileStatement, data);
-//			enteredLoops.Pop();
-//			return null;
-//		}
+		public override object VisitWhileStatement(WhileStatement whileStatement, object data)
+		{
+			enteredLoops.Push(whileStatement);
+			base.VisitWhileStatement(whileStatement, data);
+			enteredLoops.Pop();
+			return null;
+		}
 		
 		public override object VisitBlockStatement(BlockStatement blockStatement, object data)
 		{
@@ -159,9 +173,19 @@ namespace Decompiler.Transforms.Ast
 			// Replace goto with 'continue'
 			// Continue statement which moves at the very end of loop
 			if (CurrentLoop != null &&
-			    (CurrentLoop.EmbeddedStatement is BlockStatement) &&
-			    ((CurrentLoop.EmbeddedStatement as BlockStatement).LastChild as LabelStatement) != null &&
-			    ((CurrentLoop.EmbeddedStatement as BlockStatement).LastChild as LabelStatement).Label == gotoStatement.Label) {
+			    (CurrentLoopBody is BlockStatement) &&
+			    ((CurrentLoopBody as BlockStatement).LastChild as LabelStatement) != null &&
+			    ((CurrentLoopBody as BlockStatement).LastChild as LabelStatement).Label == gotoStatement.Label) {
+				gotoStatement.ReplaceWith(new ContinueStatement());
+				return null;
+			}
+			
+			// Replace goto with 'continue'
+			// Jump before while
+			if (CurrentLoop is WhileStatement &&
+			    CurrentLoop.PrevSibling != null &&
+			    CurrentLoop.PrevSibling is LabelStatement &&
+			    (CurrentLoop.PrevSibling as LabelStatement).Label == gotoStatement.Label) {
 				gotoStatement.ReplaceWith(new ContinueStatement());
 				return null;
 			}

ICSharpCode.Decompiler/ILAst/ILAstOptimizer.cs

@@ -32,12 +32,13 @@ namespace Decompiler.ControlFlow
 			}
 			
 			if (abortBeforeStep == ILAstOptimizationStep.FindLoops) return;
+			UpdateLabelRefCounts(method);
 			foreach(ILBlock block in method.GetSelfAndChildrenRecursive<ILBlock>().ToList()) {
 				ControlFlowGraph graph;
 				graph = BuildGraph(block.Body, (ILLabel)block.EntryGoto.Operand);
 				graph.ComputeDominance();
 				graph.ComputeDominanceFrontier();
-				block.Body = FindLoops(new HashSet<ControlFlowNode>(graph.Nodes.Skip(3)), graph.EntryPoint, true);
+				block.Body = FindLoops(new HashSet<ControlFlowNode>(graph.Nodes.Skip(3)), graph.EntryPoint, false);
 			}
 			
 			if (abortBeforeStep == ILAstOptimizationStep.FindConditions) return;
@@ -45,6 +46,9 @@ namespace Decompiler.ControlFlow
 			foreach(ILBlock block in method.GetSelfAndChildrenRecursive<ILBlock>().ToList()) {
 				ControlFlowGraph graph;
 				graph = BuildGraph(block.Body, (ILLabel)block.EntryGoto.Operand);
+				// TODO: Fix
+				if (graph == null)
+					continue;
 				graph.ComputeDominance();
 				graph.ComputeDominanceFrontier();
 				block.Body = FindConditions(new HashSet<ControlFlowNode>(graph.Nodes.Skip(3)), graph.EntryPoint);
@@ -94,11 +98,11 @@ namespace Decompiler.ControlFlow
 					bool added = false;
 					
 					// Insert split
-					if ((currNode is ILLabel && !(lastNode is ILLabel)) ||
+					if (currNode is ILLabel ||
 						lastNode is ILTryCatchBlock ||
 						currNode is ILTryCatchBlock ||
 					    (lastNode is ILExpression) && ((ILExpression)lastNode).IsBranch() ||
-					    (currNode is ILExpression) && ((ILExpression)currNode).IsBranch())
+					    (currNode is ILExpression) && (((ILExpression)currNode).IsBranch() && basicBlock.Body.Count > 0))
 					{
 						ILBasicBlock lastBlock = basicBlock;
 						basicBlock = new ILBasicBlock();
@@ -140,7 +144,7 @@ namespace Decompiler.ControlFlow
 			
 			// Create graph nodes
 			labelToCfNode = new Dictionary<ILLabel, ControlFlowNode>();
-			Dictionary<ILNode, ControlFlowNode>  astNodeToCfNode = new Dictionary<ILNode, ControlFlowNode>();
+			Dictionary<ILNode, ControlFlowNode> astNodeToCfNode = new Dictionary<ILNode, ControlFlowNode>();
 			foreach(ILNode node in nodes) {
 				ControlFlowNode cfNode = new ControlFlowNode(index++, -1, ControlFlowNodeType.Normal);
 				cfNodes.Add(cfNode);
@@ -153,6 +157,9 @@ namespace Decompiler.ControlFlow
 				}
 			}
 			
+			if (!labelToCfNode.ContainsKey(entryLabel))
+				return null;
+			
 			// Entry endge
 			ControlFlowNode entryNode = labelToCfNode[entryLabel];
 			ControlFlowEdge entryEdge = new ControlFlowEdge(entryPoint, entryNode, JumpType.Normal);
@@ -196,18 +203,32 @@ namespace Decompiler.ControlFlow
 			    		&& node.DominanceFrontier.Contains(node)
 			    		&& (node != entryPoint || !excludeEntryPoint))
 				{
-					HashSet<ControlFlowNode> loopContents = new HashSet<ControlFlowNode>();
-					FindLoopContents(scope, loopContents, node, node);
+					HashSet<ControlFlowNode> loopContents = FindDominatedNodes(scope, node);
+					
+					ILWhileLoop loop = new ILWhileLoop();
+					
+					ILCondition cond;
+					HashSet<ControlFlowNode> condNodes;
+					ILLabel condLabel;
+					if (TryMatchCondition(loopContents, new ControlFlowNode[]{}, node, out cond, out condNodes, out condLabel)) {
+						loopContents.ExceptWith(condNodes);
+						scope.ExceptWith(condNodes);
+						// Use loop to implement condition
+						loop.Condition      = cond.Condition;
+						loop.PreLoopLabel   = condLabel;
+						loop.PostLoopGoto   = cond.FalseBlock.EntryGoto;
+						loop.BodyBlock      = new ILBlock() { EntryGoto = cond.TrueBlock.EntryGoto };
+					} else {
+						// Give the block some explicit entry point
+						ILLabel entryLabel  = new ILLabel() { Name = "Loop_" + (nextBlockIndex++) };
+						loop.BodyBlock      = new ILBlock() { EntryGoto = new ILExpression(ILCode.Br, entryLabel) };
+						((ILBasicBlock)node.UserData).Body.Insert(0, entryLabel);
+					}
+					loop.BodyBlock.Body = FindLoops(loopContents, node, true);
 					
 					// Move the content into loop block
 					scope.ExceptWith(loopContents);
-					ILLabel entryLabel = new ILLabel() { Name = "Loop_" + (nextBlockIndex++) };
-					((ILBasicBlock)node.UserData).Body.Insert(0, entryLabel);
-					result.Add(new ILLoop() {
-						ContentBlock = new ILBlock(FindLoops(loopContents, node, true)) {
- 							EntryGoto = new ILExpression(ILCode.Br, entryLabel)
-					    }
-					});
+					result.Add(loop);
 				}
 
 				// Using the dominator tree should ensure we find the the widest loop first
@@ -224,22 +245,12 @@ namespace Decompiler.ControlFlow
 			return result;
 		}
 		
-		static void FindLoopContents(HashSet<ControlFlowNode> scope, HashSet<ControlFlowNode> loopContents, ControlFlowNode loopHead, ControlFlowNode addNode)
-		{
-			if (scope.Contains(addNode) && loopHead.Dominates(addNode) && loopContents.Add(addNode)) {
-				foreach (var edge in addNode.Incoming) {
-					FindLoopContents(scope, loopContents, loopHead, edge.Source);
-				}
-			}
-		}
-		
 		void UpdateLabelRefCounts(ILBlock method)
 		{
 			labelRefCount = new Dictionary<ILLabel, int>();
-			foreach(ILLabel label in method.GetSelfAndChildrenRecursive<ILLabel>()) {
-				labelRefCount[label] = 0;
-			}
 			foreach(ILLabel target in method.GetSelfAndChildrenRecursive<ILExpression>().SelectMany(e => e.GetBranchTargets())) {
+				if (!labelRefCount.ContainsKey(target))
+					labelRefCount[target] = 0;
 				labelRefCount[target]++;
 			}
 		}

ICSharpCode.Decompiler/ILAst/ILAstTypes.cs

@@ -326,23 +326,40 @@ namespace Decompiler
 		}
 	}
 	
-	public class ILLoop : ILNode
+	public class ILWhileLoop : ILNode
 	{
-		public ILBlock ContentBlock;
+		public ILLabel      PreLoopLabel;  // Label allowing to jump to condition
+		public ILExpression Condition;
+		public ILBlock      BodyBlock;     // BodyBlock.EntryGoto performs the goto for a met condition
+		public ILExpression PostLoopGoto;  // Performs the goto for a failed condition
 		
 		public override IEnumerable<ILNode> GetChildren()
 		{
-			if (this.ContentBlock != null)
-				yield return ContentBlock;
+			if (this.PreLoopLabel != null)
+				yield return this.PreLoopLabel;
+			if (this.Condition != null)
+				yield return this.Condition;
+			if (this.BodyBlock != null)
+				yield return this.BodyBlock;
+			if (this.PostLoopGoto != null)
+				yield return this.PostLoopGoto;
 		}
 		
 		public override void WriteTo(ITextOutput output)
 		{
-			output.WriteLine("loop {");
+			if (this.PreLoopLabel != null)
+				this.PreLoopLabel.WriteTo(output);
+			output.WriteLine("");
+			output.Write("loop (");
+			if (this.Condition != null)
+				this.Condition.WriteTo(output);
+			output.WriteLine(") {");
 			output.Indent();
-			ContentBlock.WriteTo(output);
+			this.BodyBlock.WriteTo(output);
 			output.Unindent();
 			output.WriteLine("}");
+			if (this.PostLoopGoto != null)
+				this.PostLoopGoto.WriteTo(output);
 		}
 	}
 	

NRefactory/ICSharpCode.NRefactory/CSharp/OutputVisitor/OutputVisitor.cs

@@ -382,6 +382,8 @@ namespace ICSharpCode.NRefactory.CSharp
 		
 		void WriteEmbeddedStatement(Statement embeddedStatement)
 		{
+			if (embeddedStatement.IsNull)
+				return;
 			BlockStatement block = embeddedStatement as BlockStatement;
 			if (block != null)
 				VisitBlockStatement(block, null);