further parametrize IRShow (#40459)

Ref JuliaDebug/Cthulhu.jl#149. This allows
Ctulhu to further customize how line annotations are printed. Also moves
a hook to skip printing some statements into `show_ir`, which previously
Cthulhu had to duplicate. Since Cthulhu seems to be the only downstream
consumer of `show_ir` https://juliahub.com/ui/RepoSearch?q=show_ir&r=true,
I don't expect this change to be major. I did also try to keep
compatibility for `show_ir_stmt`, which seems to be more commonly used.

Author: simeonschaub

base/compiler/ssair/show.jl

@@ -487,6 +487,25 @@ function DILineInfoPrinter(linetable::Vector, showtypes::Bool=false)
     return emit_lineinfo_update
 end
 
+# line_info_preprinter(io::IO, indent::String, idx::Int) may print relevant info
+#   at the beginning of the line, and should at least print `indent`. It returns a
+#   string that will be printed after the final basic-block annotation.
+# line_info_postprinter(io::IO, typ, used::Bool) prints the type-annotation at the end
+#   of the statement
+# should_print_stmt(idx::Int) -> Bool: whether the statement at index `idx` should be
+#   printed as part of the IR or not
+# bb_color: color used for printing the basic block brackets on the left
+struct IRShowConfig
+    line_info_preprinter
+    line_info_postprinter
+    should_print_stmt
+    bb_color::Symbol
+    function IRShowConfig(line_info_preprinter, line_info_postprinter=default_expr_type_printer;
+                          should_print_stmt=Returns(true), bb_color::Symbol=:light_black)
+        return new(line_info_preprinter, line_info_postprinter, should_print_stmt, bb_color)
+    end
+end
+
 struct _UNDEF
     global const UNDEF = _UNDEF.instance
 end
@@ -510,7 +529,7 @@ function _type(code::CodeInfo, idx::Int)
     return isassigned(types, idx) ? types[idx] : UNDEF
 end
 
-function statement_indices_to_labels(@nospecialize(stmt), cfg::CFG)
+function statement_indices_to_labels(stmt, cfg::CFG)
     # convert statement index to labels, as expected by print_stmt
     if stmt isa Expr
         if stmt.head === :enter && length(stmt.args) == 1 && stmt.args[1] isa Int
@@ -529,16 +548,17 @@ end
 
 # Show a single statement, code.stmts[idx]/code.code[idx], in the context of the whole IRCode/CodeInfo.
 # Returns the updated value of bb_idx.
-# line_info_preprinter(io::IO, indent::String, idx::Int) may print relevant info
-#   at the beginning of the line, and should at least print `indent`. It returns a
-#   string that will be printed after the final basic-block annotation.
-# line_info_postprinter(io::IO, typ, used::Bool) prints the type-annotation at the end
-#   of the statement
 # pop_new_node!(idx::Int) -> (node_idx, new_node_inst, new_node_type) may return a new
 #   node at the current index `idx`, which is printed before the statement at index
 #   `idx`. This function is repeatedly called until it returns `nothing`
+function show_ir_stmt(io::IO, code::Union{IRCode, CodeInfo}, idx::Int, config::IRShowConfig,
+                      used::BitSet, cfg::CFG, bb_idx::Int; pop_new_node! = Returns(nothing))
+    return show_ir_stmt(io, code, idx, config.line_info_preprinter, config.line_info_postprinter,
+                        used, cfg, bb_idx; pop_new_node!, config.bb_color)
+end
+
 function show_ir_stmt(io::IO, code::Union{IRCode, CodeInfo}, idx::Int, line_info_preprinter, line_info_postprinter,
-                      used::BitSet, cfg::CFG, bb_idx::Int, pop_new_node! = _ -> nothing; bb_color=:light_black)
+                      used::BitSet, cfg::CFG, bb_idx::Int; pop_new_node! = Returns(nothing), bb_color=:light_black)
     stmt = _stmt(code, idx)
     type = _type(code, idx)
     max_bb_idx_size = length(string(length(cfg.blocks)))
@@ -653,8 +673,8 @@ function ircode_new_nodes_iter(code::IRCode)
     end
 end
 
-# corresponds to `verbose_linetable=false`
-function ircode_default_linfo_printer(code::IRCode)
+# print only line numbers on the left, some of the method names and nesting depth on the right
+function inline_linfo_printer(code::IRCode)
     loc_annotations, loc_methods, loc_lineno = compute_ir_line_annotations(code)
     max_loc_width = maximum(length, loc_annotations)
     max_lineno_width = maximum(length, loc_lineno)
@@ -693,10 +713,11 @@ end
 _strip_color(s::String) = replace(s, r"\e\[\d+m" => "")
 
 # corresponds to `verbose_linetable=true`
