Add PrintDecorator hooks for Waffle's IR

src/ir/display.rs

@@ -1,20 +1,69 @@
 //! Displaying IR.
 
-use super::{FuncDecl, FunctionBody, Module, SourceLoc, ValueDef};
+use super::{Func, FuncDecl, FunctionBody, Module, SourceLoc, ValueDef};
 use crate::entity::EntityRef;
 use std::collections::HashMap;
-use std::fmt::{Display, Formatter, Result as FmtResult};
+use std::fmt::{self, Display, Formatter, Result as FmtResult};
+
+/// Hooks to print information after instruction, before and after blocks
+/// and before and after functions.
+pub trait PrintDecorator {
+    /// Print arbitrary text after an instruction.
+    ///
+    /// Invoked after every instruction in a block. The instruction has already been printed on its own line;
+    /// this method can print content after the operator, if desired.
+    fn after_inst(&self, _value: super::Value, _f: &mut fmt::Formatter) -> fmt::Result {
+        Ok(())
+    }
+
+    /// Print arbitrary text before the body of a block.
+    ///
+    /// Invoked before the block body. The block id and parameters have already been printed on its own line;
+    /// this method can print content on the line below the block id, before the body of the block is printed, if desired.
+    fn before_block(&self, _block: super::Block, _f: &mut fmt::Formatter) -> fmt::Result {
+        Ok(())
+    }
+
+    /// Print arbitrary text after the body of a block.
+    ///
+    /// Invoked after the block body, before the terminator. The block body has already been printed on its own line(s);
+    /// this method can print content on the line after the last instruction in the block body, before the terminator is printed.
+    fn after_block(&self, _block: super::Block, _f: &mut fmt::Formatter) -> fmt::Result {
+        Ok(())
+    }
+
+    /// Print arbitrary text before the body of a function.
+    ///
+    /// Invoked before the function body is printed. The function id and signature have already been printed on its own line;
+    /// this method can print content on the line before the function signature line, before the function body is printed.
+    fn before_function_body(&self, _f: &mut fmt::Formatter) -> fmt::Result {
+        Ok(())
+    }
+
+    /// Print arbitrary text after the body of a function.
+    ///
+    /// Invoked after the function body is printed. The function body has already been printed;
+    /// this method can print content on the line after the return block of the function, before the last curly brace to end the function is printed.
+    fn after_function_body(&self, _f: &mut fmt::Formatter) -> fmt::Result {
+        Ok(())
+    }
+}
+
+#[derive(Default)]
+pub(crate) struct NOPPrintDecorator();
+impl PrintDecorator for NOPPrintDecorator {}
 
 /// A wrapper around a `FunctionBody` together with some auxiliary
 /// information to perform a pretty-print of that function.
-pub struct FunctionBodyDisplay<'a> {
+pub struct FunctionBodyDisplay<'a, PD: PrintDecorator> {
     pub(crate) body: &'a FunctionBody,
     pub(crate) indent: &'a str,
     pub(crate) verbose: bool,
     pub(crate) module: Option<&'a Module<'a>>,
+    pub(crate) decorator: Option<&'a PD>,
 }
 
