Skip to main content

haloumi_ir/
printer.rs

1//! Types and traits for human-friendly printing of the IR.
2
3use std::fmt::{Formatter, Write};
4
5/// Prints a human-friendly representation of the IR meant for debugging.
6///
7/// The structure of the output emitted by the printer is never considered stable and shouldn't be
8/// relied upon as it may change unexpectedly. The purpose of the printer is to be a debugging aid
9/// for inspecting the shape of the IR and not a serialization/deserialization format.
10#[derive(Copy, Clone)]
11pub struct IRPrinter<'ir> {
12    ir: &'ir dyn IRPrintable,
13}
14
15impl std::fmt::Display for IRPrinter<'_> {
16    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
17        let mut ctx = IRPrinterCtx::new(f);
18        self.ir.fmt(&mut ctx)
19    }
20}
21
22impl std::fmt::Debug for IRPrinter<'_> {
23    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
24        f.debug_struct("IRPrinter").finish()
25    }
26}
27
28impl<'ir, T: IRPrintable> From<&'ir T> for IRPrinter<'ir> {
29    fn from(ir: &'ir T) -> Self {
30        Self { ir }
31    }
32}
33
34/// Alias type to a formatting result.
35pub type Result = std::fmt::Result;
36
37/// Trait defining the common behavior of printable IR objects.
38pub trait IRPrintable {
39    /// Format the IR object using the given context.
40    fn fmt(&self, ctx: &mut IRPrinterCtx<'_, '_>) -> Result;
41
42    /// Computes the depth of the IR tree to give hints to the printer.
43    fn depth(&self) -> usize {
44        1
45    }
46}
47
48/// Printing context.
49///
50/// Implements [`std::fmt::Write`] and it's meant to be used as the first argument of [`write!`].
51pub struct IRPrinterCtx<'a, 'f> {
52    f: &'a mut Formatter<'f>,
53    indent: Vec<usize>,
54    indent_pending: bool,
55}
56
57impl<'a, 'f> IRPrinterCtx<'a, 'f> {
58    fn new(f: &'a mut Formatter<'f>) -> Self {
59        Self {
60            f,
61            indent: vec![],
62            indent_pending: true,
63        }
64    }
65
66    /// Adds a new line to the output.
67    pub fn nl(&mut self) -> Result {
68        if !self.indent_pending {
69            self.indent_pending = true;
70            writeln!(self.f)?;
71        }
72        Ok(())
73    }
74
75    fn push_indent(&mut self, value: usize) {
76        self.indent.push(value);
77    }
78
79    fn pop_indent(&mut self) {
80        self.indent.pop();
81    }
82
83    fn do_indent(&mut self) -> Result {
84        if !self.indent_pending {
85            return Ok(());
86        }
87
88        write!(self.f, "{}", " ".repeat(self.indent.iter().sum()))?;
89        self.indent_pending = false;
90        Ok(())
91    }
92
93    /// Adds a list to the output and then a new line.
94    pub fn list_nl(
95        &mut self,
96        atom: &str,
97        body: impl FnOnce(&mut IRPrinterCtx) -> Result,
98    ) -> Result {
99        self.list(atom, body)?;
100        self.nl()
101    }
102
103    /// Adds a list to the output.
104    pub fn list(&mut self, atom: &str, body: impl FnOnce(&mut IRPrinterCtx) -> Result) -> Result {
105        write!(self, "(")?;
106        if !atom.is_empty() {
107            write!(self, "{atom} ")?;
108        }
109        body(self)?;
110        write!(self, ")")
111    }
112
113    /// Adds a list to the output indenting the output inside it.
114    pub fn block(&mut self, atom: &str, body: impl FnOnce(&mut IRPrinterCtx) -> Result) -> Result {
115        self.list(atom, |ctx| {
116            ctx.push_indent(2 + atom.len());
117            body(ctx)?;
118            ctx.pop_indent();
119            Ok(())
120        })
121    }
122
123    pub(crate) fn fmt_call<I: IRPrintable, O: IRPrintable>(
124        &mut self,
125        callee: &str,
126        inputs: &[I],
127        outputs: &[O],
128        id: Option<usize>,
129    ) -> Result {
130        self.block("call", |ctx| {
131            if let Some(id) = id {
132                write!(ctx, "{id} ")?;
133            }
134            writeln!(ctx, "\"{}\" ", callee)?;
135            ctx.block("inputs", |ctx| {
136                let do_nl = inputs.iter().any(|expr| expr.depth() > 1);
137                let mut is_first = true;
138                for expr in inputs {
139                    if do_nl && !is_first {
140                        ctx.nl()?;
141                    }
142                    is_first = false;
143                    expr.fmt(ctx)?;
144                }
145                Ok(())
146            })?;
147            ctx.nl()?;
148            ctx.list("outputs", |ctx| {
149                for output in outputs {
150                    output.fmt(ctx)?;
151                }
152                Ok(())
153            })
154        })
155    }
156}
157
158impl Write for IRPrinterCtx<'_, '_> {
159    fn write_str(&mut self, s: &str) -> Result {
160        let ends_with_nl = s.ends_with('\n');
161        let mut lines = s.lines().peekable();
162        loop {
163            let Some(line) = lines.next() else {
164                self.indent_pending = ends_with_nl;
165                return Ok(());
166            };
167            let not_done = lines.peek().is_some();
168            self.do_indent()?;
169
170            write!(self.f, "{}", line)?;
171            if not_done || ends_with_nl {
172                writeln!(self.f)?;
173            }
174        }
175    }
176}
177
178impl std::fmt::Debug for IRPrinterCtx<'_, '_> {
179    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
180        f.debug_struct("IRPrinterCtx")
181            .field("indent", &self.indent)
182            .field("indent_pending", &self.indent_pending)
183            .finish()
184    }
185}