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}