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 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 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
62 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 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 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 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 #[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 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 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#[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 {
408 let mut l = &mut newright;
410 let mut r = &mut newleft;
411
412 loop {
413 match comparator(key, &node.key) {
414 Ordering::Equal => break,
416
417 Ordering::Less => {
418 let mut left = match node.pop_left() {
419 Some(left) => left,
420 None => break,
421 };
422 if comparator(key, &left.key) == Ordering::Less {
424 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 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}