Clean up control flow generation

atrexus · Aug 13, 2024 · ae96652 · ae96652
1 parent 2224b82
commit ae96652
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 20 deletions.
diff --git a/src/Unluau/IR/ProtoTypes/ControlFlow/BasicBlock.cs b/src/Unluau/IR/ProtoTypes/ControlFlow/BasicBlock.cs
@@ -20,7 +20,7 @@ public enum BranchType
         /// <summary>
         /// This block never branches.
         /// </summary>
-        Never
+        Never,
     }
 
     /// <summary>
@@ -31,7 +31,7 @@ public class BasicBlock : Node
         /// <summary>
         /// The kind of branch that is taken.
         /// </summary>
-        public BranchType Branch { get; set; } = BranchType.Always;
+        public BranchType? Branch { get; set; }
 
         /// <summary>
         /// The list of instructions in the basic block.
@@ -45,6 +45,8 @@ public class BasicBlock : Node
 
         public void AddEdge(BasicBlock target) => OutgoingEdges.Add(new Edge(this, target));
 
+        public bool IsDead => OutgoingEdges.Count == 0 && Branch == null;
+
         /// <inheritdoc/>
         public override void Accept(Visitor visitor)
         {

diff --git a/src/Unluau/IR/ProtoTypes/ControlFlow/Graph.cs b/src/Unluau/IR/ProtoTypes/ControlFlow/Graph.cs
@@ -48,7 +48,7 @@ private void AddInstruction(Instruction instruction)
 
                 if (instruction.Context.Pc == pc + 1)
                 {
-                    if (_currentBlock == oldBlock)
+                    if (_currentBlock == oldBlock && _currentBlock.Instructions.Count > 0)
                         _currentBlock = new();
 
                     block.AddEdge(_currentBlock);
@@ -59,6 +59,9 @@ private void AddInstruction(Instruction instruction)
 
             _currentBlock.Instructions.Add(instruction);
 
+            if (_currentBlock != oldBlock && oldBlock.IsDead)
+                oldBlock.AddEdge(_currentBlock);
+
             switch (instruction.Code)
             {
                 case OpCode.JumpX:
@@ -68,8 +71,6 @@ private void AddInstruction(Instruction instruction)
                 }
                 case OpCode.LoadB:
                 {
-                    var branchType = instruction.C > 0 ? BranchType.Always : BranchType.Never;
-
                     AddEdge(instruction.C, instruction.Context.Pc, _currentBlock, BranchType.Always);
                     break;
                 }
@@ -79,21 +80,25 @@ private void AddInstruction(Instruction instruction)
                     AddEdge(instruction.D, instruction.Context.Pc, _currentBlock, BranchType.Always);
                     break;
                 }
+                case OpCode.ForNPrep:
+                {
+                    AddEdge(0, instruction.Context.Pc, _currentBlock, BranchType.Always);
+                    break;
+                }
                 case OpCode.Return:
                 {
                     AddEdge(0, instruction.Context.Pc, _currentBlock, BranchType.Never);
-                    goto default;
+                    break;
                 }
                 case OpCode.JumpXEqKNil:
                 case OpCode.JumpXEqKB:
                 case OpCode.JumpXEqKN:
                 case OpCode.JumpXEqKS:
                 {
-                    var branchType = instruction.D > 0 ? BranchType.Can : BranchType.Never;
-
-                    AddEdge(instruction.D == 1 ? 0 : instruction.D, instruction.Context.Pc, _currentBlock, branchType);
+                    AddEdge(instruction.D == 1 ? 0 : instruction.D, instruction.Context.Pc, _currentBlock, BranchType.Can);
                     break;
                 }
+                case OpCode.ForNLoop:
                 case OpCode.JumpIf:
                 case OpCode.JumpIfNot:
                 case OpCode.JumpIfEq:
@@ -103,15 +108,7 @@ private void AddInstruction(Instruction instruction)
                 case OpCode.JumpIfLt:
                 case OpCode.JumpIfNotLt:
                 {
-                    var branchType = instruction.D > 0 ? BranchType.Can : BranchType.Never;
-
-                    AddEdge(instruction.D, instruction.Context.Pc, _currentBlock, branchType);
-                    break;
-                }
-                default:
-                {
-                    if (_currentBlock != oldBlock)
-                        oldBlock.AddEdge(_currentBlock);
+                    AddEdge(instruction.D, instruction.Context.Pc, _currentBlock, BranchType.Can);
                     break;
                 }
             }

diff --git a/src/Unluau/IR/Writers/DotWriter.cs b/src/Unluau/IR/Writers/DotWriter.cs
@@ -73,7 +73,7 @@ public override bool Visit(ProtoType protoType)
             _graph!.Add(_protoSubGraph);
 
             // Visit all basic blocks in the control flow graph.
-            protoType.ControlFlow.Accept(this);
+            protoType.ControlFlow!.Accept(this);
 
             return false;
         }
@@ -144,7 +144,9 @@ public override bool Visit(Edge edge)
                 .From($"block_{edge.Source.GetHashCode()}").To($"block_{edge.Target.GetHashCode()}")
                 .WithArrowHead(DotEdgeArrowType.Vee)
                 .WithArrowTail(DotEdgeArrowType.None)
-                .WithAttribute("fontname", "Helvetica");
+                .WithAttribute("fontname", "Helvetica")
+                .WithAttribute("headport", "n")
+                .WithAttribute("tailport", "s");
 
             if (_edgeName != null)
                 myEdge = myEdge.WithLabel(_edgeName);