From ae9665238615f49704290d948560e61259856d19 Mon Sep 17 00:00:00 2001
From: atrexus <110138304+atrexus@users.noreply.github.com>
Date: Tue, 13 Aug 2024 14:54:58 -0400
Subject: [PATCH] Clean up control flow generation

---
 .../IR/ProtoTypes/ControlFlow/BasicBlock.cs   |  6 ++--
 src/Unluau/IR/ProtoTypes/ControlFlow/Graph.cs | 29 +++++++++----------
 src/Unluau/IR/Writers/DotWriter.cs            |  6 ++--
 3 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/src/Unluau/IR/ProtoTypes/ControlFlow/BasicBlock.cs b/src/Unluau/IR/ProtoTypes/ControlFlow/BasicBlock.cs
index 6fa9d3b..9f2bbcf 100644
--- 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
index fb8a758..afab6ae 100644
--- 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
index 35f41b7..0d4e833 100644
--- 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);