Put BBCode in module (#512)

* BBCode in module

* Formatting and BasicBlockCode module name

* Shuffle some names around

* Update src/interpreter/bbcode.jl

Co-authored-by: Hong Ge <[email protected]>
Signed-off-by: Will Tebbutt <[email protected]>

* Bump patch version

---------

Signed-off-by: Will Tebbutt <[email protected]>
Co-authored-by: Hong Ge <[email protected]>

Project.toml

@@ -1,7 +1,7 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.4.102"
+version = "0.4.103"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/Mooncake.jl

@@ -6,7 +6,6 @@ using ADTypes,
     ChainRules,
     DiffRules,
     ExprTools,
-    Graphs,
     InteractiveUtils,
     LinearAlgebra,
     MistyClosures,
@@ -105,11 +104,13 @@ include("codual.jl")
 include("debug_mode.jl")
 include("stack.jl")
 
+include(joinpath("interpreter", "bbcode.jl"))
+using .BasicBlockCode
+
 include(joinpath("interpreter", "contexts.jl"))
 include(joinpath("interpreter", "abstract_interpretation.jl"))
 include(joinpath("interpreter", "patch_for_319.jl"))
 include(joinpath("interpreter", "ir_utils.jl"))
-include(joinpath("interpreter", "bbcode.jl"))
 include(joinpath("interpreter", "ir_normalisation.jl"))
 include(joinpath("interpreter", "zero_like_rdata.jl"))
 include(joinpath("interpreter", "s2s_reverse_mode_ad.jl"))

src/interpreter/bbcode.jl

@@ -1,7 +1,72 @@
-# See the docstring for `BBCode` for some context on this file.
+"""
+    module BasicBlockCode
+
+See the docstring for the `BBCode` `struct` for info on this file.
+"""
+module BasicBlockCode
+
+using Graphs
+
+using Core.Compiler:
+    ReturnNode,
+    PhiNode,
+    GotoIfNot,
+    GotoNode,
+    NewInstruction,
+    IRCode,
+    SSAValue,
+    PiNode,
+    Argument
+const CC = Core.Compiler
+
+export ID,
+    seed_id!,
+    IDPhiNode,
+    IDGotoNode,
+    IDGotoIfNot,
+    Switch,
+    BBlock,
+    phi_nodes,
+    terminator,
+    insert_before_terminator!,
+    collect_stmts,
+    compute_all_predecessors,
+    BBCode,
+    remove_unreachable_blocks!,
+    characterise_used_ids,
+    characterise_unique_predecessor_blocks,
+    sort_blocks!,
+    InstVector,
+    IDInstPair,
+    __line_numbers_to_block_numbers!,
+    is_reachable_return_node,
+    new_inst
 
 const _id_count::Dict{Int,Int32} = Dict{Int,Int32}()
 
+"""
+    new_inst(stmt, type=Any, flag=CC.IR_FLAG_REFINED)::NewInstruction
+
+Create a `NewInstruction` with fields:
+- `stmt` = `stmt`
+- `type` = `type`
+- `info` = `CC.NoCallInfo()`
+- `line` = `Int32(1)`
+- `flag` = `flag`
+"""
+function new_inst(@nospecialize(stmt), @nospecialize(type)=Any, flag=CC.IR_FLAG_REFINED)
+    return NewInstruction(stmt, type, CC.NoCallInfo(), Int32(1), flag)
+end
+
+"""
+    const InstVector = Vector{NewInstruction}
+
+Note: the `CC.NewInstruction` type is used to represent instructions because it has the
+correct fields. While it is only used to represent new instrucdtions in `Core.Compiler`, it
+is used to represent all instructions in `BBCode`.
+"""
+const InstVector = Vector{NewInstruction}
+
 """
     ID()
 
@@ -394,6 +459,47 @@ function _control_flow_graph(blks::Vector{BBlock})::Core.Compiler.CFG
     return Core.Compiler.CFG(basic_blocks, index[2:(end - 1)])
 end
 
+"""
+    _instructions_to_blocks(insts::InstVector, cfg::CC.CFG)::InstVector
+
+Pulls out the instructions from `insts`, and calls `__line_numbers_to_block_numbers!`.
+"""
+function _lines_to_blocks(insts::InstVector, cfg::CC.CFG)::InstVector
+    stmts = __line_numbers_to_block_numbers!(Any[x.stmt for x in insts], cfg)
+    return map((inst, stmt) -> NewInstruction(inst; stmt), insts, stmts)
+end
+
+"""
+    __line_numbers_to_block_numbers!(insts::Vector{Any}, cfg::CC.CFG)
+
+Converts any edges in `GotoNode`s, `GotoIfNot`s, `PhiNode`s, and `:enter` expressions which
+refer to line numbers into references to block numbers. The `cfg` provides the information
+required to perform this conversion.
+
+For context, `CodeInfo` objects have references to line numbers, while `IRCode` uses
+block numbers.
+
+This code is copied over directly from the body of `Core.Compiler.inflate_ir!`.
+"""
+function __line_numbers_to_block_numbers!(insts::Vector{Any}, cfg::CC.CFG)
+    for i in eachindex(insts)
+        stmt = insts[i]
+        if isa(stmt, GotoNode)
+            insts[i] = GotoNode(CC.block_for_inst(cfg, stmt.label))
+        elseif isa(stmt, GotoIfNot)
+            insts[i] = GotoIfNot(stmt.cond, CC.block_for_inst(cfg, stmt.dest))
+        elseif isa(stmt, PhiNode)
+            insts[i] = PhiNode(
+                Int32[CC.block_for_inst(cfg, Int(edge)) for edge in stmt.edges], stmt.values
+            )
+        elseif Meta.isexpr(stmt, :enter)
+            stmt.args[1] = CC.block_for_inst(cfg, stmt.args[1]::Int)
+            insts[i] = stmt
+        end
+    end
+    return insts
+end
+
 #
 # Converting from IRCode to BBCode
 #
@@ -431,7 +537,8 @@ end
 Convert an `Compiler.InstructionStream` into a list of `Compiler.NewInstruction`s.
 """
 function new_inst_vec(x::CC.InstructionStream)
-    return map((v...,) -> NewInstruction(v...), stmt(x), x.type, x.info, x.line, x.flag)
+    stmt = @static VERSION < v"1.11.0-rc4" ? x.inst : x.stmt
+    return map((v...,) -> NewInstruction(v...), stmt, x.type, x.info, x.line, x.flag)
 end
 
 # Maps from positional names (SSAValues for nodes, Integers for basic blocks) to IDs.
@@ -673,7 +780,7 @@ function _distance_to_entry(blks::Vector{BBlock})::Vector{Int}
 end
 
 """
-    _sort_blocks!(ir::BBCode)::BBCode
+    sort_blocks!(ir::BBCode)::BBCode
 
 Ensure that blocks appear in order of distance-from-entry-point, where distance the
 distance from block b to the entry point is defined to be the minimum number of basic
@@ -687,7 +794,7 @@ WARNING: use with care. Only use if you are confident that arbitrary re-ordering
 blocks in `ir` is valid. Notably, this does not hold if you have any `IDGotoIfNot` nodes in
 `ir`.
 """
-function _sort_blocks!(ir::BBCode)::BBCode
+function sort_blocks!(ir::BBCode)::BBCode
     I = sortperm(_distance_to_entry(ir.blocks))
     ir.blocks .= ir.blocks[I]
     return ir
@@ -765,6 +872,15 @@ function characterise_unique_predecessor_blocks(
     return is_unique_pred, pred_is_unique_pred
 end
 
+"""
+    is_reachable_return_node(x::ReturnNode)
+
+Determine whether `x` is a `ReturnNode`, and if it is, if it is also reachable. This is
+purely a function of whether or not its `val` field is defined or not.
+"""
+is_reachable_return_node(x::ReturnNode) = isdefined(x, :val)
+is_reachable_return_node(x) = false
+
 """
     characterise_used_ids(stmts::Vector{IDInstPair})::Dict{ID, Bool}
 
@@ -882,26 +998,4 @@ function _remove_unreachable_blocks!(blks::Vector{BBlock})
     return remaining_blks
 end
 
-"""
-    inc_args(stmt)
-
-Increment by `1` the `n` field of any `Argument`s present in `stmt`.
-"""
-inc_args(x::Expr) = Expr(x.head, map(__inc, x.args)...)
-inc_args(x::ReturnNode) = isdefined(x, :val) ? ReturnNode(__inc(x.val)) : x
-inc_args(x::IDGotoIfNot) = IDGotoIfNot(__inc(x.cond), x.dest)
-inc_args(x::IDGotoNode) = x
-function inc_args(x::IDPhiNode)
-    new_values = Vector{Any}(undef, length(x.values))
-    for n in eachindex(x.values)
-        if isassigned(x.values, n)
-            new_values[n] = __inc(x.values[n])
-        end
-    end
-    return IDPhiNode(x.edges, new_values)
 end
-inc_args(::Nothing) = nothing
-inc_args(x::GlobalRef) = x
-
-__inc(x::Argument) = Argument(x.n + 1)
-__inc(x) = x

src/interpreter/ir_utils.jl

@@ -1,12 +1,3 @@
-"""
-    const InstVector = Vector{NewInstruction}
-
-Note: the `CC.NewInstruction` type is used to represent instructions because it has the
-correct fields. While it is only used to represent new instrucdtions in `Core.Compiler`, it
-is used to represent all instructions in `BBCode`.
-"""
-const InstVector = Vector{NewInstruction}
-
 """
     stmt(ir::CC.InstructionStream)
 
@@ -44,47 +35,6 @@ function ircode(
     return CC.IRCode(stmts, cfg, linetable, argtypes, meta, CC.VarState[])
 end
 
-"""
-    __line_numbers_to_block_numbers!(insts::Vector{Any}, cfg::CC.CFG)
-
-Converts any edges in `GotoNode`s, `GotoIfNot`s, `PhiNode`s, and `:enter` expressions which
-refer to line numbers into references to block numbers. The `cfg` provides the information
-required to perform this conversion.
-
-For context, `CodeInfo` objects have references to line numbers, while `IRCode` uses
-block numbers.
-
-This code is copied over directly from the body of `Core.Compiler.inflate_ir!`.
-"""
-function __line_numbers_to_block_numbers!(insts::Vector{Any}, cfg::CC.CFG)
-    for i in eachindex(insts)
-        stmt = insts[i]
-        if isa(stmt, GotoNode)
-            insts[i] = GotoNode(CC.block_for_inst(cfg, stmt.label))
-        elseif isa(stmt, GotoIfNot)
-            insts[i] = GotoIfNot(stmt.cond, CC.block_for_inst(cfg, stmt.dest))
-        elseif isa(stmt, PhiNode)
-            insts[i] = PhiNode(
-                Int32[CC.block_for_inst(cfg, Int(edge)) for edge in stmt.edges], stmt.values
-            )
-        elseif Meta.isexpr(stmt, :enter)
-            stmt.args[1] = CC.block_for_inst(cfg, stmt.args[1]::Int)
-            insts[i] = stmt
-        end
-    end
-    return insts
-end
-
-"""
-    _instructions_to_blocks(insts::InstVector, cfg::CC.CFG)::InstVector
-
-Pulls out the instructions from `insts`, and calls `__line_numbers_to_block_numbers!`.
-"""
-function _lines_to_blocks(insts::InstVector, cfg::CC.CFG)::InstVector
-    stmts = __line_numbers_to_block_numbers!(Any[x.stmt for x in insts], cfg)
-    return map((inst, stmt) -> NewInstruction(inst; stmt), insts, stmts)
-end
-
 """
     __insts_to_instruction_stream(insts::Vector{Any})
 
@@ -270,15 +220,6 @@ function lookup_ir(
     return CC.typeinf_ircode(interp, mi.def, mi.specTypes, mi.sparam_vals, optimize_until)
 end
 
-"""
-    is_reachable_return_node(x::ReturnNode)
-
-Determine whether `x` is a `ReturnNode`, and if it is, if it is also reachable. This is
-purely a function of whether or not its `val` field is defined or not.
-"""
-is_reachable_return_node(x::ReturnNode) = isdefined(x, :val)
-is_reachable_return_node(x) = false
-
 """
     is_unreachable_return_node(x::ReturnNode)
 
@@ -305,20 +246,6 @@ Throw an `UnhandledLanguageFeatureException` with message `msg`.
 """
 unhandled_feature(msg::String) = throw(UnhandledLanguageFeatureException(msg))
 
-"""
-    new_inst(stmt, type=Any, flag=CC.IR_FLAG_REFINED)::NewInstruction
-
-Create a `NewInstruction` with fields:
-- `stmt` = `stmt`
-- `type` = `type`
-- `info` = `CC.NoCallInfo()`
-- `line` = `Int32(1)`
-- `flag` = `flag`
-"""
-function new_inst(@nospecialize(stmt), @nospecialize(type)=Any, flag=CC.IR_FLAG_REFINED)
-    return NewInstruction(stmt, type, CC.NoCallInfo(), Int32(1), flag)
-end
-
 """
     replace_uses_with!(stmt, def::Union{Argument, SSAValue}, val)
 

src/interpreter/s2s_reverse_mode_ad.jl

@@ -353,6 +353,31 @@ function comms_channel(info::ADStmtInfo)
     return only(filter(x -> x[1] == info.comms_id, info.fwds))
 end
 
+"""
+    inc_args(stmt)
+
+Increment by `1` the `n` field of any `Argument`s present in `stmt`.
+Used in `make_ad_stmts!`.
+"""
+inc_args(x::Expr) = Expr(x.head, map(__inc, x.args)...)
+inc_args(x::ReturnNode) = isdefined(x, :val) ? ReturnNode(__inc(x.val)) : x
+inc_args(x::IDGotoIfNot) = IDGotoIfNot(__inc(x.cond), x.dest)
+inc_args(x::IDGotoNode) = x
+function inc_args(x::IDPhiNode)
+    new_values = Vector{Any}(undef, length(x.values))
+    for n in eachindex(x.values)
+        if isassigned(x.values, n)
+            new_values[n] = __inc(x.values[n])
+        end
+    end
+    return IDPhiNode(x.edges, new_values)
+end
+inc_args(::Nothing) = nothing
+inc_args(x::GlobalRef) = x
+
+__inc(x::Argument) = Argument(x.n + 1)
+__inc(x) = x
+
 """
     make_ad_stmts!(inst::NewInstruction, line::ID, info::ADInfo)::ADStmtInfo
 
@@ -1489,7 +1514,7 @@ function pullback_ir(
     # avoid annoying the Julia compiler.
     blks = vcat(entry_block, main_blocks, exit_block)
     pb_ir = BBCode(blks, arg_types, ir.sptypes, ir.linetable, ir.meta)
-    return remove_unreachable_blocks!(_sort_blocks!(pb_ir))
+    return remove_unreachable_blocks!(sort_blocks!(pb_ir))
 end
 
 """