-impl<'a> Display for FunctionBodyDisplay<'a> {
+impl<'a, PD: PrintDecorator> Display for FunctionBodyDisplay<'a, PD> {
     fn fmt(&self, f: &mut Formatter) -> FmtResult {
         let arg_tys = self
             .body
@@ -37,25 +86,31 @@ impl<'a> Display for FunctionBodyDisplay<'a> {
             ret_tys.join(", ")
         )?;
 
+        if let Some(decorator) = self.decorator {
+            decorator.before_function_body(f)?;
+        }
+
         for (value, value_def) in self.body.values.entries() {
             match value_def {
-                ValueDef::Operator(op, args, tys) if self.verbose => writeln!(
-                    f,
-                    "{}    {} = {} {} # {}",
-                    self.indent,
-                    value,
-                    op,
-                    self.body.arg_pool[*args]
-                        .iter()
-                        .map(|arg| format!("{}", arg))
-                        .collect::<Vec<_>>()
-                        .join(", "),
-                    self.body.type_pool[*tys]
-                        .iter()
-                        .map(|arg| format!("{}", arg))
-                        .collect::<Vec<_>>()
-                        .join(", ")
-                )?,
+                ValueDef::Operator(op, args, tys) if self.verbose => {
+                    writeln!(
+                        f,
+                        "{}    {} = {} {} # {} ",
+                        self.indent,
+                        value,
+                        op,
+                        self.body.arg_pool[*args]
+                            .iter()
+                            .map(|arg| format!("{}", arg))
+                            .collect::<Vec<_>>()
+                            .join(", "),
+                        self.body.type_pool[*tys]
+                            .iter()
+                            .map(|arg| format!("{}", arg))
+                            .collect::<Vec<_>>()
+                            .join(", "),
+                    )?;
+                }
                 ValueDef::BlockParam(block, idx, ty) if self.verbose => writeln!(
                     f,
                     "{}    {} = blockparam {}, {} # {}",
@@ -91,6 +146,11 @@ impl<'a> Display for FunctionBodyDisplay<'a> {
                 block_params.join(", "),
                 block.desc
             )?;
+
+            if let Some(decorator) = self.decorator {
+                decorator.before_block(block_id, f)?
+            };
+
             writeln!(
                 f,
                 "{}    # preds: {}",
@@ -143,16 +203,20 @@ impl<'a> Display for FunctionBodyDisplay<'a> {
                         } else {
                             "".to_owned()
                         };
-                        writeln!(
+                        write!(
                             f,
-                            "{}    {} = {} {} # {} {}",
+                            "{}    {} = {} {} # {} {} ",
                             self.indent,
                             inst,
                             op,
                             args.join(", "),
                             tys.join(", "),
                             loc,
                         )?;
+                        if let Some(decorator) = self.decorator {
+                            decorator.after_inst(inst, f)?;
+                        }
+                        writeln!(f, "")?;
                     }
                     ValueDef::PickOutput(val, idx, ty) => {
                         writeln!(f, "{}    {} = {}.{} # {}", self.indent, inst, val, idx, ty)?;
@@ -163,20 +227,27 @@ impl<'a> Display for FunctionBodyDisplay<'a> {
                     _ => unreachable!(),
                 }
             }
+            if let Some(decorator) = self.decorator {
+                decorator.after_block(block_id, f)?;
+            }
             writeln!(f, "{}    {}", self.indent, block.terminator)?;
         }
 
+        if let Some(decorator) = self.decorator {
+            decorator.after_function_body(f)?;
+        }
         writeln!(f, "{}}}", self.indent)?;
 
         Ok(())
     }
 }
 
-pub struct ModuleDisplay<'a> {
+pub struct ModuleDisplay<'a, PD: PrintDecorator> {
     pub(crate) module: &'a Module<'a>,
+    pub(crate) decorators: Option<Box<dyn Fn(Func) -> PD>>,
 }
 
-impl<'a> Display for ModuleDisplay<'a> {
+impl<'a, PD: PrintDecorator> Display for ModuleDisplay<'a, PD> {
     fn fmt(&self, f: &mut Formatter) -> FmtResult {
         writeln!(f, "module {{")?;
         if let Some(func) = self.module.start_func {
@@ -250,7 +321,15 @@ impl<'a> Display for ModuleDisplay<'a> {
                         sig,
                         sig_strs.get(&sig).unwrap()
                     )?;
-                    writeln!(f, "{}", body.display("    ", Some(self.module)))?;
+
+                    if let Some(decorator) = &self.decorators {
+                        let decorator = &(*decorator)(func);
+                        writeln!(
+                            f,
+                            "{}",
+                            body.display_with_decorator("    ", Some(self.module), decorator)
+                        )?;
+                    }
                 }
                 FuncDecl::Lazy(sig, name, reader) => {
                     writeln!(

src/ir/func.rs

@@ -1,4 +1,7 @@
-use super::{Block, FunctionBodyDisplay, Local, Module, Signature, Type, Value, ValueDef};
+use super::{
+    Block, FunctionBodyDisplay, Local, Module, NOPPrintDecorator, PrintDecorator, Signature, Type,
+    Value, ValueDef,
+};
 use crate::backend::WasmFuncBackend;
 use crate::cfg::CFGInfo;
 use crate::entity::{EntityRef, EntityVec, PerEntity};
@@ -450,12 +453,32 @@ impl FunctionBody {
         &'a self,
         indent: &'a str,
         module: Option<&'a Module>,
-    ) -> FunctionBodyDisplay<'a> {
+    ) -> FunctionBodyDisplay<'a, impl PrintDecorator> {
+        FunctionBodyDisplay::<NOPPrintDecorator> {
+            body: self,
+            indent,
+            verbose: false,
+            module,
+            decorator: None,
+        }
+    }
+
+    /// Prety-print this function body with some additional information.
+    /// `indent` is prepended to each line of output.
+    /// `module`, if provided, allows printing source locations as comments at each operator.
+    /// `decorator` describes how the additional information should be printed in the IR.
+    pub fn display_with_decorator<'a, PD: PrintDecorator>(
+        &'a self,
+        indent: &'a str,
+        module: Option<&'a Module>,
+        decorator: &'a PD,
+    ) -> FunctionBodyDisplay<'a, PD> {
         FunctionBodyDisplay {
             body: self,
             indent,
             verbose: false,
             module,
+            decorator: Some(&decorator),
         }
     }
 
@@ -466,12 +489,13 @@ impl FunctionBody {
         &'a self,
         indent: &'a str,
         module: Option<&'a Module>,
-    ) -> FunctionBodyDisplay<'a> {
-        FunctionBodyDisplay {
+    ) -> FunctionBodyDisplay<'a, impl PrintDecorator> {
+        FunctionBodyDisplay::<NOPPrintDecorator> {
             body: self,
             indent,
             verbose: true,
             module,
+            decorator: None,
         }
     }
 

src/ir/module.rs

@@ -1,4 +1,7 @@
-use super::{Func, FuncDecl, Global, Memory, ModuleDisplay, Signature, Table, Type};
+use super::{
+    Func, FuncDecl, Global, Memory, ModuleDisplay, NOPPrintDecorator, PrintDecorator, Signature,
+    Table, Type,
+};
 use crate::entity::{EntityRef, EntityVec};
 use crate::ir::{Debug, DebugMap, FunctionBody};
 use crate::{backend, frontend};
@@ -327,11 +330,30 @@ impl<'a> Module<'a> {
 
     /// Return a wrapper that implements Display on this module,
     /// pretty-printing it as textual IR.
-    pub fn display<'b>(&'b self) -> ModuleDisplay<'b>
+    pub fn display<'b>(&'b self) -> ModuleDisplay<'b, impl PrintDecorator>
     where
         'b: 'a,
     {
-        ModuleDisplay { module: self }
+        ModuleDisplay::<NOPPrintDecorator> {
+            module: self,
+            decorators: None,
+        }
+    }
+
+    /// Return a wrapper that implements Display on this module,
+    /// pretty-printing it as textual IR with some additional text whose
+    /// printing is described in Decorator.
+    pub fn display_with_decorator<'b, PD: PrintDecorator>(
+        &'b self,
+        decorators: Box<dyn Fn(Func) -> PD>,
+    ) -> ModuleDisplay<'b, PD>
+    where
+        'b: 'a,
+    {
+        ModuleDisplay {
+            module: self,
+            decorators: Some(decorators),
+        }
     }
 
     /// Internal (used during parsing): create an empty module, with