do some registers

analysis/src/ir.rs

@@ -33,7 +33,7 @@
 
 mod custom;
 pub mod info;
-mod pretty;
+pub mod pretty;
 mod validate;
 mod visit;
 
@@ -299,3 +299,21 @@ impl Branch {
         }
     }
 }
+
+impl Location {
+    pub fn start() -> Self {
+        Self {
+            bb: BbIdx(0),
+            stmt: Some(0),
+        }
+    }
+    pub fn terminator(bb: BbIdx) -> Self {
+        Self { bb, stmt: None }
+    }
+    pub fn stmt(bb: BbIdx, stmt: usize) -> Self {
+        Self {
+            bb,
+            stmt: Some(stmt),
+        }
+    }
+}

analysis/src/ir/pretty.rs

@@ -1,39 +1,63 @@
-use std::fmt::{Display, Formatter, Result, Write};
+use std::{
+    cell::Cell,
+    fmt::{self, Display, Formatter, Result, Write},
+};
 
-use super::{BbIdx, BinKind, Branch, ConstValue, Func, Ir, Operand, StatementKind, UnaryKind};
+use super::{
+    BbIdx, BinKind, Branch, ConstValue, Func, Ir, Location, Operand, StatementKind, UnaryKind,
+};
 use crate::ir::Register;
 
-pub fn ir_to_string(ir: &Ir<'_>) -> String {
+pub fn ir_to_string<'a>(ir: &'a Ir<'a>, custom: &impl Customizer<'a>) -> String {
     let mut buf = String::new();
-    PrettyPrinter { out: &mut buf }.ir(ir).unwrap();
+    PrettyPrinter { out: &mut buf }.ir(ir, custom).unwrap();
     buf
 }
 
-pub fn func_to_string(func: &Func<'_>) -> String {
+pub fn func_to_string<'a>(func: &'a Func<'a>, custom: &impl Customizer<'a>) -> String {
     let mut buf = String::new();
-    PrettyPrinter { out: &mut buf }.func(func).unwrap();
+    PrettyPrinter { out: &mut buf }.func(func, custom).unwrap();
     buf
 }
 
