kir/
print.rs

1use std::{collections::HashMap, fmt::Display, ops::Range};
2
3use indexmap::IndexMap;
4use num::{BigInt, BigUint};
5use slotmap::SecondaryMap;
6use std::hash::Hash;
7
8use crate::ir::*;
9
10fn escape_str(s: &str) -> String {
11  let mut res = String::new();
12  for c in s.chars() {
13    match c {
14      '\\' => res.push_str("\\\\"),
15      '"' => res.push_str("\\\""),
16      '\n' => res.push_str("\\n"),
17      '\r' => res.push_str("\\r"),
18      '\t' => res.push_str("\\t"),
19      _ => res.push(c),
20    }
21  }
22  res
23}
24
25pub struct ValuePrinter<'r> {
26  pub values: &'r ValueMap,
27  pub used: HashMap<String, usize>,
28  pub resolved: SecondaryMap<ValueId, ValueName<'r>>,
29  pub next_value_id: usize,
30}
31
32#[derive(Clone, Copy)]
33pub enum ValueName<'r> {
34  Named(&'r str, usize),
35  Unnamed(usize),
36}
37impl Display for ValueName<'_> {
38  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39    match self {
40      ValueName::Named(name, 0) => write!(f, "{}", name),
41      ValueName::Named(name, n) => write!(f, "{}_{}", name, n),
42      ValueName::Unnamed(n) => write!(f, "{}", n),
43    }
44  }
45}
46
47impl<'r> ValuePrinter<'r> {
48  pub fn new(values: &'r ValueMap) -> Self {
49    ValuePrinter {
50      values,
51      used: HashMap::new(),
52      resolved: SecondaryMap::new(),
53      next_value_id: 0,
54    }
55  }
56  pub fn find_next_name(&mut self, name: &str) -> usize {
57    if let Some(v) = self.used.get(name) {
58      let mut v = *v;
59      while self.used.contains_key(&format!("{}_{}", name, v)) {
60        v += 1;
61      }
62      *self.used.get_mut(name).unwrap() = v;
63      v
64    } else {
65      self.used.insert(name.to_string(), 1);
66      0
67    }
68  }
69  pub fn next_unnamed(&mut self, vid: ValueId) -> ValueName<'r> {
70    // let id = self.next_value_id;
71    // self.next_value_id += 1;
72    // ValueName::Unnamed(id)
73
74    // instead of using next_value_id, just use the vid as the unnamed id
75    // FIXME: this is not a good idea, because the vid might(?) be not unique
76    // TODO: add a checker to ensure the vid is unique
77    ValueName::Unnamed(vid.as_ffi() as usize)
78  }
79  pub fn resolve_name(&mut self, vid: ValueId) -> ValueName<'r> {
80    if let Some(name) = self.resolved.get(vid) {
81      return *name;
82    }
83    let res = match self.values[vid].name {
84      Some(ref name) => {
85        let id = self.find_next_name(name);
86        ValueName::Named(name, id)
87      }
88      None => self.next_unnamed(vid),
89    };
90    self.resolved.insert(vid, res);
91    res
92  }
93}
94
95pub struct Printer<'v> {
96  pub buf: String,
97  pub indent: i32,
98  pub space: bool,
99  pub vp: Option<ValuePrinter<'v>>,
100}
101
102impl<'v> Printer<'v> {
103  pub fn new() -> Self {
104    Printer {
105      buf: String::new(),
106      indent: 0,
107      space: false,
108      vp: None,
109    }
110  }
111  pub fn write_fmt(&mut self, fmt: std::fmt::Arguments) {
112    use std::fmt::Write;
113    if self.space && !fmt.to_string().starts_with(')') {
114      self.buf.push_str(" ");
115      self.space = false;
116    }
117    self.buf.write_fmt(fmt).unwrap();
118    if !fmt.to_string().ends_with('(') {
119      self.space = true;
120      // println!("space, last fmt is {:?}", fmt);
121    } else {
122      // println!("no space, last fmt is {:?}", fmt);
123    }
124  }
125  pub fn newline(&mut self) {
126    self.space = false;
127    self.buf.push_str("\n");
128    for _ in 0..self.indent {
129      self.buf.push_str("  ");
130    }
131  }
132  pub fn indent(&mut self, diff: i32) {
133    self.indent += diff;
134  }
135  pub fn print_unescaped_str(&mut self, s: &str) {
136    self.write_fmt(format_args!("\"{}\"", escape_str(s)));
137  }
138  pub fn print<T: Print>(&mut self, item: &'v T) {
139    item.print(self);
140  }
141  pub fn print_value(&mut self, vid: ValueId) {
142    let vp = self.vp.as_mut().unwrap();
143    let name = vp.resolve_name(vid);
144    let ty = vp.values[vid].ty.clone();
145    match name {
146      ValueName::Named(name, 0) => {
147        self.write_fmt(format_args!("%{}{}", name, type_str(&ty)))
148      }
149      ValueName::Named(name, id) => {
150        self.write_fmt(format_args!("%{}_{}{}", name, id, type_str(&ty)))
151      }
152      ValueName::Unnamed(id) => {
153        self.write_fmt(format_args!("%{}{}", id, type_str(&ty)))
154      }
155    }
156  }
157  pub fn print_list<'r: 'v, T: Print + 'r>(
158    &mut self,
159    kw: &str,
160    left: &str,
161    sep: &str,
162    right: &str,
163    newline: bool,
164    last: bool,
165    list: impl 'r + IntoIterator<Item = &'r T>,
166  ) {
167    let mut first = true;
168    self.write_fmt(format_args!("{}", left));
169    if kw != "" {
170      self.write_fmt(format_args!("{}", kw));
171    }
172    if newline {
173      self.indent(1);
174    } else {
175      self.space = false;
176    }
177    for item in list {
178      if first {
179        first = false;
180      } else {
181        self.buf.push_str(sep);
182      }
183      if newline {
184        self.newline();
185      }
186      item.print(self);
187    }
188    if !first && last {
189      if last {
190        self.buf.push_str(sep);
191      }
192    }
193    if newline {
194      self.indent(-1);
195    }
196    if !first && newline {
197      self.newline();
198    }
199    self.space = false;
200    self.write_fmt(format_args!("{}", right));
201  }
202
203  pub fn print_list_tuple<'r: 'v, K: Print + 'r, D: Print + 'r>(
204    &mut self,
205    kw: &str,
206    left: &str,
207    sep: &str,
208    right: &str,
209    newline: bool,
210    last: bool,
211    list: impl 'r + IntoIterator<Item = (&'r K, &'r D)>,
212  ) {
213    let mut first = true;
214    self.write_fmt(format_args!("{}", left));
215    if kw != "" {
216      self.write_fmt(format_args!("{}", kw));
217    }
218    if newline {
219      self.indent(1);
220    } else {
221      self.space = false;
222    }
223    for (k, d) in list {
224      if first {
225        first = false;
226      } else {
227        self.buf.push_str(sep);
228      }
229      if newline {
230        self.newline();
231      }
232      k.print(self);
233      write!(self, " : ");
234      d.print(self);
235    }
236    if !first && last {
237      if last {
238        self.buf.push_str(sep);
239      }
240    }
241    if newline {
242      self.indent(-1);
243    }
244    if !first && newline {
245      self.newline();
246    }
247    self.space = false;
248    self.write_fmt(format_args!("{}", right));
249  }
250
251  pub fn set_printer(
252    &mut self,
253    printer: Option<ValuePrinter<'v>>,
254  ) -> Option<ValuePrinter<'v>> {
255    std::mem::replace(&mut self.vp, printer)
256  }
257}
258
259pub fn ir_dump(n: &impl Print) -> String {
260  let mut p = Printer::new();
261  p.print(n);
262  p.buf
263}
264pub fn ir_dump_with(values: &ValueMap, n: &impl Print) -> String {
265  let mut p = Printer::new();
266  p.vp = Some(ValuePrinter::new(values));
267  p.print(n);
268  p.buf
269}
270
271pub trait Print {
272  fn print<'p>(&'p self, p: &mut Printer<'p>);
273}
274
275pub trait IRDump: Sized {
276  fn ir_dump(&self) -> String;
277  fn ir_dump_with(&self, values: &ValueMap) -> String;
278}
279impl<T: Print + Sized> IRDump for T {
280  fn ir_dump(&self) -> String {
281    ir_dump(self)
282  }
283  fn ir_dump_with(&self, values: &ValueMap) -> String {
284    ir_dump_with(values, self)
285  }
286}
287
288impl Print for String {
289  fn print<'p>(&'p self, p: &mut Printer<'p>) {
290    p.print_unescaped_str(self);
291  }
292}
293
294macro_rules! impl_print_for_number {
295    ($($ty:ty),*) => {
296        $(
297            impl Print for $ty {
298                fn print<'p>(&'p self, p: &mut Printer<'p>) {
299                    write!(p, "{}", self);
300                }
301            }
302        )*
303    };
304}
305impl_print_for_number!(
306  bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, isize, usize
307);
308impl_print_for_number!(BigUint, BigInt);
309
310impl Print for Type {
311  fn print<'p>(&'p self, p: &mut Printer<'p>) {
312    write!(p, "{}", self.to_string());
313  }
314}
315
316impl Print for ValueId {
317  fn print<'p>(&'p self, p: &mut Printer<'p>) {
318    p.print_value(*self);
319  }
320}
321
322impl<T: Print> Print for Vec<T> {
323  fn print<'p>(&'p self, p: &mut Printer<'p>) {
324    p.print_list("", "(", "", ")", false, false, self.iter());
325  }
326}
327
328impl<T: Print, const N: usize> Print for [T; N] {
329  fn print<'p>(&'p self, p: &mut Printer<'p>) {
330    p.print_list("", "(", "", ")", true, false, self.iter());
331  }
332}
333
334impl<T: Print> Print for Option<T> {
335  fn print<'p>(&'p self, p: &mut Printer<'p>) {
336    // p.print_list("", "(", "", ")", false, false, self.iter())
337    match self {
338      Some(v) => v.print(p),
339      None => write!(p, "_"),
340    }
341  }
342}
343
344impl<T: Print> Print for Box<T> {
345  fn print<'p>(&'p self, p: &mut Printer<'p>) {
346    self.as_ref().print(p);
347  }
348}
349
350impl<T: Print> Print for Range<T> {
351  fn print<'p>(&'p self, p: &mut Printer<'p>) {
352    self.start.print(p);
353    write!(p, "..");
354    self.end.print(p);
355  }
356}
357
358impl<K: Print + Eq + Hash + Clone, D: Print + Clone> Print for IndexMap<K, D> {
359  fn print<'p>(&'p self, p: &mut Printer<'p>) {
360    p.print_list_tuple("", "{", ",", "}", true, false, self);
361  }
362}
363
364fn type_str(ty: &Option<Type>) -> String {
365  match ty {
366    Some(ty) => format!(":{}", ty.to_string()),
367    None => "".to_string(),
368    // None => ":_tbd".to_string(),
369  }
370}
371
372const MAX_DEPTH: usize = 1;
373
374impl Print for json::object::Object {
375  fn print<'p>(&'p self, p: &mut Printer<'p>) {
376    p.write_fmt(format_args!(
377      "{}",
378      json::stringify(
379        json::JsonValue::Object(self.clone()).limit_depth(MAX_DEPTH)
380      )
381    ));
382  }
383}
384
385trait ObjectManipulate {
386  fn limit_depth(self, depth: usize) -> Self;
387}
388
389impl ObjectManipulate for json::JsonValue {
390  fn limit_depth(self, depth: usize) -> Self {
391    match self {
392      json::JsonValue::Object(obj) => {
393        if depth > 0 {
394          let mut res = json::object::Object::new();
395          for (k, v) in obj.iter() {
396            res.insert(k, v.clone().limit_depth(depth - 1));
397          }
398          json::JsonValue::Object(res)
399        } else {
400          json::JsonValue::String("...".to_string())
401        }
402      }
403      _ => self,
404    }
405  }
406}