fugue_ir/disassembly/
partmap.rs

1use std::borrow::Borrow;
2use std::collections::BTreeMap as Map;
3use std::collections::btree_map::{Range, RangeMut};
4use std::ops::Bound::Excluded;
5use std::ops::RangeBounds;
6
7#[derive(Debug, Clone)]
8pub enum BoundKind<'a, K, V> {
9    None(&'a V),
10    Lower(&'a K, &'a V),
11    Upper(&'a K, &'a V),
12    Both(&'a K, &'a K, &'a V),
13}
14
15impl<'a, K, V> BoundKind<'a, K, V> {
16    pub fn lower(&self) -> Option<&'a K> {
17        match self {
18            Self::None(_) | Self::Upper(_, _) => None,
19            Self::Lower(k, _) | Self::Both(k, _, _) => Some(k)
20        }
21    }
22
23    pub fn upper(&self) -> Option<&'a K> {
24        match self {
25            Self::None(_) | Self::Lower(_, _) => None,
26            Self::Upper(k, _) | Self::Both(_, k, _) => Some(k)
27        }
28    }
29
30    pub fn value(&self) -> &'a V {
31        match self {
32            Self::None(v) | Self::Lower(_, v) |
33            Self::Upper(_, v) | Self::Both(_, _, v) => v,
34        }
35    }
36}
37
38#[derive(Debug, Clone)]
39#[derive(serde::Deserialize, serde::Serialize)]
40pub struct PartMap<K: Ord, V> {
41    mapping: Map<K, V>,
42    default: V,
43}
44
45impl<K, V> PartMap<K, V>
46where K: Clone + Ord,
47      V: Clone {
48
49    pub fn new(default: V) -> Self {
50        Self {
51            mapping: Map::new(),
52            default,
53        }
54    }
55
56    pub fn default_value(&self) -> &V {
57        &self.default
58    }
59
60    pub fn default_value_mut(&mut self) -> &mut V {
61        &mut self.default
62    }
63
64    pub fn is_empty(&self) -> bool {
65        self.mapping.is_empty()
66    }
67
68    pub fn bounds<'a>(&self, point: &'a K) -> BoundKind<K, V> {
69        let lb = self.mapping.range(..=point).rev().next();
70        let ub = self.mapping
71            .range(point..)
72            .find_map(|(k, _)| if k > point { Some(k) } else { None });
73
74        match (lb, ub) {
75            (None, None) => BoundKind::None(self.default_value()),
76            (Some((l, v)), None) => BoundKind::Lower(l, v),
77            (None, Some(u)) => BoundKind::Upper(u, self.default_value()),
78            (Some((l, v)), Some(u)) => BoundKind::Both(l, u, v),
79        }
80    }
81
82    pub fn clear(&mut self) {
83        self.mapping = Map::new();
84    }
85
86    pub fn clear_range<'a>(&mut self, start: &'a K, end: &'a K) -> &mut V {
87        self.split(start);
88        self.split(end);
89
90        let keys = self.mapping
91            .range((Excluded(start), Excluded(end)))
92            .map(|(k, _)| k.clone())
93            .collect::<Vec<K>>();
94
95        for key in keys {
96            self.mapping.remove(&key);
97        }
98
99        self.mapping.get_mut(start).unwrap()
100    }
101
102    pub fn get<'a>(&self, point: &'a K) -> Option<&V> {
103        self.mapping.range(..=point).rev().next().map(|(_, v)| v)
104    }
105
106    pub fn get_or_default<'a>(&self, point: &'a K) -> &V {
107        self.get(point).unwrap_or(self.default_value())
108    }
109
110    pub fn iter(&self) -> impl Iterator<Item=(&K, &V)> {
111        self.mapping.iter()
112    }
113
114    pub fn begin<'a>(&'a self, point: &'a K) -> Range<'a, K, V> {
115        self.range(point..)
116    }
117
118    pub fn begin_mut<'a>(&'a mut self, point: &'a K) -> RangeMut<'a, K, V> {
119        self.range_mut(point..)
120    }
121
122    pub fn range<'a, T, R>(&'a self, range: R) -> Range<'a, K, V>
123    where K: Borrow<T> + 'a,
124          R: RangeBounds<T> + 'a,
125          T: Ord + ?Sized + 'a, {
126        self.mapping.range(range)
127    }
128
129    pub fn range_mut<'a, T, R>(&'a mut self, range: R) -> RangeMut<'a, K, V>
130    where K: Borrow<T> + 'a,
131          R: RangeBounds<T> + 'a,
132          T: Ord + ?Sized + 'a, {
133        self.mapping.range_mut(range)
134    }
135
136    pub fn split<'a>(&'a mut self, at: &'a K) -> &'a V {
137        self.split_mut(at)
138    }
139
140    pub fn split_mut<'a>(&'a mut self, at: &'a K) -> &'a mut V {
141        let value = if let Some(point) = self.mapping.range(..=at).rev().next() {
142            if point.0 == at {
143                None
144            } else {
145                Some(point.1.clone())
146            }
147        } else {
148            Some(self.default.clone())
149        };
150
151        if let Some(value) = value {
152            self.mapping.entry(at.clone()).or_insert(value)
153        } else {
154            self.mapping.get_mut(&at).unwrap()
155        }
156    }
157}
158
159#[cfg(test)]
160mod test {
161    use super::*;
162
163    #[test]
164    fn simple_partmap() {
165        let mut map = PartMap::<isize, usize>::new(0);
166
167        *map.split_mut(&5) = 5;
168        *map.split_mut(&2) = 2;
169        *map.split_mut(&3) = 4;
170        *map.split_mut(&3) = 3;
171
172        assert_eq!(map.get(&6), Some(&5));
173        assert_eq!(map.get(&8), Some(&5));
174        assert_eq!(map.get(&4), Some(&3));
175        assert_eq!(map.get(&1), None);
176    }
177}