+pub trait Customizer<'a> {
+    fn start_func(&self, func: &'a Func<'a>);
+    fn fmt_reg(&self, reg: Register, f: &mut fmt::Formatter<'_>, loc: Location) -> fmt::Result;
+}
+
+#[derive(Default)]
+pub struct DefaultCustomizer<'a>(Cell<Option<&'a Func<'a>>>);
+
+impl<'a> Customizer<'a> for DefaultCustomizer<'a> {
+    fn start_func(&self, func: &'a Func<'a>) {
+        self.0.set(Some(func));
+    }
+
+    fn fmt_reg(&self, reg: Register, f: &mut fmt::Formatter<'_>, loc: Location) -> fmt::Result {
+        match self.0.get().unwrap().regs[reg.0 as usize].name {
+            None => write!(f, "%{}", reg.0),
+            Some(name) => write!(f, "%{name}"),
+        }
+    }
+}
+
 pub struct PrettyPrinter<W> {
     out: W,
 }
 
 impl<W: Write> PrettyPrinter<W> {
-    pub fn ir(&mut self, ir: &Ir<'_>) -> Result {
+    pub fn ir<'a>(&mut self, ir: &'a Ir<'a>, custom: &impl Customizer<'a>) -> Result {
         for func in ir.funcs.values() {
-            self.func(func)?;
+            self.func(func, custom)?;
         }
         Ok(())
     }
 
-    pub fn func(&mut self, func: &Func<'_>) -> Result {
-        let print_reg = |reg: Register| {
-            display_fn(move |f| match func.regs[reg.0 as usize].name {
-                None => write!(f, "%{}", reg.0),
-                Some(name) => write!(f, "%{name}"),
-            })
-        };
+    pub fn func<'a>(&mut self, func: &'a Func<'a>, custom: &impl Customizer<'a>) -> Result {
+        custom.start_func(func);
+
+        let print_reg =
+            |reg: Register, loc: Location| display_fn(move |f| custom.fmt_reg(reg, f, loc));
 
         write!(self.out, "def {}(", func.name)?;
         for param in 0..func.arity {
@@ -42,7 +66,7 @@ impl<W: Write> PrettyPrinter<W> {
                 self.out,
                 "{} {}",
                 reg.tyl.ty,
-                print_reg(Register(param as _))
+                print_reg(Register(param as _), Location::start())
             )?;
             if (param + 1) != func.arity {
                 write!(self.out, ", ")?;
@@ -50,20 +74,25 @@ impl<W: Write> PrettyPrinter<W> {
         }
         writeln!(self.out, ") {{",)?;
 
-        let print_op = |op: Operand| {
+        let print_op = |op: Operand, loc: Location| {
             display_fn(move |f| match op {
                 Operand::Const(c) => Display::fmt(&c, f),
-                Operand::Reg(reg) => Display::fmt(&print_reg(reg), f),
+                Operand::Reg(reg) => Display::fmt(&print_reg(reg, loc), f),
             })
         };
 
         for (i, bb) in func.bbs.iter().enumerate() {
+            let bb_idx = BbIdx::from_usize(i);
             if i > 0 {
                 writeln!(self.out)?;
             }
-            writeln!(self.out, "  {}:", BbIdx::from_usize(i))?;
+            writeln!(self.out, "  {}:", bb_idx)?;
+
+            for (stmt_idx, stmt) in bb.statements.iter().enumerate() {
+                let loc = Location::stmt(bb_idx, stmt_idx);
+                let print_reg = |reg| print_reg(reg, loc);
+                let print_op = |op| print_op(op, loc);
 
-            for stmt in &bb.statements {
                 match stmt.kind {
                     StatementKind::Alloca {
                         result: reg,
@@ -177,14 +206,15 @@ impl<W: Write> PrettyPrinter<W> {
                 }?;
             }
 
+            let loc = Location::terminator(bb_idx);
             match bb.term {
                 Branch::Goto(bbn) => writeln!(self.out, "    goto {}", bbn)?,
                 Branch::Switch { cond, yes, no } => writeln!(
                     self.out,
                     "    switch {}, then {yes}, else {no}",
-                    print_op(cond)
+                    print_op(cond, loc)
                 )?,
-                Branch::Ret(op) => writeln!(self.out, "    ret {}", print_op(op))?,
+                Branch::Ret(op) => writeln!(self.out, "    ret {}", print_op(op, loc))?,
             }
         }
 

analysis/src/lower.rs

@@ -37,7 +37,9 @@ pub fn lower_translation_unit<'cx>(
 
                 let (ref declarator, def_span) = decl.init_declarators[0];
 
-                let ast::DirectDeclarator::WithParams { ident, params } = &declarator.declarator.decl else {
+                let ast::DirectDeclarator::WithParams { ident, params } =
+                    &declarator.declarator.decl
+                else {
                     unreachable!("function def needs withparams declarator");
                 };
 
@@ -144,7 +146,10 @@ impl<'a, 'cx> FnLoweringCtxt<'a, 'cx> {
             todo!("complex lvalues")
         };
         let Some(var) = self.resolve_ident(ident) else {
-            return Err(Error::new(format!("cannot find variable {ident}"), ident_span));
+            return Err(Error::new(
+                format!("cannot find variable {ident}"),
+                ident_span,
+            ));
         };
         Ok(match var.kind {
             VariableInfoKind::Local { ptr_to } => (Operand::Reg(ptr_to), var.tyl),
@@ -245,7 +250,10 @@ impl<'a, 'cx> FnLoweringCtxt<'a, 'cx> {
             ast::Expr::Atom(ast::Atom::Float(_)) => todo!("no floats"),
             ast::Expr::Atom(ast::Atom::Ident((ident, ident_span))) => {
                 let Some(var) = self.resolve_ident(*ident) else {
-                    return Err(Error::new(format!("cannot find variable {ident}"), *ident_span));
+                    return Err(Error::new(
+                        format!("cannot find variable {ident}"),
+                        *ident_span,
+                    ));
                 };
                 let tyl = var.tyl;
                 match var.kind {
@@ -445,13 +453,7 @@ fn lower_func<'cx>(
 ) -> Result<Func<'cx>, Error> {
     let mut cx = FnLoweringCtxt {
         scopes: vec![Default::default()],
-        build: FuncBuilder::new(
-            name,
-            def_span,
-            ret_ty,
-            lcx,
-            params.len(),
-        ),
+        build: FuncBuilder::new(name, def_span, ret_ty, lcx, params.len()),
         lcx,
     };
 

analysis/src/lower/builder.rs

@@ -166,7 +166,10 @@ impl<'a, 'cx> FuncBuilder<'a, 'cx> {
     }
 
     pub fn finish(self) -> Func<'cx> {
-        println!("{}", ir::func_to_string(&self.ir));
+        println!(
+            "{}",
+            ir::func_to_string(&self.ir, &ir::pretty::DefaultCustomizer::default())
+        );
 
         self.ir
     }