1 files changed, 54 insertions, 38 deletions
diff --git a/ARMeilleure/IntermediateRepresentation/BasicBlock.cs b/ARMeilleure/IntermediateRepresentation/BasicBlock.cs
index ac48ac8e..640978fe 100644
--- a/ARMeilleure/IntermediateRepresentation/BasicBlock.cs
+++ b/ARMeilleure/IntermediateRepresentation/BasicBlock.cs
@@ -1,9 +1,12 @@
+using System;
 using System.Collections.Generic;
 
 namespace ARMeilleure.IntermediateRepresentation
 {
     class BasicBlock : IIntrusiveListNode<BasicBlock>
     {
+        private readonly List<BasicBlock> _successors = new List<BasicBlock>();
+
         public int Index { get; set; }
 
         public BasicBlock ListPrevious { get; set; }
@@ -11,69 +14,82 @@ namespace ARMeilleure.IntermediateRepresentation
 
         public IntrusiveList<Node> Operations { get; }
 
-        private BasicBlock _next;
-        private BasicBlock _branch;
-
-        public BasicBlock Next
-        {
-            get => _next;
-            set => _next = AddSuccessor(_next, value);
-        }
-
-        public BasicBlock Branch
-        {
-            get => _branch;
-            set => _branch = AddSuccessor(_branch, value);
-        }
-
         public List<BasicBlock> Predecessors { get; }
 
         public HashSet<BasicBlock> DominanceFrontiers { get; }
-
         public BasicBlock ImmediateDominator { get; set; }
 
-        public BasicBlock()
+        public int SuccessorCount => _successors.Count;
+
+        public BasicBlock() : this(index: -1) { }
+
+        public BasicBlock(int index)
         {
             Operations = new IntrusiveList<Node>();
-
             Predecessors = new List<BasicBlock>();
-
             DominanceFrontiers = new HashSet<BasicBlock>();
 
-            Index = -1;
+            Index = index;
         }
 
-        public BasicBlock(int index) : this()
+        public void AddSuccessor(BasicBlock block)
         {
-            Index = index;
+            if (block == null)
+            {
+                throw new ArgumentNullException(nameof(block));
+            }
+
+            block.Predecessors.Add(this);
+
+            _successors.Add(block);
         }
 
-        private BasicBlock AddSuccessor(BasicBlock oldBlock, BasicBlock newBlock)
+        public void RemoveSuccessor(int index)
         {
-            oldBlock?.Predecessors.Remove(this);
-            newBlock?.Predecessors.Add(this);
+            BasicBlock oldBlock = _successors[index];
 
-            return newBlock;
+            oldBlock.Predecessors.Remove(this);
+
+            _successors.RemoveAt(index);
         }
 
-        public void Append(Node node)
+        public BasicBlock GetSuccessor(int index)
         {
-            // If the branch block is not null, then the list of operations
-            // should end with a branch instruction. We insert the new operation
-            // before this branch.
-            if (_branch != null || (Operations.Last != null && IsLeafBlock()))
-            {
-                Operations.AddBefore(Operations.Last, node);
-            }
-            else
+            return _successors[index];
+        }
+
+        public void SetSuccessor(int index, BasicBlock block)
+        {
+            if (block == null)
             {
-                Operations.AddLast(node);
+                throw new ArgumentNullException(nameof(block));
             }
+
+            BasicBlock oldBlock = _successors[index];
+
+            oldBlock.Predecessors.Remove(this);
+            block.Predecessors.Add(this);
+
+            _successors[index] = block;
         }
 
-        private bool IsLeafBlock()
+        public void Append(Node node)
         {
-            return _branch == null && _next == null;
+            var lastOp = Operations.Last as Operation;
+
+            // Append node before terminal or to end if no terminal.
+            switch (lastOp?.Instruction)
+            {
+                case Instruction.Return:
+                case Instruction.Tailcall:
+                case Instruction.BranchIf:
+                    Operations.AddBefore(lastOp, node);
+                    break;
+
+                default:
+                    Operations.AddLast(node);
+                    break;
+            }
         }
 
         public Node GetLastOp()