From 2e56d20724bae2f601f7ecc26e49c6c75012aade Mon Sep 17 00:00:00 2001
From: Will Tebbutt <wct23@cam.ac.uk>
Date: Thu, 16 May 2024 09:24:51 +0100
Subject: [PATCH] Compute CFG (#158)

* Compute CFG

* Fix typo

* Bump patch

* Remove dead comment
---
 Project.toml               |  2 +-
 src/interpreter/bbcode.jl  | 33 ++++++++++++++++++++++++++++++++-
 test/interpreter/bbcode.jl | 10 ++++++++++
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index 02aa12c3..234539c9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Tapir"
 uuid = "07d77754-e150-4737-8c94-cd238a1fb45b"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.2.14"
+version = "0.2.15"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
diff --git a/src/interpreter/bbcode.jl b/src/interpreter/bbcode.jl
index bb7402bf..383d7990 100644
--- a/src/interpreter/bbcode.jl
+++ b/src/interpreter/bbcode.jl
@@ -342,6 +342,37 @@ end
 concatenate_ids(bb_code::BBCode) = reduce(vcat, map(b -> b.inst_ids, bb_code.blocks))
 concatenate_stmts(bb_code::BBCode) = reduce(vcat, map(b -> b.insts, bb_code.blocks))
 
+"""
+    control_flow_graph(bb_code::BBCode)::Core.Compiler.CFG
+
+Computes the `Core.Compiler.CFG` object associated to this `bb_code`.
+"""
+control_flow_graph(bb_code::BBCode)::Core.Compiler.CFG = _control_flow_graph(bb_code.blocks)
+
+# Internal function, used to implement control_flow_graph, but easier to write test cases
+# for because there is no need to construct an ensure BBCode object.
+function _control_flow_graph(blks::Vector{BBlock})::Core.Compiler.CFG
+
+    # Get IDs of predecessors and successors.
+    preds_ids = _compute_all_predecessors(blks)
+    succs_ids = _compute_all_successors(blks)
+
+    # Construct map from block ID to block number.
+    block_ids = map(b -> b.id, blks)
+    id_to_num = Dict{ID, Int}(zip(block_ids, collect(eachindex(block_ids))))
+
+    # Convert predecessor and successor IDs to numbers.
+    preds = map(id -> sort(map(p -> id_to_num[p], preds_ids[id])), block_ids)
+    succs = map(id -> sort(map(s -> id_to_num[s], succs_ids[id])), block_ids)
+
+    index = vcat(0, cumsum(map(length, blks))) .+ 1
+    basic_blocks = map(eachindex(blks)) do n
+        stmt_range = Core.Compiler.StmtRange(index[n], index[n+1] - 1)
+        return Core.Compiler.BasicBlock(stmt_range, preds[n], succs[n])
+    end
+    return Core.Compiler.CFG(basic_blocks, index[2:end-1])
+end
+
 #
 # Converting from IRCode to BBCode
 #
@@ -454,7 +485,7 @@ function CC.IRCode(bb_code::BBCode)
     bb_code = _lower_switch_statements(bb_code)
     bb_code = _remove_double_edges(bb_code)
     insts = _ids_to_line_positions(bb_code)
-    cfg = _compute_basic_blocks(insts)
+    cfg =  control_flow_graph(bb_code)
     insts = _lines_to_blocks(insts, cfg)
     return IRCode(
         CC.InstructionStream(
diff --git a/test/interpreter/bbcode.jl b/test/interpreter/bbcode.jl
index a54ffe95..95876eb8 100644
--- a/test/interpreter/bbcode.jl
+++ b/test/interpreter/bbcode.jl
@@ -53,6 +53,16 @@ end
         @test length(Tapir.collect_stmts(bb_code)) == length(ir.stmts.inst)
         @test Tapir.id_to_line_map(bb_code) isa Dict{ID, Int}
     end
+    @testset "control_flow_graph" begin
+        ir = Base.code_ircode_by_type(Tuple{typeof(sin), Float64})[1][1]
+        bb = BBCode(ir)
+        new_ir = Core.Compiler.IRCode(bb)
+        cfg = Tapir.control_flow_graph(bb)
+        @test all(map((l, r) -> l.stmts == r.stmts, ir.cfg.blocks, cfg.blocks))
+        @test all(map((l, r) -> sort(l.preds) == sort(r.preds), ir.cfg.blocks, cfg.blocks))
+        @test all(map((l, r) -> sort(l.succs) == sort(r.succs), ir.cfg.blocks, cfg.blocks))
+        @test ir.cfg.index == cfg.index
+    end
     @testset "_characterise_unique_predecessor_blocks" begin
         @testset "single block" begin
             blk_id = ID()