koopa/back/
generator.rs

1//! Koopa IR generator ([`Generator`]), name manager ([`NameManager`]) and
2//! Koopa IR visitor trait ([`Visitor`]) related implementations.
3//!
4//! The Koopa IR generator converts in-memory Koopa IR programs into
5//! other forms by using IR visitors. IR visitors can use name manager
6//! to generate function/basic block/value names when visiting IR programs.
7
8use crate::ir::entities::{BasicBlockData, FunctionData, Program, ValueData};
9use std::borrow::Borrow;
10use std::collections::{HashMap, HashSet};
11use std::fs::File;
12use std::io::{Result, Write};
13use std::num::NonZeroUsize;
14use std::path::Path;
15use std::rc::Rc;
16
17/// A manager for storing names and allocating unique temporary
18/// names of global values, functions, basic blocks and local values.
19#[derive(Default)]
20pub struct NameManager {
21  next_id: usize,
22  cur_scope: ScopeKind,
23  prefix: Prefix,
24  global_names: HashSet<StringRc>,
25  bb_names: HashSet<StringRc>,
26  global_vars: HashMap<*const ValueData, Rc<String>>,
27  funcs: HashMap<*const FunctionData, Rc<String>>,
28  bbs: HashMap<*const BasicBlockData, Rc<String>>,
29  values: HashMap<*const ValueData, Rc<String>>,
30}
31
32impl NameManager {
33  /// Creates a new name manager.
34  pub fn new() -> Self {
35    Self::default()
36  }
37
38  /// Enters the function scope.
39  /// Call this method when generating basic blocks and local values.
40  ///
41  /// # Panics
42  ///
43  /// Panics when the current scope is already function scope.
44  pub fn enter_func_scope(&mut self) {
45    assert!(
46      matches!(self.cur_scope, ScopeKind::Global),
47      "already in function scope"
48    );
49    self.cur_scope = ScopeKind::Function;
50    self.values.clear();
51  }
52
53  /// Exits the function scope.
54  /// Call this method when quitting from function generation.
55  ///
56  /// # Panics
57  ///
58  /// Panics when the current scope is not function scope.
59  pub fn exit_func_scope(&mut self) {
60    assert!(
61      matches!(self.cur_scope, ScopeKind::Function),
62      "not in function scope"
63    );
64    self.cur_scope = ScopeKind::Global;
65    self.bb_names.clear();
66    self.bbs.clear();
67    for name in self.values.values() {
68      self.global_names.remove(name);
69    }
70  }
71
72  /// Sets the prefix of generated names.
73  pub fn set_prefix(&mut self, prefix: Prefix) {
74    self.prefix = prefix;
75  }
76
77  /// Returns the name of the given function.
78  pub fn func_name(&mut self, func: &FunctionData) -> Rc<String> {
79    let ptr: *const FunctionData = func;
80    if let Some(name) = self.funcs.get(&ptr) {
81      name.clone()
82    } else {
83      let name = self.next_name_str(func.name(), |s| &mut s.global_names);
84      self.funcs.insert(ptr, name);
85      self.funcs[&ptr].clone()
86    }
87  }
88
89  /// Returns the name of the given basic block.
90  pub fn bb_name(&mut self, bb: &BasicBlockData) -> Rc<String> {
91    let ptr: *const BasicBlockData = bb;
92    if let Some(name) = self.bbs.get(&ptr) {
93      name.clone()
94    } else {
95      let name = self.next_name(bb.name(), |s| &mut s.bb_names);
96      self.bbs.insert(ptr, name);
97      self.bbs[&ptr].clone()
98    }
99  }
100
101  /// Returns the name of the given value.
102  ///
103  /// # Panics
104  ///
105  /// Panics if the given value is a constant.
106  pub fn value_name(&mut self, value: &ValueData) -> Rc<String> {
107    assert!(!value.kind().is_const(), "can not name constants");
108    if value.kind().is_global_alloc() {
109      self.value_name_impl(value, |s| &mut s.global_vars)
110    } else {
111      self.value_name_impl(value, |s| &mut s.values)
112    }
113  }
114
115  fn value_name_impl<F>(&mut self, value: &ValueData, value_set: F) -> Rc<String>
116  where
117    F: for<'a> Fn(&'a mut Self) -> &'a mut HashMap<*const ValueData, Rc<String>>,
118  {
119    let ptr: *const ValueData = value;
120    if let Some(name) = value_set(self).get(&ptr) {
121      name.clone()
122    } else {
123      let name = self.next_name(value.name(), |s| &mut s.global_names);
124      let values = value_set(self);
125      values.insert(ptr, name);
126      values[&ptr].clone()
127    }
128  }
129
130  /// Returns a temporary value name.
131  pub fn temp_value_name(&mut self) -> Rc<String> {
132    self.next_name(&None, |s| &mut s.global_names)
133  }
134
135  /// Generates the next name by the given `Option<String>`
136  /// and stores it to the given name set.
137  fn next_name<F>(&mut self, name: &Option<String>, name_set: F) -> Rc<String>
138  where
139    F: for<'a> FnOnce(&'a mut Self) -> &'a mut HashSet<StringRc>,
140  {
141    // check if there is a name
142    if let Some(name) = name {
143      self.next_name_str(name, name_set)
144    } else {
145      // generate a temporary name
146      let name = self.prefix.temp_name(self.next_id);
147      self.next_id += 1;
148      let names = name_set(self);
149      names.insert(name.clone().into());
150      names.get(&name).unwrap().to_rc()
151    }
152  }
153
154  /// Generates the next name by the given string
155  /// and stores it to the given name set.
156  fn next_name_str<F>(&mut self, name: &str, name_set: F) -> Rc<String>
157  where
158    F: for<'a> FnOnce(&'a mut Self) -> &'a mut HashSet<StringRc>,
159  {
160    let name = self.prefix.name(name);
161    let names = name_set(self);
162    // check for duplicate names
163    if !names.contains(&name) {
164      names.insert(name.clone().into());
165      names.get(&name).unwrap().to_rc()
166    } else {
167      // generate a new name
168      for id in 0.. {
169        let new_name = format!("{}_{}", name, id);
170        if !names.contains(&new_name) {
171          names.insert(new_name.clone().into());
172          return names.get(&new_name).unwrap().to_rc();
173        }
174      }
175      unreachable!()
176    }
177  }
178}
179
180/// Kind of scope.
181#[derive(Clone, Copy, Default)]
182enum ScopeKind {
183  #[default]
184  Global,
185  Function,
186}
187
188/// Prefix of name.
189#[derive(Default)]
190pub enum Prefix {
191  /// Default prefix,
192  /// named variables start with '@' and
193  /// temporary variables start with '%'.
194  #[default]
195  Default,
196  /// Custom prefix,
197  /// named variables start with `named` and
198  /// temporary variables start with `temp`.
199  Custom {
200    /// Named variable prefix.
201    named: String,
202    /// Temporary variable prefix.
203    temp: String,
204    /// Maximum length of named variable.
205    /// Truncates the variable name if its length is too long.
206    /// `None` if no limit.
207    max_len: Option<NonZeroUsize>,
208  },
209}
210
211impl Prefix {
212  /// Returns the name according to the prefix setting.
213  fn name(&self, name: &str) -> String {
214    match self {
215      Prefix::Default => name.into(),
216      Prefix::Custom {
217        named,
218        temp,
219        max_len,
220      } => {
221        let mut name = if let Some(name) = name.strip_prefix('@') {
222          format!("{}{}", named, name)
223        } else {
224          format!("{}{}", temp, &name[1..])
225        };
226        if let Some(max_len) = max_len {
227          name.truncate(max_len.get());
228        }
229        name
230      }
231    }
232  }
233
234  /// Returns a temp name by the given id.
235  fn temp_name(&self, id: usize) -> String {
236    match self {
237      Prefix::Default => format!("%{}", id),
238      Prefix::Custom { temp, .. } => format!("{}{}", temp, id),
239    }
240  }
241}
242
243/// `Rc<String>` that implements `Borrow<str>`.
244#[derive(Clone, PartialEq, Eq, Hash)]
245struct StringRc(Rc<String>);
246
247impl StringRc {
248  fn to_rc(&self) -> Rc<String> {
249    self.0.clone()
250  }
251}
252
253impl From<String> for StringRc {
254  fn from(s: String) -> Self {
255    Self(Rc::new(s))
256  }
257}
258
259impl From<&str> for StringRc {
260  fn from(s: &str) -> Self {
261    Self(Rc::new(s.into()))
262  }
263}
264
265impl Borrow<Rc<String>> for StringRc {
266  fn borrow(&self) -> &Rc<String> {
267    &self.0
268  }
269}
270
271impl Borrow<String> for StringRc {
272  fn borrow(&self) -> &String {
273    self.0.as_ref()
274  }
275}
276
277impl Borrow<str> for StringRc {
278  fn borrow(&self) -> &str {
279    self.0.as_ref()
280  }
281}
282
283/// A generator for traversing and generating Koopa IR into other forms.
284pub struct Generator<W: Write, V: Visitor<W>> {
285  writer: W,
286  visitor: V,
287  name_man: NameManager,
288}
289
290impl<W: Write, V: Visitor<W>> Generator<W, V> {
291  /// Creates a new generator.
292  pub fn new(writer: W) -> Self
293  where
294    V: Default,
295  {
296    Self {
297      writer,
298      visitor: V::default(),
299      name_man: NameManager::new(),
300    }
301  }
302
303  /// Creates a new generator with the given visitor.
304  pub fn with_visitor(writer: W, visitor: V) -> Self {
305    Self {
306      writer,
307      visitor,
308      name_man: NameManager::new(),
309    }
310  }
311
312  /// Consumes and returns the writer inside of the current generator.
313  pub fn writer(self) -> W {
314    self.writer
315  }
316
317  /// Generates on the given Koopa IR program.
318  pub fn generate_on(&mut self, program: &Program) -> Result<V::Output> {
319    self
320      .visitor
321      .visit(&mut self.writer, &mut self.name_man, program)
322  }
323}
324
325impl<V: Visitor<File>> Generator<File, V> {
326  /// Creates a new generator from the given path.
327  pub fn from_path<P>(path: P) -> Result<Self>
328  where
329    V: Default,
330    P: AsRef<Path>,
331  {
332    File::create(path).map(Generator::new)
333  }
334}
335
336/// A visitor trait for all Koopa IR visitors.
337pub trait Visitor<W: Write> {
338  /// The output type of all visitor methods.
339  type Output;
340
341  /// Visits the given Koopa IR program.
342  fn visit(&mut self, w: &mut W, nm: &mut NameManager, program: &Program) -> Result<Self::Output>;
343}
koopa/back/generator.rs

koopa/back/
generator.rs