-function ircode_verbose_linfo_printer(code::IRCode, used::BitSet)
+function ircode_verbose_linfo_printer(code::IRCode)
     stmts = code.stmts
     max_depth = maximum(compute_inlining_depth(code.linetable, stmts[i][:line]) for i in 1:length(stmts.line))
     last_stack = Ref(Int[])
+    used = stmts_used(code, false)
     maxlength_idx = if isempty(used)
         0
     else
@@ -733,16 +754,27 @@ function ircode_verbose_linfo_printer(code::IRCode, used::BitSet)
     end
 end
 
-function show_ir(io::IO, code::IRCode, expr_type_printer=default_expr_type_printer; verbose_linetable=false)
+function statementidx_lineinfo_printer(f, code::IRCode)
+    printer = f(code.linetable)
+    function (io::IO, indent::String, idx::Int)
+        printer(io, indent, idx > 0 ? code.stmts[idx][:line] : typemin(Int32))
+    end
+end
+function statementidx_lineinfo_printer(f, code::CodeInfo)
+    printer = f(code.linetable)
+    function (io::IO, indent::String, idx::Int)
+        printer(io, indent, idx > 0 ? code.codelocs[idx] : typemin(Int32))
+    end
+end
+statementidx_lineinfo_printer(code) = statementidx_lineinfo_printer(DILineInfoPrinter, code)
+
+function stmts_used(code::IRCode, warn_unset_entry=true)
     stmts = code.stmts
-    isempty(stmts) && return # unlikely, but avoid errors from reducing over empty sets
     used = BitSet()
-    cfg = code.cfg
     for stmt in stmts
         scan_ssa_use!(push!, used, stmt[:inst])
     end
     new_nodes = code.new_nodes.stmts
-    warn_unset_entry = true
     for nn in 1:length(new_nodes)
         if isassigned(new_nodes.inst, nn)
             scan_ssa_use!(push!, used, new_nodes[nn][:inst])
