1  // This code is bad!
  2  
  3  use crate::ir::*;
  4  use std::fmt::{Display, Formatter, Result};
  5  
  6  type F<'a, 'b> = &'a mut Formatter<'b>;
  7  
  8  const DISPLAY_TYS: bool = false;
  9  
 10  struct SepBy<'a, T, U: ?Sized>(T, &'a U);
 11  impl<T, I, U> Display for SepBy<'_, T, U>
 12  where
 13      T: Clone + IntoIterator<Item = I>,
 14      I: Display,
 15      U: Display + ?Sized,
 16  {
 17      fn fmt(&self, f: F) -> Result {
 18          for (i, item) in self.0.clone().into_iter().enumerate() {
 19              if i != 0 {
 20                  write!(f, "{}", self.1)?;
 21              }
 22              write!(f, "{item}")?;
 23          }
 24          Ok(())
 25      }
 26  }
 27  
 28  struct Commas<T>(T);
 29  impl<T, I> Display for Commas<T>
 30  where
 31      T: Clone + IntoIterator<Item = I>,
 32      I: Display,
 33  {
 34      fn fmt(&self, f: F) -> Result {
 35          write!(f, "{}", SepBy(self.0.clone(), ", "))
 36      }
 37  }
 38  
 39  struct WithTy<'a, T: ?Sized>(&'a T, &'a Ty);
 40  impl<T: Display + ?Sized> Display for WithTy<'_, T> {
 41      fn fmt(&self, f: F) -> Result {
 42          write!(f, "{}", self.0)?;
 43          if DISPLAY_TYS {
 44              write!(f, " {}", self.1)?;
 45          }
 46          Ok(())
 47      }
 48  }
 49  
 50  impl<'a, T> Clone for WithTy<'a, T> {
 51      fn clone(&self) -> Self {
 52          *self
 53      }
 54  }
 55  
 56  impl<'a, T> Copy for WithTy<'a, T> {}
 57  
 58  struct Returns<T>(T);
 59  impl<T, I> Display for Returns<T>
 60  where
 61      T: Clone + IntoIterator<Item = I>,
 62      I: Display,
 63  {
 64      fn fmt(&self, f: F) -> Result {
 65          let oops: Vec<_> = self.0.clone().into_iter().collect();
 66          match oops.len() {
 67              0 => Ok(()),
 68              1 => write!(f, "{}", oops[0]),
 69              _ => write!(f, "{{{}}}", Commas(&oops)),
 70          }
 71      }
 72  }
 73  
 74  struct ReturnsSpace<T>(T);
 75  impl<T, I> Display for ReturnsSpace<T>
 76  where
 77      T: Clone + IntoIterator<Item = I>,
 78      I: Display,
 79  {
 80      fn fmt(&self, f: F) -> Result {
 81          let returns = format!("{}", Returns(self.0.clone()));
 82          if returns.is_empty() {
 83              Ok(())
 84          } else {
 85              write!(f, " {returns}")
 86          }
 87      }
 88  }
 89  
 90  trait DisplayWithName {
 91      fn fmt_with_name(&self, f: F, name: &str) -> Result;
 92      fn with_name<'a>(&'a self, name: &'a str) -> impl Display + 'a
 93      where
 94          Self: Sized,
 95      {
 96          struct WithName<'a, T>(&'a T, &'a str);
 97          impl<T: DisplayWithName> Display for WithName<'_, T> {
 98              fn fmt(&self, f: F) -> Result {
 99                  self.0.fmt_with_name(f, self.1)
100              }
101          }
102          WithName(self, name)
103      }
104  }
105  
106  impl DisplayWithName for BlockId {
107      fn fmt_with_name(&self, f: F, name: &str) -> Result {
108          write!(f, "{name}_{}", self.0)
109      }
110  }
111  
112  impl DisplayWithName for Exit {
113      fn fmt_with_name(&self, f: F, name: &str) -> Result {
114          match self {
115              Self::Jump(loc) => write!(f, "goto {}", loc.with_name(name)),
116              Self::CondJump(cond, loc_true, loc_false) => write!(
117                  f,
118                  "if {cond}: goto {} else goto {}",
119                  loc_true.with_name(name),
120                  loc_false.with_name(name)
121              ),
122              Self::Return(regs) => write!(f, "return{}", ReturnsSpace(regs)),
123          }
124      }
125  }
126  
127  impl DisplayWithName for Function {
128      fn fmt_with_name<'a>(&'a self, f: F, name: &str) -> Result {
129          let reg_def = |r: &'a Register| WithTy(r, self.tys.get(r).unwrap_or(&Ty::Int(IntKind::U8)));
130          write!(
131              f,
132              "fn {name}({})",
133              Commas(self.parameters.iter().map(reg_def)),
134          )?;
135          for (id, block) in self.iter() {
136              write!(f, "\n{}:", id.with_name(name))?;
137              for inst in &block.insts {
138                  write!(f, "\n    ")?;
139                  match inst {
140                      Inst::Store(r, sk) => {
141                          use StoreKind as Sk;
142                          write!(f, "{} = ", reg_def(r))?;
143                          match sk {
144                              Sk::StackAlloc(ty) => write!(f, "StackAlloc({ty})"),
145                              Sk::Int(i, kind) => write!(f, "{i}_{kind}"),
146                              Sk::Copy(r) => write!(f, "{r}"),
147                              Sk::Read(r) => write!(f, "{r}^"),
148                              Sk::UnaryOp(op, inner) => write!(
149                                  f,
150                                  "{}{inner}",
151                                  match op {
152                                      UnaryOp::Neg => "-",
153                                  }
154                              ),
155                              Sk::BinOp(op, lhs, rhs) => write!(
156                                  f,
157                                  "{lhs} {} {rhs}",
158                                  match op {
159                                      BinOp::Add => "+",
160                                      BinOp::Sub => "-",
161                                      BinOp::Mul => "*",
162                                      BinOp::CmpLe => "<=",
163                                  }
164                              ),
165                              Sk::IntCast(inner, ty) => write!(f, "{inner} as {ty}"),
166                              Sk::PtrOffset(lhs, rhs) => write!(f, "{lhs}[{rhs}]@"),
167                              Sk::FieldOffset(inner, field) => write!(f, "{inner}.{field}@"),
168                              Sk::Phi(regs) => write!(
169                                  f,
170                                  "Phi({})",
171                                  Commas(
172                                      regs.iter()
173                                          .map(|(id, r)| format!("{}: {r}", id.with_name(name)))
174                                  )
175                              ),
176  
177                              Sk::Function(name) => write!(f, "{name}"),
178                          }
179                      }
180                      Inst::Write(dst, src) => write!(f, "{dst}^ = {src}"),
181                      Inst::Call {
182                          callee,
183                          returns,
184                          args,
185                      } => write!(
186                          f,
187                          "[{}] = {callee}({})",
188                          Returns(returns.iter().map(reg_def)),
189                          Commas(args)
190                      ),
191                      Inst::Nop => write!(f, "Nop"),
192                  }?;
193              }
194              write!(f, "\n    {}", block.exit.with_name(name))?;
195          }
196          Ok(())
197      }
198  }
199  
200  impl Display for Register {
201      fn fmt(&self, f: F) -> Result {
202          write!(f, "%{}", self.0)
203      }
204  }
205  
206  impl Display for Condition {
207      fn fmt(&self, f: F) -> Result {
208          match self {
209              Self::NonZero(r) => write!(f, "NonZero({r})"),
210          }
211      }
212  }
213  
214  impl Display for Program {
215      fn fmt(&self, f: F) -> Result {
216          for (i, (name, function)) in self.functions.iter().enumerate() {
217              if i != 0 {
218                  write!(f, "\n\n")?;
219              }
220              write!(f, "{}", function.with_name(name))?;
221          }
222          Ok(())
223      }
224  }
225  
226  impl Display for Ty {
227      fn fmt(&self, f: F) -> Result {
228          match self {
229              Self::Int(kind) => write!(f, "{kind}"),
230              Self::Pointer(inner) => write!(f, "{inner}^"),
231              Self::Function(params, returns) => {
232                  write!(f, "fn({}){}", Commas(params), ReturnsSpace(returns))
233              }
234              Self::Struct(fields) => {
235                  write!(
236                      f,
237                      "struct({})",
238                      Commas(fields.iter().map(|(name, ty)| format!("{name} {ty}")))
239                  )
240              }
241          }
242      }
243  }
244  
245  impl Display for IntKind {
246      fn fmt(&self, f: F) -> Result {
247          let name = match self {
248              Self::Usize => "usize",
249              Self::U8 => "u8",
250          };
251          write!(f, "{name}")
252      }
253  }
254  
255  impl Display for Cfg {
256      fn fmt(&self, f: &mut Formatter<'_>) -> Result {
257          write!(f, "digraph CFG {{")?;
258          for (index, node) in &self.map {
259              for succ in &node.successors {
260                  write!(f, "    \"n{}\" -> \"n{}\"", index.0, succ.0)?;
261              }
262          }
263          write!(f, "}}")?;
264  
265          write!(f, "digraph DomTree {{")?;
266          for (index, node) in &self.map {
267              let Some(parent) = &node.immediate_dominator else {
268                  continue;
269              };
270              write!(f, "    \"n{}\" -> \"n{}\"", parent.0, index.0)?;
271          }
272          write!(f, "}}")?;
273          Ok(())
274      }
275  }