1use std::fmt::{Formatter, Write};
4
5#[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
34pub type Result = std::fmt::Result;
36
37pub trait IRPrintable {
39 fn fmt(&self, ctx: &mut IRPrinterCtx<'_, '_>) -> Result;
41
42 fn depth(&self) -> usize {
44 1
45 }
46}
47
48pub 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 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 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 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 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}