Maintain node links at all levels, not just at the top one.

src/ControlFlow/Node.cs

@@ -19,6 +19,29 @@ namespace Decompiler.ControlFlow
 				base.InsertItem(index, item);
 			}
 		}
+		
+		public bool ContainsRecursive(Node node)
+		{
+			if (this.Contains(node)) {
+				return true;
+			}
+			foreach(Node item in this.Items) {
+				if (item.Childs.ContainsRecursive(node)) {
+					return true;
+				}
+			}
+			return false;
+		}
+		
+		public Node FindContainer(Node node)
+		{
+			foreach(Node item in this.Items) {
+				if (item == node || item.Childs.ContainsRecursive(node)) {
+					return item;
+				}
+			}
+			return null;
+		}
 	}
 	
 	public abstract class Node
@@ -83,6 +106,46 @@ namespace Decompiler.ControlFlow
 			this.id = NextNodeID++;
 		}
 		
+		void GetBasicBlockSuccessors(NodeCollection accumulator)
+		{
+			BasicBlock me = this as BasicBlock;
+			if (me != null) {
+				if (me.FallThroughBasicBlock != null) {
+					accumulator.Add(me.FallThroughBasicBlock);
+				}
+				if (me.BranchBasicBlock != null) {
+					accumulator.Add(me.BranchBasicBlock);
+				}
+			} else {
+				foreach(Node child in this.Childs) {
+					child.GetBasicBlockSuccessors(accumulator);
+				}
+			}
+		}
+		
+		public void RebuildNodeLinks()
+		{
+			foreach(Node child in this.Childs) {
+				NodeCollection successorBasicBlocks = new NodeCollection();
+				child.GetBasicBlockSuccessors(successorBasicBlocks);
+				NodeCollection successorNodes = new NodeCollection();
+				foreach(Node successorBasicBlock in successorBasicBlocks) {
+					Node container = this.Childs.FindContainer(successorBasicBlock);
+					if (container != null) {
+						successorNodes.Add(container);
+					}
+				}
+				// Remove self link
+				if (successorNodes.Contains(child)) {
+					successorNodes.Remove(child);
+				}
+				foreach(Node target in successorNodes) {
+					child.Successors.Add(target);
+					target.Predecessors.Add(child);
+				}
+			}
+		}
+		
 		public void Optimize()
 		{
 			for(int i = 0; i < this.Childs.Count;) {
@@ -165,6 +228,8 @@ namespace Decompiler.ControlFlow
 			Relink(head, mergedNode);
 			Relink(tail, mergedNode);
 			
+			mergedNode.RebuildNodeLinks();
+			
 			// Remove the old nodes and add the merged node - replace head with the merged node
 			container.Childs.Remove(tail);
 			int headIndex = container.Childs.IndexOf(head);

src/ControlFlow/Nodes.cs

@@ -7,7 +7,6 @@ namespace Decompiler.ControlFlow
 {
 	public class MethodBodyGraph: Node
 	{
-		// TODO: Add links between the generated BasicBlocks
 		public MethodBodyGraph(StackExpressionCollection exprs): base(null)
 		{
 			if (exprs.Count == 0) throw new ArgumentException("Count == 0", "exprs");
@@ -26,7 +25,7 @@ namespace Decompiler.ControlFlow
 				exprs[i].BasicBlock = basicBlock;
 			}
 			
-			// Add fall-through links
+			// Add fall-through links to BasicBlocks
 			for(int i = 0; i < exprs.Count - 1; i++) {
 				BasicBlock node = exprs[i].BasicBlock;
 				BasicBlock target = exprs[i + 1].BasicBlock;
@@ -38,22 +37,21 @@ namespace Decompiler.ControlFlow
 				if (exprs[i].LastByteCode.OpCode.Code == Code.Br) continue;
 				
 				node.FallThroughBasicBlock = target;
-				node.Successors.Add(target);
-				target.Predecessors.Add(node);
 			}
 			
-			// Add branch links
+			// Add branch links to BasicBlocks
 			for(int i = 0; i < exprs.Count; i++) {
 				if (exprs[i].BranchTarget != null) {
 					BasicBlock node = exprs[i].BasicBlock;
 					BasicBlock target = exprs[i].BranchTarget.BasicBlock;
 					
 					node.BranchBasicBlock = target;
-					node.Successors.Add(target);
-					target.Predecessors.Add(node);
 				}
 			}
 			
+			// Link all nodes
+			RebuildNodeLinks();
+			
 			this.HeadChild = this.Childs[0];
 		}
 	}