@@ -751,44 +783,42 @@ function show_ir(io::IO, code::IRCode, expr_type_printer=default_expr_type_print
             warn_unset_entry = false
         end
     end
-    bb_idx = 1
-
-    pop_new_node! = ircode_new_nodes_iter(code)
-
-    if verbose_linetable
-        line_info_preprinter = ircode_verbose_linfo_printer(code, used)
-    else
-        line_info_preprinter = ircode_default_linfo_printer(code)
-    end
-
-    for idx in 1:length(stmts)
-        bb_idx = show_ir_stmt(io, code, idx, line_info_preprinter, expr_type_printer,
-                              used, cfg, bb_idx, pop_new_node!; bb_color=:normal)
-    end
-    nothing
-end
-
-function statementidx_lineinfo_printer(f, code::CodeInfo)
-    printer = f(code.linetable)
-    return (io::IO, indent::String, idx::Int) -> printer(io, indent, idx > 0 ? code.codelocs[idx] : typemin(Int32))
+    return used
 end
-statementidx_lineinfo_printer(code::CodeInfo) = statementidx_lineinfo_printer(DILineInfoPrinter, code)
 
-function show_ir(io::IO, code::CodeInfo, line_info_preprinter=statementidx_lineinfo_printer(code), line_info_postprinter=default_expr_type_printer)
+function stmts_used(code::CodeInfo)
     stmts = code.code
     used = BitSet()
-    cfg = compute_basic_blocks(stmts)
     for stmt in stmts
         scan_ssa_use!(push!, used, stmt)
     end
+    return used
+end
+
+function default_config(code::IRCode; verbose_linetable=false)
+    return IRShowConfig(verbose_linetable ? ircode_verbose_linfo_printer(code)
+                                          : inline_linfo_printer(code);
+                        bb_color=:normal)
+end
+default_config(code::CodeInfo) = IRShowConfig(statementidx_lineinfo_printer(code))
+
+function show_ir(io::IO, code::Union{IRCode, CodeInfo}, config::IRShowConfig=default_config(code);
+                 pop_new_node! = code isa IRCode ? ircode_new_nodes_iter(code) : Returns(nothing))
+    stmts = code isa IRCode ? code.stmts : code.code
+    used = stmts_used(code)
+    cfg = code isa IRCode ? code.cfg : compute_basic_blocks(stmts)
     bb_idx = 1
 
     for idx in 1:length(stmts)
-        bb_idx = show_ir_stmt(io, code, idx, line_info_preprinter, line_info_postprinter, used, cfg, bb_idx)
+        if config.should_print_stmt(code, idx, used)
+            bb_idx = show_ir_stmt(io, code, idx, config, used, cfg, bb_idx; pop_new_node!)
+        elseif bb_idx <= length(cfg.blocks) && idx == cfg.blocks[bb_idx].stmts.stop
+            bb_idx += 1
+        end
     end
 
     max_bb_idx_size = length(string(length(cfg.blocks)))
-    line_info_preprinter(io, " "^(max_bb_idx_size + 2), 0)
+    config.line_info_preprinter(io, " "^(max_bb_idx_size + 2), 0)
     nothing
 end
 

base/reflection.jl

@@ -1262,7 +1262,11 @@ function print_statement_costs(io::IO, @nospecialize(tt::Type);
         maxcost = Core.Compiler.statement_costs!(cst, code.code, code, Any[match.sparams...], false, params)
         nd = ndigits(maxcost)
         println(io, meth)
-        IRShow.show_ir(io, code, (io, linestart, idx) -> (print(io, idx > 0 ? lpad(cst[idx], nd+1) : " "^(nd+1), " "); return ""))
+        irshow_config = IRShow.IRShowConfig() do io, linestart, idx
+            print(io, idx > 0 ? lpad(cst[idx], nd+1) : " "^(nd+1), " ")
+            return ""
+        end
+        IRShow.show_ir(io, code, irshow_config)
         println(io)
     end
 end

base/show.jl

@@ -2476,7 +2476,7 @@ function show(io::IO, src::CodeInfo; debuginfo::Symbol=:source)
         # TODO: static parameter values?
         # only accepts :source or :none, we can't have a fallback for default since
         # that would break code_typed(, debuginfo=:source) iff IRShow.default_debuginfo[] = :none
-        IRShow.show_ir(lambda_io, src, IRShow.__debuginfo[debuginfo](src))
+        IRShow.show_ir(lambda_io, src, IRShow.IRShowConfig(IRShow.__debuginfo[debuginfo](src)))
     else
         # this is a CodeInfo that has not been used as a method yet, so its locations are still LineNumberNodes
         body = Expr(:block)

stdlib/InteractiveUtils/src/codeview.jl

@@ -127,7 +127,8 @@ function code_warntype(io::IO, @nospecialize(f), @nospecialize(t); debuginfo::Sy
         print(io, "Body")
         warntype_type_printer(io, rettype, true)
         println(io)
-        Base.IRShow.show_ir(lambda_io, src, lineprinter(src), warntype_type_printer)
+        irshow_config = Base.IRShow.IRShowConfig(lineprinter(src), warntype_type_printer)
+        Base.IRShow.show_ir(lambda_io, src, irshow_config)
         println(io)
     end
     nothing

test/show.jl

@@ -1931,13 +1931,13 @@ let src = code_typed(my_fun28173, (Int,), debuginfo=:source)[1][1]
         @test repr(src) == repr_ir
     end
     lines1 = split(repr(ir), '\n')
-    @test isempty(pop!(lines1))
+    @test all(isspace, pop!(lines1))
     Core.Compiler.insert_node!(ir, 1, Core.Compiler.NewInstruction(QuoteNode(1), Val{1}), false)
     Core.Compiler.insert_node!(ir, 1, Core.Compiler.NewInstruction(QuoteNode(2), Val{2}), true)
     Core.Compiler.insert_node!(ir, length(ir.stmts.inst), Core.Compiler.NewInstruction(QuoteNode(3), Val{3}), false)
     Core.Compiler.insert_node!(ir, length(ir.stmts.inst), Core.Compiler.NewInstruction(QuoteNode(4), Val{4}), true)
     lines2 = split(repr(ir), '\n')
-    @test isempty(pop!(lines2))
+    @test all(isspace, pop!(lines2))
     @test popfirst!(lines2) == "2  1 ──       $(QuoteNode(1))"
     @test popfirst!(lines2) == "   │          $(QuoteNode(2))" # TODO: this should print after the next statement
     let line1 = popfirst!(lines1)
@@ -1958,7 +1958,7 @@ let src = code_typed(my_fun28173, (Int,), debuginfo=:source)[1][1]
 
     # verbose linetable
     io = IOBuffer()
-    Base.IRShow.show_ir(io, ir; verbose_linetable=true)
+    Base.IRShow.show_ir(io, ir, Base.IRShow.default_config(ir; verbose_linetable=true))
     seekstart(io)
     @test count(contains(r"my_fun28173 at a{80}:\d+"), eachline(io)) == 9
 end
@@ -1970,7 +1970,7 @@ let src = code_typed(gcd, (Int, Int), debuginfo=:source)[1][1]
     ir = Core.Compiler.inflate_ir(src)
     push!(ir.stmts.inst, Core.Compiler.ReturnNode())
     lines = split(sprint(show, ir), '\n')
-    @test isempty(pop!(lines))
+    @test all(isspace, pop!(lines))
     @test pop!(lines) == "   !!! ──       unreachable::#UNDEF"
 end