griffin_core/uplc/optimize/
interner.rs

1use alloc::{rc::Rc, string::String};
2use hashbrown::HashMap;
3
4use crate::uplc::ast::{Name, Program, Term, Unique};
5
6#[derive(Eq, Hash, PartialEq, Clone)]
7pub struct InternKey {
8    name: String,
9    previous_unique: Unique,
10}
11
12pub struct CodeGenInterner {
13    identifiers: HashMap<InternKey, Unique>,
14    current: Unique,
15}
16
17impl Default for CodeGenInterner {
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23/// Interner that uses previous uniques to prevent future unique collisions
24/// when performing optimizations
25impl CodeGenInterner {
26    pub fn new() -> Self {
27        Self {
28            identifiers: HashMap::new(),
29            current: Unique::new(0),
30        }
31    }
32
33    pub fn program(&mut self, program: &mut Program<Name>) {
34        self.term(&mut program.term);
35    }
36
37    pub fn term(&mut self, term: &mut Term<Name>) {
38        match term {
39            Term::Var(name) => {
40                let name = Rc::make_mut(name);
41                name.unique = self.intern(name.text.clone(), name.unique);
42            }
43            Term::Delay(term) => self.term(Rc::make_mut(term)),
44            Term::Lambda {
45                parameter_name,
46                body,
47            } => {
48                let parameter_name = Rc::make_mut(parameter_name);
49                parameter_name.unique =
50                    self.intern(parameter_name.text.clone(), parameter_name.unique);
51                self.term(Rc::make_mut(body));
52            }
53            Term::Apply { function, argument } => {
54                self.term(Rc::make_mut(function));
55                self.term(Rc::make_mut(argument));
56            }
57            Term::Constant(_) => (),
58            Term::Force(term) => self.term(Rc::make_mut(term)),
59            Term::Error => (),
60            Term::Builtin(_) => (),
61            Term::Constr { fields, .. } => {
62                for field in fields {
63                    self.term(field);
64                }
65            }
66            Term::Case { constr, branches } => {
67                self.term(Rc::make_mut(constr));
68
69                for branch in branches {
70                    self.term(branch);
71                }
72            }
73        }
74    }
75
76    pub fn intern(&mut self, text: String, previous_unique: Unique) -> Unique {
77        let key = InternKey {
78            name: text,
79            previous_unique,
80        };
81
82        if let Some(u) = self.identifiers.get(&key) {
83            *u
84        } else {
85            let unique = self.current;
86
87            self.identifiers.insert(key, self.current_unique());
88
89            self.current.increment();
90
91            unique
92        }
93    }
94
95    pub fn current_unique(&self) -> Unique {
96        self.current
97    }
98}