1use std::collections::HashMap;
3use std::hash::Hash;
4use std::iter::FromIterator;
5use std::rc::Rc;
6
7use crate::metrics::{increment, sample};
8
9#[derive(Debug, PartialEq)]
25pub struct Environment<K: Hash + Eq, V: PartialEq> {
26 current: Rc<HashMap<K, V>>,
27 previous: Option<Rc<Environment<K, V>>>,
28}
29
30impl<K: Hash + Eq, V: PartialEq> Clone for Environment<K, V> {
31 fn clone(&self) -> Self {
34 increment!("Environment::clone");
35 Self {
36 current: self.current.clone(),
37 previous: self.previous.clone(),
38 }
39 }
40}
41
42impl<K: Hash + Eq, V: PartialEq> Default for Environment<K, V> {
43 fn default() -> Self {
44 Self {
45 current: Rc::new(HashMap::new()),
46 previous: None,
47 }
48 }
49}
50
51impl<K: Hash + Eq, V: PartialEq> Environment<K, V> {
52 pub fn new() -> Self {
54 Self::default()
55 }
56
57 pub fn insert(&mut self, key: K, value: V) {
59 increment!("Environment::insert");
60 match Rc::get_mut(&mut self.current) {
61 Some(cur) => {
62 cur.insert(key, value);
63 }
64 None => {
65 let mut new = HashMap::new();
66 new.insert(key, value);
67 self.set_current(new);
68 }
69 }
70 }
71
72 pub fn get(&self, key: &K) -> Option<&V> {
74 increment!("Environment::get");
75
76 let mut layer_count = 0;
77 #[allow(clippy::let_and_return)]
79 let result = self.iter_layers().find_map(|hmap| {
80 sample!("Environment.hashmap_size_get", hmap.len() as f64);
81 layer_count += 1;
82 hmap.get(key)
83 });
84
85 sample!("Environment.get_layers_traversed", layer_count as f64);
86 result
87 }
88
89 pub fn iter_layers(&self) -> EnvLayerIter<'_, K, V> {
92 EnvLayerIter { env: Some(self) }
93 }
94
95 pub fn iter_elems(&self) -> impl Iterator<Item = (&K, &V)> {
100 let env: Vec<&HashMap<K, V>> = self.iter_layers().collect();
101 env.into_iter().rev().flatten()
102 }
103
104 pub fn ptr_eq(this: &Self, that: &Self) -> bool {
107 let prev_layers_eq = match (&this.previous, &that.previous) {
108 (Some(ptr_this), Some(ptr_that)) => Rc::ptr_eq(ptr_this, ptr_that),
109 (None, None) => true,
110 _ => false,
111 };
112
113 let curr_layers_eq = (this.current.is_empty() && that.current.is_empty())
114 || Rc::ptr_eq(&this.current, &that.current);
115
116 prev_layers_eq && curr_layers_eq
117 }
118
119 pub fn is_empty(&self) -> bool {
122 self.current.is_empty() && self.previous.is_none()
123 }
124
125 fn set_current(&mut self, new_current: HashMap<K, V>) {
128 let old_current = std::mem::replace(&mut self.current, Rc::new(new_current));
129 if !old_current.is_empty() {
130 self.previous = Some(Rc::new(Environment {
131 current: old_current,
132 previous: self.previous.clone(),
133 }));
134 }
135 }
136}
137
138impl<K: Hash + Eq, V: PartialEq> FromIterator<(K, V)> for Environment<K, V> {
139 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
140 Self {
141 current: Rc::new(HashMap::from_iter(iter)),
142 previous: None,
143 }
144 }
145}
146
147impl<K: Hash + Eq, V: PartialEq> Extend<(K, V)> for Environment<K, V> {
148 fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
149 match Rc::get_mut(&mut self.current) {
152 Some(current) => current.extend(iter),
153 None => {
154 self.set_current(HashMap::from_iter(iter));
155 }
156 }
157 }
158}
159
160pub struct EnvLayerIter<'a, K: 'a + Hash + Eq, V: 'a + PartialEq> {
167 env: Option<&'a Environment<K, V>>,
168}
169
170impl<'a, K: 'a + Hash + Eq, V: 'a + PartialEq> Iterator for EnvLayerIter<'a, K, V> {
171 type Item = &'a HashMap<K, V>;
172
173 fn next(&mut self) -> Option<Self::Item> {
174 self.env.map(|env| {
175 let res = &*env.current;
176 self.env = env.previous.as_deref();
177 res
178 })
179 }
180}
181
182#[cfg(test)]
183mod tests {
184 use super::*;
185
186 impl<K: Hash + Eq, V: PartialEq> Environment<K, V> {
187 pub fn depth(&self) -> usize {
188 1 + self.previous.as_ref().map_or(0, |p| p.depth())
189 }
190 }
191
192 #[test]
193 fn test_env_base() {
194 let mut env_base = Environment::new();
195 env_base.insert(1, 'a');
196 assert_eq!(env_base.get(&1), Some(&'a'));
197 assert_eq!(env_base.get(&5), None);
198 assert_eq!(env_base.depth(), 1);
199 }
200
201 #[test]
202 fn test_clone() {
203 let mut env_base = Environment::new();
204 env_base.insert(1, 'a');
205
206 let mut env2 = env_base.clone();
207 env2.insert(2, 'b');
208 assert_eq!(env2.get(&1), Some(&'a'));
209 assert_eq!(env2.get(&2), Some(&'b'));
210 env_base.insert(3, 'c');
211 assert_eq!(env2.get(&3), None);
212 assert_eq!(env_base.get(&3), Some(&'c'));
213 env_base.insert(2, 'z');
214 assert_eq!(env_base.get(&2), Some(&'z'));
215
216 assert_eq!(env_base.depth(), 2);
217 assert_eq!(env2.depth(), 2);
218 }
219
220 #[test]
221 fn test_deepness() {
222 let mut env_base = Environment::<u8, char>::new();
223 assert_eq!(env_base.depth(), 1);
224
225 let mut env2 = env_base.clone();
226 assert_eq!(env_base.depth(), 1);
227 assert_eq!(env2.depth(), 1);
228
229 env2.insert(1, 'a');
230 let env3 = env2.clone();
231 assert_eq!(env_base.depth(), 1);
232 assert_eq!(env2.depth(), 1);
233 assert_eq!(env3.depth(), 1);
234
235 env2.insert(1, 'b');
236 assert_eq!(env2.depth(), 2);
237 assert_eq!(env3.depth(), 1);
238
239 let mut env4 = env2.clone();
240 assert_eq!(env_base.depth(), 1);
241 assert_eq!(env4.depth(), 2);
242 env4.insert(1, 'y');
243 assert_eq!(env4.depth(), 3);
244
245 env_base.insert(1, 'z');
246 assert_eq!(env_base.depth(), 1);
247 assert_eq!(env2.depth(), 2);
248 assert_eq!(env3.depth(), 1);
249 assert_eq!(env4.depth(), 3);
250 }
251
252 #[test]
253 fn test_iter_layer() {
254 let mut env_base = Environment::<u8, char>::new();
255 assert_eq!(env_base.iter_layers().count(), 1);
256 env_base.insert(1, 'a');
257 assert_eq!(env_base.iter_layers().count(), 1);
258 assert_eq!(env_base.iter_layers().next().unwrap().get(&1), Some(&'a'));
259 assert_eq!(env_base.iter_layers().nth(1), None);
260 let _ = env_base.clone();
261 assert_eq!(env_base.iter_layers().count(), 1);
262 env_base.insert(2, 'b');
263 assert_eq!(env_base.iter_layers().count(), 1);
264
265 let _env2 = env_base.clone();
266 env_base.insert(3, 'c');
267 assert_eq!(env_base.iter_layers().count(), 2);
268 let mut iter = env_base.iter_layers();
269 let map1 = iter.next().unwrap();
270 assert_eq!(map1.get(&3), Some(&'c'));
271 assert_eq!(map1.get(&2), None);
272 assert_eq!(map1.get(&1), None);
273 let map2 = iter.next().unwrap();
274 assert_eq!(map2.get(&3), None);
275 assert_eq!(map2.get(&2), Some(&'b'));
276 assert_eq!(map2.get(&1), Some(&'a'));
277 assert!(iter.next().is_none());
278 }
279
280 #[test]
281 fn test_iter_elem() {
282 let env_base = Environment::<u8, char>::new();
283 assert!(env_base.iter_elems().next().is_none());
284
285 let mut env_base = Environment::new();
286 env_base.insert(1, 'a');
287 env_base.insert(2, 'b');
288 let mut iter_elems = env_base.iter_elems();
289 assert!(iter_elems.next().is_some());
290 assert!(iter_elems.next().is_some());
291 assert!(iter_elems.next().is_none());
292 assert!({
293 let mut vec: Vec<_> = env_base.iter_elems().collect();
294 vec.sort_unstable();
295 vec == vec![(&1, &'a'), (&2, &'b')]
296 });
297 drop(iter_elems);
298
299 let mut env2 = env_base.clone();
300 env_base.insert(1, 'z');
301 env2.insert(1, 'y');
302 assert_eq!(env_base.iter_elems().count(), 3);
303 assert_eq!(env2.iter_elems().count(), 3);
304 assert_eq!(env_base.iter_elems().nth(2), Some((&1, &'z')));
305 assert_eq!(env2.iter_elems().nth(2), Some((&1, &'y')));
306
307 let mut iter_elems_base = env_base.iter_elems();
308 let _ = iter_elems_base.next();
309 let mut iter_elems2 = env2.iter_elems();
310 let _ = iter_elems_base.next();
311 let _ = iter_elems2.next();
312
313 let mut env_base = Environment::new();
314 env_base.insert(1, 'a');
315 env_base.insert(2, 'b');
316 let _ = env_base.clone();
317 env_base.insert(1, 'z');
318 env_base.insert(3, 'c');
319 let hmap: HashMap<_, _> = env_base.iter_elems().collect();
320 assert_eq!(hmap.len(), 3);
321 assert_eq!(hmap[&1], &'z');
322 assert_eq!(hmap[&2], &'b');
323 assert_eq!(hmap[&3], &'c');
324 }
325}