layout/adt/
map.rs

1//! This module implements the scoped map.
2
3use std::cmp::Eq;
4use std::collections::HashMap;
5use std::hash::Hash;
6
7/// Scoped map that supports inserting and removing lots of key-val pairs
8/// at once.
9#[derive(Debug)]
10pub struct ScopedMap<K, V> {
11    stack: Vec<Vec<(K, V)>>,
12}
13
14impl<K: PartialEq + Clone + Hash + Eq, V: Clone> Default for ScopedMap<K, V> {
15    fn default() -> Self {
16        Self::new()
17    }
18}
19
20impl<K: PartialEq + Clone + Hash + Eq, V: Clone> ScopedMap<K, V> {
21    pub fn new() -> Self {
22        ScopedMap { stack: Vec::new() }
23    }
24
25    pub fn push(&mut self) {
26        self.stack.push(Vec::new());
27    }
28
29    pub fn pop(&mut self) {
30        if !self.is_empty() {
31            self.stack.pop();
32        }
33    }
34
35    pub fn len(&self) -> usize {
36        self.stack.len()
37    }
38
39    pub fn is_empty(&self) -> bool {
40        self.stack.is_empty()
41    }
42
43    pub fn insert(&mut self, key: &K, val: &V) {
44        assert!(!self.is_empty());
45        let scope = self.stack.last_mut().unwrap();
46        for pair in scope {
47            if pair.0 == *key {
48                pair.1 = val.clone();
49                return;
50            }
51        }
52
53        self.stack
54            .last_mut()
55            .unwrap()
56            .push((key.clone(), val.clone()));
57    }
58
59    pub fn flatten(&self) -> HashMap<K, V> {
60        let mut map: HashMap<K, V> = HashMap::new();
61
62        for scope in self.stack.iter() {
63            for pair in scope.iter() {
64                map.insert(pair.0.clone(), pair.1.clone());
65            }
66        }
67
68        map
69    }
70
71    pub fn get(&self, key: &K) -> Option<V> {
72        // For each scope, in reverse:
73        for scope in self.stack.iter().rev() {
74            for pair in scope {
75                if pair.0 == *key {
76                    return Option::Some(pair.1.clone());
77                }
78            }
79        }
80        Option::None
81    }
82
83    pub fn has(&self, key: &K) -> bool {
84        matches!(self.get(key), Option::Some(_))
85    }
86}
87
88#[test]
89fn test_scoped_map() {
90    let mut map: ScopedMap<usize, usize> = ScopedMap::new();
91
92    assert!(map.is_empty());
93    map.push();
94    assert_eq!(map.len(), 1);
95
96    map.insert(&1, &1);
97    map.insert(&2, &2);
98    map.insert(&3, &3);
99
100    assert_eq!(map.get(&1).unwrap(), 1);
101    assert_eq!(map.get(&2).unwrap(), 2);
102    assert_eq!(map.get(&3).unwrap(), 3);
103
104    assert!(map.has(&1));
105    assert!(map.has(&2));
106    assert!(map.has(&3));
107
108    map.push();
109
110    assert!(map.has(&1));
111    assert!(map.has(&2));
112    assert!(map.has(&3));
113
114    map.insert(&1, &4);
115    map.insert(&2, &5);
116    map.insert(&3, &6);
117
118    assert_eq!(map.get(&1).unwrap(), 4);
119    assert_eq!(map.get(&2).unwrap(), 5);
120    assert_eq!(map.get(&3).unwrap(), 6);
121
122    map.pop();
123    assert_eq!(map.get(&1).unwrap(), 1);
124    assert_eq!(map.get(&2).unwrap(), 2);
125    assert_eq!(map.get(&3).unwrap(), 3);
126    map.pop();
127    assert!(!map.has(&1));
128    assert!(!map.has(&2));
129    assert!(!map.has(&3));
130}
131
132#[test]
133fn test_scoped_map2() {
134    let mut map: ScopedMap<usize, usize> = ScopedMap::new();
135    map.push();
136    map.insert(&1, &1);
137    map.push();
138    map.insert(&1, &2);
139    map.insert(&2, &3);
140    map.push();
141
142    let flat = map.flatten();
143    assert!(flat.contains_key(&1));
144    assert!(flat.contains_key(&2));
145    assert!(!flat.contains_key(&3));
146    assert_eq!(*flat.get(&1).unwrap(), 2);
147}