ryna/structures/
variable_map.rs

1use rustc_hash::{FxHashMap, FxHashSet};
2
3use crate::types::Type;
4
5#[derive(Clone)]
6struct VariableContext {
7    is_base: bool,
8    names: FxHashMap<String, usize>,
9    order: Vec<String>,                 // In declaration order
10    vars: Vec<(usize, Type)>            // In declaration order
11}
12
13impl VariableContext {
14    fn new(is_base: bool) -> Self {
15        Self { is_base, vars: vec!(), order: vec!(), names: FxHashMap::default() }
16    }
17
18    fn contains(&self, name: &String) -> bool {
19        self.names.contains_key(name)
20    }
21
22    fn get(&self, name: &String) -> Option<&(usize, Type)> {
23        match self.names.get(name) {
24            Some(idx) => Some(&self.vars[*idx]),
25            None => None,
26        }
27    }
28
29    fn define_var(&mut self, name: String, idx: usize, t: Type) {
30        if !self.contains(&name) {
31            self.names.entry(name.clone()).or_insert(self.vars.len());
32            self.order.push(name);
33            self.vars.push((idx, t));
34        }
35    }
36}
37
38#[derive(Clone)]
39pub struct VariableMap {
40    counter: usize,
41    contexts: Vec<VariableContext>
42}
43
44impl Default for VariableMap {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50impl VariableMap {
51    pub fn new() -> Self {
52        VariableMap { contexts: vec!(), counter: 0 }
53    }
54
55    pub fn depth(&self) -> usize {
56        self.contexts.len()
57    }
58
59    pub fn add_context(&mut self, is_base: bool) {
60        self.contexts.push(VariableContext::new(is_base));
61    } 
62
63    pub fn remove_context(&mut self) {
64        self.contexts.pop().unwrap();
65    }
66
67    pub fn define_var(&mut self, name: String, idx: usize, t: Type) -> usize {
68        self.contexts.last_mut().unwrap().define_var(name, idx, t);
69
70        self.contexts.len() - 1 // Will never underflow, otherwise it would panic in the previous line
71    }
72
73    pub fn count_up(&mut self) -> usize {
74        self.counter += 1;
75        self.counter
76    }
77
78    pub fn is_var_defined(&mut self, name: &String) -> bool {
79        for ctx in self.contexts.iter().rev() {
80            if ctx.contains(name) {
81                return true;
82            }
83        }
84
85        false
86    }
87
88    pub fn is_var_defined_in_last_ctx(&mut self, name: &String) -> bool {
89        self.contexts.last().unwrap().contains(name)
90    }
91
92    pub fn get_var(&self, name: &String) -> Option<(usize, &(usize, Type))> {
93        for (i, ctx) in self.contexts.iter().enumerate().rev() {
94            if let Some(v) = ctx.get(name) {
95                return Some((i, v));
96            }
97        }
98
99        None
100    }
101
102    pub fn num_vars(&self) -> usize {
103        self.contexts.iter().map(|i| i.vars.len()).sum()
104    }
105
106    pub fn for_each_last_ctx<T: FnMut(usize, &String, &Type)>(&self, mut f: T) {
107        let last_ctx = self.contexts.last().unwrap();
108
109        for n in last_ctx.order.iter().rev() {
110            let (i, t) = &last_ctx.vars[*last_ctx.names.get(n).unwrap()];
111            f(*i, n, t);
112        }
113    }
114
115    pub fn for_each_until_base<T: FnMut(usize, &String, &Type)>(&self, mut f: T) {
116        for ctx in self.contexts.iter().rev() {
117            for n in ctx.order.iter().rev() {
118                let (i, t) = &ctx.vars[*ctx.names.get(n).unwrap()];
119                f(*i, n, t);
120            }
121
122            if ctx.is_base {
123                break;
124            }
125        }
126    }
127
128    pub fn var_names(&self) -> FxHashSet<&String> {
129        self.contexts.iter().rev().flat_map(|i| i.names.keys()).collect()
130    }
131}