diff --git a/src/ControlFlow/Node.cs b/src/ControlFlow/Node.cs index 69a81c501..d65b6b601 100644 --- a/src/ControlFlow/Node.cs +++ b/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); diff --git a/src/ControlFlow/Nodes.cs b/src/ControlFlow/Nodes.cs index 25c17ccdb..15b623d67 100644 --- a/src/ControlFlow/Nodes.cs +++ b/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]; } }