Add option to dump optimized HIR

Author: tekknolagi

zjit/src/codegen.rs

@@ -3,8 +3,9 @@ use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::invariants::{iseq_escapes_ep, track_no_ep_escape_assumption};
 use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, SP};
 use crate::hir::{self, Block, BlockId, BranchEdge, CallInfo};
-use crate::hir::{Const, FrameState, Function, Insn, InsnId};
+use crate::hir::{Const, FrameState, Function, Insn, InsnId, FunctionPrinter};
 use crate::hir_type::{types::Fixnum, Type};
+use crate::options::{get_option, DumpHIR};
 
 /// Ephemeral code generation state
 struct JITState {
@@ -88,6 +89,12 @@ fn iseq_gen_entry_point(iseq: IseqPtr) -> *const u8 {
             }
         };
         ssa.optimize();
+        match get_option!(dump_hir_opt) {
+            Some(DumpHIR::WithoutSnapshot) => println!("HIR:\n{}", FunctionPrinter::without_snapshot(&ssa)),
+            Some(DumpHIR::All) => println!("HIR:\n{}", FunctionPrinter::with_snapshot(&ssa)),
+            Some(DumpHIR::Raw) => println!("HIR:\n{:#?}", &ssa),
+            None => {},
+        }
 
         // Compile High-level IR into machine code
         let cb = ZJITState::get_code_block();

zjit/src/hir.rs

@@ -413,22 +413,22 @@ impl Block {
     }
 }
 
-struct FunctionPrinter<'a> {
+pub struct FunctionPrinter<'a> {
     fun: &'a Function,
     display_snapshot: bool,
     ptr_map: PtrPrintMap,
 }
 
 impl<'a> FunctionPrinter<'a> {
-    fn without_snapshot(fun: &'a Function) -> Self {
+    pub fn without_snapshot(fun: &'a Function) -> Self {
         let mut ptr_map = PtrPrintMap::identity();
         if cfg!(test) {
             ptr_map.map_ptrs = true;
         }
         Self { fun, display_snapshot: false, ptr_map }
     }
 
-    fn with_snapshot(fun: &'a Function) -> FunctionPrinter<'a> {
+    pub fn with_snapshot(fun: &'a Function) -> FunctionPrinter<'a> {
         let mut printer = Self::without_snapshot(fun);
         printer.display_snapshot = true;
         printer

zjit/src/options.rs

@@ -15,6 +15,9 @@ pub struct Options {
     /// Dump High-level IR generated from ISEQ.
     pub dump_hir: Option<DumpHIR>,
 
+    /// Dump High-level IR after optimization, right before codegen.
+    pub dump_hir_opt: Option<DumpHIR>,
+
     /// Dump all compiled machine code.
     pub dump_disasm: bool,
 }
@@ -56,6 +59,7 @@ pub fn init_options() -> Options {
     Options {
         debug: false,
         dump_hir: None,
+        dump_hir_opt: None,
         dump_disasm: false,
     }
 }
@@ -97,6 +101,10 @@ fn parse_option(options: &mut Options, str_ptr: *const std::os::raw::c_char) ->
         ("dump-hir", "all") => options.dump_hir = Some(DumpHIR::All),
         ("dump-hir", "raw") => options.dump_hir = Some(DumpHIR::Raw),
 
+        ("dump-hir-opt", "") => options.dump_hir_opt = Some(DumpHIR::WithoutSnapshot),
+        ("dump-hir-opt", "all") => options.dump_hir_opt = Some(DumpHIR::All),
+        ("dump-hir-opt", "raw") => options.dump_hir_opt = Some(DumpHIR::Raw),
+
         ("dump-disasm", "") => options.dump_disasm = true,
 
         _ => return None, // Option name not recognized