geo_booleanop/splay/
tree.rs

1use super::node::Node;
2use std::cell::UnsafeCell;
3use std::cmp::Ordering;
4use std::fmt;
5use std::mem;
6use std::ops::{Index, IndexMut};
7
8pub struct SplayTree<K, V, C>
9where
10    C: Fn(&K, &K) -> Ordering,
11{
12    comparator: C,
13    root: UnsafeCell<Option<Box<Node<K, V>>>>,
14    size: usize,
15}
16
17impl<K, V, C> SplayTree<K, V, C>
18where
19    C: Fn(&K, &K) -> Ordering,
20{
21    pub fn new(comparator: C) -> SplayTree<K, V, C> {
22        SplayTree {
23            comparator,
24            root: UnsafeCell::new(None),
25            size: 0,
26        }
27    }
28    pub fn len(&self) -> usize {
29        self.size
30    }
31
32    pub fn is_empty(&self) -> bool {
33        self.len() == 0
34    }
35
36    pub fn clear(&mut self) {
37        self.root_mut().take();
38        self.size = 0;
39    }
40
41    pub fn contains(&self, key: &K) -> bool {
42        self.find_key(key).is_some()
43    }
44
45    pub fn get(&self, key: &K) -> Option<&V> {
46        // Splay trees are self-modifying, which is the cause of this ugly mess
47        match self.root_mut() {
48            Some(ref mut root) => {
49                splay(key, root, &self.comparator);
50                if (self.comparator)(key, &root.key) == Ordering::Equal {
51                    Some(&root.value)
52                } else {
53                    None
54                }
55            }
56            None => None,
57        }
58    }
59
60    /// Return a mutable reference to the value corresponding to the key
61    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
62        // Splay trees are self-modifying, which is the cause of this ugly mess
63        match self.root_mut() {
64            Some(ref mut root) => {
65                splay(key, root, &self.comparator);
66                if (self.comparator)(key, &root.key) == Ordering::Equal {
67                    Some(&mut root.value)
68                } else {
69                    None
70                }
71            }
72            None => None,
73        }
74    }
75
76    pub fn find_key(&self, key: &K) -> Option<&K> {
77        // Splay trees are self-modifying, which is the cause of this ugly mess
78        match self.root_mut() {
79            Some(ref mut root) => {
80                splay(key, root, &self.comparator);
81                if (self.comparator)(key, &root.key) == Ordering::Equal {
82                    Some(&root.key)
83                } else {
84                    None
85                }
86            }
87            None => None,
88        }
89    }
90
91    pub fn next(&self, key: &K) -> Option<(&K, &V)> {
92        // Splay trees are self-modifying, which is the cause of this ugly mess
93        let mut node: &Node<K, V> = match self.root_mut() {
94            Some(ref mut root) => {
95                splay(key, root, &self.comparator);
96                root
97            }
98            None => return None,
99        };
100
101        let mut successor: Option<(&K, &V)> = None;
102
103        loop {
104            match (self.comparator)(key, &node.key) {
105                Ordering::Less => {
106                    successor = Some((&node.key, &node.value));
107                    match node.left {
108                        Some(ref left) => node = left,
109                        None => break,
110                    }
111                }
112                Ordering::Equal | Ordering::Greater => match node.right {
113                    Some(ref right) => node = right,
114                    None => break,
115                },
116            }
117        }
118
119        successor
120    }
121
122    pub fn prev(&self, key: &K) -> Option<(&K, &V)> {
123        // Splay trees are self-modifying, which is the cause of this ugly mess
124        let mut node: &Node<K, V> = match self.root_mut() {
125            Some(ref mut root) => {
126                splay(key, root, &self.comparator);
127                root
128            }
129            None => return None,
130        };
131
132        let mut predecessor: Option<(&K, &V)> = None;
133
134        loop {
135            match (self.comparator)(key, &node.key) {
136                Ordering::Equal | Ordering::Less => match node.left {
137                    Some(ref left) => node = left,
138                    None => break,
139                },
140                Ordering::Greater => {
141                    predecessor = Some((&node.key, &node.value));
142                    match node.right {
143                        Some(ref right) => node = right,
144                        None => break,
145                    }
146                }
147            }
148        }
149
150        predecessor
151    }
152
153    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
154        match self.root_mut() {
155            Some(ref mut root) => {
156                splay(&key, root, &self.comparator);
157
158                match (self.comparator)(&key, &root.key) {
159                    Ordering::Equal => {
160                        let old = mem::replace(&mut root.value, value);
161                        return Some(old);
162                    }
163                    Ordering::Less => {
164                        let left = root.pop_left();
165                        let new = Node::new_boxed(key, value, left, None);
166                        let prev = mem::replace(root, new);
167                        root.right = Some(prev);
168                    }
169                    Ordering::Greater => {
170                        let right = root.pop_right();
171                        let new = Node::new_boxed(key, value, None, right);
172                        let prev = mem::replace(root, new);
173                        root.left = Some(prev);
174                    }
175                }
176            }
177            slot => {
178                *slot = Some(Node::new_boxed(key, value, None, None));
179            }
180        }
181        self.size += 1;
182        None
183    }
184
185    pub fn remove(&mut self, key: &K) -> Option<V> {
186        match *self.root_mut() {
187            None => {
188                return None;
189            }
190            Some(ref mut root) => {
191                splay(key, root, &self.comparator);
192                if (self.comparator)(key, &root.key) != Ordering::Equal {
193                    return None;
194                }
195            }
196        }
197
198        let Node { left, right, value, .. } = *self.root_mut().take().unwrap();
199
200        *self.root_mut() = match left {
201            None => right,
202            Some(mut node) => {
203                splay(key, &mut node, &self.comparator);
204                node.right = right;
205                Some(node)
206            }
207        };
208
209        self.size -= 1;
210        Some(value)
211    }
212
213    pub fn min(&self) -> Option<&K> {
214        self.min_node().map(|node| &node.key)
215    }
216
217    pub fn max(&self) -> Option<&K> {
218        self.max_node().map(|node| &node.key)
219    }
220
221    fn min_node(&self) -> Option<&Node<K, V>> {
222        match self.root_ref() {
223            Some(ref root) => {
224                let mut node = root;
225
226                while let Some(ref left) = node.left {
227                    node = left
228                }
229                Some(node)
230            }
231            None => None,
232        }
233    }
234
235    fn max_node(&self) -> Option<&Node<K, V>> {
236        match self.root_ref() {
237            Some(ref root) => {
238                let mut node = root;
239
240                while let Some(ref right) = node.right {
241                    node = right
242                }
243                Some(node)
244            }
245            None => None,
246        }
247    }
248
249    // Messy code ahead
250    #[allow(clippy::mut_from_ref)]
251    fn root_mut(&self) -> &mut Option<Box<Node<K, V>>> {
252        unsafe { &mut *self.root.get() }
253    }
254
255    fn root_ref(&self) -> &Option<Box<Node<K, V>>> {
256        unsafe { &*self.root.get() }
257    }
258}
259
260impl<'a, K, V, C> Index<&'a K> for SplayTree<K, V, C>
261where
262    C: Fn(&K, &K) -> Ordering,
263{
264    type Output = V;
265
266    fn index(&self, index: &'a K) -> &V {
267        self.get(index).expect("key not present in SplayMap")
268    }
269}
270impl<'a, K, V, C> IndexMut<&'a K> for SplayTree<K, V, C>
271where
272    C: Fn(&K, &K) -> Ordering,
273{
274    fn index_mut(&mut self, index: &K) -> &mut V {
275        self.get_mut(index).expect("key not present in SplayMap")
276    }
277}
278
279impl<K, V, C> Extend<(K, V)> for SplayTree<K, V, C>
280where
281    C: Fn(&K, &K) -> Ordering,
282{
283    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, i: I) {
284        for (k, v) in i {
285            self.insert(k, v);
286        }
287    }
288}
289
290impl<K, V, C> IntoIterator for SplayTree<K, V, C>
291where
292    C: Fn(&K, &K) -> Ordering,
293{
294    type Item = (K, V);
295    type IntoIter = IntoIter<K, V>;
296
297    fn into_iter(self) -> Self::IntoIter {
298        IntoIter {
299            cur: self.root_mut().take(),
300            remaining: self.size,
301        }
302    }
303}
304
305impl<K, V, C> Drop for SplayTree<K, V, C>
306where
307    C: Fn(&K, &K) -> Ordering,
308{
309    fn drop(&mut self) {
310        self.clear();
311    }
312}
313
314impl<K, V, C> fmt::Debug for SplayTree<K, V, C>
315where
316    K: fmt::Debug,
317    V: fmt::Debug,
318    C: Fn(&K, &K) -> Ordering,
319{
320    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
321        write!(f, "{:?}", self.root_ref())
322    }
323}
324
325pub struct IntoIter<K, V> {
326    cur: Option<Box<Node<K, V>>>,
327    remaining: usize,
328}
329
330impl<K, V> Iterator for IntoIter<K, V> {
331    type Item = (K, V);
332    fn next(&mut self) -> Option<(K, V)> {
333        let mut cur = match self.cur.take() {
334            Some(cur) => cur,
335            None => return None,
336        };
337        loop {
338            match cur.pop_left() {
339                Some(node) => {
340                    let mut node = node;
341                    cur.left = node.pop_right();
342                    node.right = Some(cur);
343                    cur = node;
344                }
345
346                None => {
347                    self.cur = cur.pop_right();
348                    // left and right fields are both None
349                    let node = *cur;
350                    let Node { key, value, .. } = node;
351                    self.remaining -= 1;
352                    return Some((key, value));
353                }
354            }
355        }
356    }
357
358    fn size_hint(&self) -> (usize, Option<usize>) {
359        (self.remaining, Some(self.remaining))
360    }
361}
362
363impl<K, V> DoubleEndedIterator for IntoIter<K, V> {
364    fn next_back(&mut self) -> Option<(K, V)> {
365        let mut cur = match self.cur.take() {
366            Some(cur) => cur,
367            None => return None,
368        };
369        loop {
370            match cur.pop_right() {
371                Some(node) => {
372                    let mut node = node;
373                    cur.right = node.pop_left();
374                    node.left = Some(cur);
375                    cur = node;
376                }
377
378                None => {
379                    self.cur = cur.pop_left();
380                    // left and right fields are both None
381                    let node = *cur;
382                    let Node { key, value, .. } = node;
383                    self.remaining -= 1;
384                    return Some((key, value));
385                }
386            }
387        }
388    }
389}
390
391impl<K, V> ExactSizeIterator for IntoIter<K, V> {}
392
393/// Performs a top-down splay operation on a tree rooted at `node`. This will
394/// modify the pointer to contain the new root of the tree once the splay
395/// operation is done. When finished, if `key` is in the tree, it will be at the
396/// root. Otherwise the closest key to the specified key will be at the root.
397#[allow(clippy::borrowed_box)]
398fn splay<K, V, C>(key: &K, node: &mut Box<Node<K, V>>, comparator: &C)
399where
400    C: Fn(&K, &K) -> Ordering,
401{
402    let mut newleft = None;
403    let mut newright = None;
404
405    // Eplicitly grab a new scope so the loans on newleft/newright are
406    // terminated before we move out of them.
407    {
408        // Yes, these are backwards, that's intentional.
409        let mut l = &mut newright;
410        let mut r = &mut newleft;
411
412        loop {
413            match comparator(key, &node.key) {
414                // Found it, yay!
415                Ordering::Equal => break,
416
417                Ordering::Less => {
418                    let mut left = match node.pop_left() {
419                        Some(left) => left,
420                        None => break,
421                    };
422                    // rotate this node right if necessary
423                    if comparator(key, &left.key) == Ordering::Less {
424                        // A bit odd, but avoids drop glue
425                        mem::swap(&mut node.left, &mut left.right);
426                        mem::swap(&mut left, node);
427                        let none = mem::replace(&mut node.right, Some(left));
428                        match mem::replace(&mut node.left, none) {
429                            Some(l) => {
430                                left = l;
431                            }
432                            None => break,
433                        }
434                    }
435
436                    *r = Some(mem::replace(node, left));
437                    let tmp = r;
438                    r = &mut tmp.as_mut().unwrap().left;
439                }
440
441                // If you look closely, you may have seen some similar code
442                // before
443                Ordering::Greater => {
444                    let mut right = match node.pop_right() {
445                        Some(right) => right,
446                        None => break,
447                    };
448
449                    if comparator(key, &right.key) == Ordering::Greater {
450                        mem::swap(&mut node.right, &mut right.left);
451                        mem::swap(&mut right, node);
452                        let none = mem::replace(&mut node.left, Some(right));
453                        match mem::replace(&mut node.right, none) {
454                            Some(r) => {
455                                right = r;
456                            }
457                            None => break,
458                        }
459                    }
460                    *l = Some(mem::replace(node, right));
461                    let tmp = l;
462                    l = &mut tmp.as_mut().unwrap().right;
463                }
464            }
465        }
466
467        mem::swap(l, &mut node.left);
468        mem::swap(r, &mut node.right);
469    }
470
471    node.left = newright;
472    node.right = newleft;
473}