range_tree/
lib.rs

1use std::collections::BTreeMap;
2use std::ops::{Range, RangeBounds, Bound};
3
4#[derive(Debug)]
5pub struct RangeTree<V: Value> {
6    inner: BTreeMap<V::K, Slot<V>>,
7}
8
9#[derive(Debug)]
10struct Slot<V: Value> {
11    len: V::K,
12    value: V,
13}
14
15impl<K: Key + 'static, V: Value<K = K> + 'static> RangeTree<V> {
16    pub const fn new() -> Self {
17        Self {
18            inner: BTreeMap::new(),
19        }
20    }
21    fn real_start_bound(&self, start_bound: Bound<K>) -> Bound<K> {
22        self.inner.range((Bound::Unbounded, start_bound))
23            .next_back()
24            .filter(|(k, v)| match start_bound {
25                // already included in the next range
26                Bound::Unbounded => false,
27                Bound::Included(i) => **k + v.len > i,
28                Bound::Excluded(e) => **k + v.len >= e,
29            })
30            .map_or(start_bound, |(k, _)| Bound::Included(*k))
31    }
32    pub fn conflicts(&self, range: impl RangeBounds<K>) -> impl Iterator<Item = (Range<K>, &V)> + '_ {
33        let real_start_bound = self.real_start_bound(range.start_bound().cloned());
34        let is_empty = real_start_bound == range.end_bound().cloned();
35        //dbg!(range.start_bound(), range.end_bound(), is_empty, real_start_bound);
36
37        (!is_empty).then(|| self.inner.range((real_start_bound, range.end_bound().cloned())).map(|(k, v)| (*k..*k + v.len, &v.value))).into_iter().flatten()
38    }
39    pub fn conflicts_mut(&mut self, range: impl RangeBounds<K>) -> impl Iterator<Item = (Range<K>, &mut V)> + '_ {
40        let real_start_bound = self.real_start_bound(range.start_bound().cloned());
41
42        self.inner.range_mut((real_start_bound, range.end_bound().cloned())).map(|(k, v)| (*k..*k + v.len, &mut v.value))
43    }
44    pub fn contains(&self, key: K) -> bool {
45        self.conflicts(key..=key).next().is_some()
46    }
47    pub fn insert(&mut self, mut start: K, mut len: K, mut value: V) -> Result<(), Occupied> {
48        if self.conflicts(start..start + len).next().is_some() {
49            return Err(Occupied);
50        }
51
52        if let Some((prev_base, prev)) = self.inner.range(..start).next_back() {
53            if *prev_base + prev.len == start {
54                match value.try_merge_backwards(&prev.value) {
55                    Ok(v) => {
56                        len += prev.len;
57                        start = *prev_base;
58                        value = v;
59
60                        let _ = self.inner.remove(&start).unwrap();
61                    }
62                    Err(v) => value = v,
63                }
64            }
65        }
66        let end = start + len;
67        if let Some(next) = self.inner.get(&end) {
68            match value.try_merge_forward(&next.value) {
69                Ok(v) => {
70                    len += next.len;
71                    value = v;
72
73                    let _ = self.inner.remove(&end).unwrap();
74                }
75                Err(v) => value = v,
76            }
77        }
78        let _ = self.inner.insert(start, Slot { len, value });
79
80
81        Ok(())
82    }
83    pub fn remove(&mut self, range: Range<K>) -> Vec<(Range<K>, V)> {
84        let mut ret = Vec::new();
85
86        let mut start = range.start_bound().cloned();
87        let end = range.end_bound().cloned();
88        //dbg!(self.conflicts((start, end)).collect::<Vec<_>>());
89
90        while let Some((part_range_ref, _)) = { let x = self.conflicts((start, end)).next(); x } {
91            let k = part_range_ref.start;
92            let Slot { len, value } = self.inner.remove(&k).unwrap();
93            let part_range = k..k+len;
94
95            let (range_before, range_to_remove, range_after) = extract(range.clone(), part_range.clone());
96            let (v_before, value, v_after) = value.split(range_before.clone(), range_to_remove.clone(), range_after.clone());
97
98            if let (Some(range_before), Some(v_before)) = (range_before, v_before) {
99                let k = range_before.start;
100                let len = range_before.end - range_before.start;
101                let _ = self.insert(k, len, v_before);
102            }
103            if let (Some(range_after), Some(v_after)) = (range_after, v_after) {
104                let k = range_after.start;
105                let len = range_after.end - range_after.start;
106                let _ = self.insert(k, len, v_after);
107            }
108
109            ret.push((range_to_remove, value));
110
111            start = part_range.end_bound().cloned();
112        }
113        //dbg!(&ret);
114
115        ret
116    }
117    pub fn remove_and_unused(&mut self, range: Range<K>) -> Vec<(Range<K>, Option<V>)> {
118        let removed = self.remove(range.clone());
119        let mut ret = Vec::new();
120        let mut iter = removed.into_iter().peekable();
121
122        let Some((first_range, first_v)) = iter.next() else {
123            return vec!((range, None));
124        };
125
126        let mut prev_end = first_range.start;
127
128        if first_range.start != range.start {
129            ret.push((K::from(0)..first_range.start, None));
130        }
131        ret.push((first_range.clone(), Some(first_v)));
132
133        while let Some((entry_range, v)) = iter.next() {
134            if prev_end != entry_range.start {
135                ret.push((prev_end..entry_range.start, None));
136            }
137
138            prev_end = entry_range.end;
139            ret.push((entry_range, Some(v)));
140
141            if iter.peek().is_none() {
142                ret.push((prev_end..range.end, None));
143            }
144        }
145
146        ret
147    }
148    pub fn end(&self) -> K {
149        self.inner.last_key_value().map_or(K::from(0), |(k, v)| *k + v.len)
150    }
151    pub fn len(&self) -> usize {
152        self.inner.len()
153    }
154    pub fn is_empty(&self) -> bool {
155        self.len() == 0
156    }
157    pub fn iter(&self) -> impl Iterator<Item = (Range<K>, &V)> + '_ {
158        self.inner.iter().map(|(k, Slot { len, value })| (*k..*k + *len, value))
159    }
160    pub fn iter_mut(&mut self) -> impl Iterator<Item = (Range<K>, &mut V)> + '_ {
161        self.inner.iter_mut().map(|(k, Slot { len, value })| (*k..*k + *len, value))
162    }
163}
164
165fn extract<K: Ord + Copy + std::fmt::Debug>(outer: Range<K>, inner: Range<K>) -> (Option<Range<K>>, Range<K>, Option<Range<K>>) {
166    //dbg!(&outer, &inner);
167    let smallest_start = core::cmp::min(outer.start, inner.start);
168    let greatest_end = core::cmp::max(outer.end, inner.end);
169
170    let middle_start = core::cmp::max(inner.start, outer.start);
171    let middle_end = core::cmp::min(inner.end, outer.end);
172
173    let middle_start = core::cmp::min(middle_start, middle_end);
174    let middle_end = core::cmp::max(middle_start, middle_end);
175
176    (
177        Some(smallest_start..middle_start).filter(|r| !r.is_empty()),
178        middle_start..middle_end,
179        Some(middle_end..greatest_end).filter(|r| !r.is_empty()),
180    )
181}
182
183#[derive(Debug)]
184pub struct Occupied;
185
186pub trait Key: std::ops::Add<Output = Self> + std::ops::Sub<Output = Self> + std::ops::AddAssign + Copy + Ord + From<u8> + std::fmt::Debug {}
187impl Key for u64 {}
188impl Key for usize {}
189
190pub trait Value: Sized + std::fmt::Debug {
191    type K: Key;
192
193    fn try_merge_backwards(self, other: &Self) -> Result<Self, Self>;
194    fn try_merge_forward(self, other: &Self) -> Result<Self, Self>;
195
196    fn split(self, prev_range: Option<Range<Self::K>>, range: Range<Self::K>, next_range: Option<Range<Self::K>>) -> (Option<Self>, Self, Option<Self>);
197}
range_tree/lib.rs

range_tree/
lib.rs
