kravltree/
avltree.rs

1use std::borrow::Borrow;
2use std::marker::PhantomData;
3use std::{
4    cmp::Ordering,
5    fmt::{Debug, Formatter},
6    ptr::NonNull,
7};
8
9use crate::entry::*;
10
11pub struct AvlTreeMap<K, V> {
12    tree: Option<EntryPtr<K, V>>,
13    count: usize,
14    first: Option<EntryPtr<K, V>>, // Points to an entry inside the Tree
15    last: Option<EntryPtr<K, V>>,  // Points to an entry inside the Tree
16    modified: bool,                // TODO: We really need this?
17    insert_path: [bool; 48],
18}
19
20impl<K, V> AvlTreeMap<K, V> {
21    pub fn new() -> Self {
22        Self {
23            tree: None,
24            count: 0,
25            first: None,
26            last: None,
27            modified: false,
28            insert_path: [bool::default(); 48],
29        }
30    }
31
32    pub fn len(&self) -> usize {
33        self.count
34    }
35
36    pub fn is_empty(&self) -> bool {
37        self.len() == 0
38    }
39
40
41    pub fn first_key(&self) -> Option<&K> {
42        if let Some(t) = self.first {
43            let k = unsafe { (*t.ptr.as_ptr()).key() };
44            Some(k)
45        } else {
46            None
47        }
48    }
49
50    pub fn last_key(&self) -> Option<&K> {
51        if let Some(t) = self.last {
52            let k = unsafe { (*t.ptr.as_ptr()).key() };
53            Some(k)
54        } else {
55            None
56        }
57    }
58
59    pub fn clear(&mut self) {
60        self.count = 0;
61        // Drop tree
62        let mut start = self.tree;
63        if let Some(start) = start.take() {
64            let _ = unsafe { Box::from_raw(start.as_ptr()) };
65        }
66        // /Drop tree
67        self.tree = None;
68        self.first = None;
69        self.last = None;
70    }
71}
72
73impl<K, V> Clone for AvlTreeMap<K, V> where K: Clone + Ord, V: Clone {
74    fn clone(&self) -> Self {
75        // Clone here just creates a new AvlTreeMap and populates it.
76        // Cloning Entries is not safe since they hold pointers to other
77        // entries in the map, and cloning them just create a new entry
78        // but still points to entries in the original map.
79        // A proper clone implementation would need additional work
80        // to change pointers to point to entries in the new map instead of the original one.
81        let mut new_map = AvlTreeMap::new();
82
83        for (k, v) in self.iter() {
84            new_map.insert(k.clone(), v.clone());
85        }
86
87        new_map
88    }
89}
90
91impl<K, V> AvlTreeMap<K, V>
92where
93    K: PartialOrd + Ord,
94{
95    /// Returns the entry corresponding to the given key, if it is in the tree.
96    pub(crate) fn find_key<Q: ?Sized>(&self, key: &Q) -> Option<&Entry<K, V>>
97    where
98        K: Borrow<Q>,
99        Q: PartialOrd,
100    {
101        let mut tree = self.tree.as_ref().map(|e| unsafe { e.ptr.as_ref() });
102        while let Some(entry) = tree {
103            let cmp = key.partial_cmp(entry.key().borrow());
104            if let Some(cmp) = cmp {
105                if cmp == Ordering::Equal {
106                    return Some(entry);
107                } else if cmp == Ordering::Less {
108                    tree = entry.left();
109                } else {
110                    tree = entry.right();
111                }
112            }
113        }
114        tree
115    }
116
117    pub(crate) fn find_entry<Q: ?Sized>(&mut self, key: &Q) -> Option<&mut Entry<K, V>>
118        where
119            K: Borrow<Q>,
120            Q: PartialOrd,
121    {
122        let mut tree = self.tree;
123        while let Some(mut entry) = tree {
124            let cmp = key.partial_cmp(entry.key().borrow());
125            if let Some(cmp) = cmp {
126                if cmp == Ordering::Equal {
127                    return Some(unsafe { entry.ptr.as_mut() });
128                } else if cmp == Ordering::Less {
129                    tree = entry.raw_left();
130                } else {
131                    tree = entry.raw_right();
132                }
133            }
134        }
135        unsafe { tree.map(|mut p| p.ptr.as_mut()) }
136    }
137
138    /// Locates a key, if the key is present, than the entry for the key is returned,
139    /// otherwise, it will be either the smallest greater key or the greatest smaller key.
140    #[allow(dead_code)]
141    fn locate_key(&self, key: &K) -> Option<&Entry<K, V>> {
142        let mut last = self.tree.as_ref().map(|e| unsafe { e.ptr.as_ref() });
143        let mut tree = self.tree.as_ref().map(|e| unsafe { e.ptr.as_ref() });
144        let mut cmp = None;
145        while let Some(entry) = tree {
146            cmp = key.partial_cmp(entry.key());
147            if let Some(cmp) = cmp {
148                if cmp == Ordering::Equal {
149                    return Some(entry);
150                } else if cmp == Ordering::Less {
151                    last = tree;
152                    tree = entry.left();
153                } else {
154                    last = tree;
155                    tree = entry.right();
156                }
157            }
158        }
159
160        if cmp == Some(Ordering::Equal) {
161            tree
162        } else {
163            last
164        }
165    }
166
167    #[allow(dead_code)]
168    fn reset_insert_path(&mut self) {
169        for i in 0..self.insert_path.len() {
170            self.insert_path[i] = false;
171        }
172    }
173
174    pub fn contains_value<Q: ?Sized>(&self, value: &Q) -> bool
175    where
176        V: Borrow<Q>,
177        Q: PartialEq,
178    {
179        let mut iter = self.iter();
180        while let Some(n) = iter.next() {
181            if value.eq(n.1.borrow()) {
182                return true;
183            }
184        }
185
186        false
187    }
188
189    pub fn contains_key<Q: ?Sized>(&self, key: &Q) -> bool
190    where
191        K: Borrow<Q>,
192        Q: PartialOrd,
193    {
194        self.find_key(key).is_some()
195    }
196
197    pub fn get<Q: ?Sized>(&self, key: &Q) -> Option<&V>
198        where
199            K: Borrow<Q>,
200            Q: PartialOrd,
201    {
202        self.find_key(key).map(|v| v.value())
203    }
204
205    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
206        //let mut modified = false;
207        let e;
208
209        if let None = self.tree {
210            let entry = Entry::new(key, value); // TODO: Avoid cloning
211            let entry = Box::new(entry);
212            let ptr = EntryPtr::from_box(entry);
213            self.first = Some(ptr);
214            self.last = self.first;
215            self.tree = Some(ptr);
216            self.count += 1;
217            return None;
218        } else {
219            let mut p = self.tree;
220            let mut y = p;
221            let mut q = None;
222            let mut z = None;
223            let w;
224            let mut cmp;
225            let mut i = 0usize;
226
227            // Carefuly maintain this loop to never end up with an infinite loop
228            loop {
229                // TODO: Provisory, make p and y meet the borrow checker rules
230                //cmp = *(&key.cmp(p.key()));
231                cmp = *(&key.cmp(p.unwrap().key()));
232                if cmp == Ordering::Equal {
233                    return Some(p.unwrap().set_value(value));
234                }
235
236                // p.balance()
237                if p.unwrap().balance() != 0 {
238                    i = 0;
239                    z = q;
240                    y = p;
241                }
242
243                let greater = cmp == Ordering::Greater;
244                self.insert_path[i] = greater;
245
246                i += 1;
247
248                if greater {
249                    // p.succ()
250                    if p.unwrap().succ() {
251                        self.count += 1;
252                        let entry = Entry::new(key, value); // TODO: avoid cloning
253                        let entry = Box::new(entry);
254                        let raw_entry = Box::into_raw(entry);
255                        let mut entry_ptr = EntryPtr::new(NonNull::new(raw_entry).unwrap());
256                        e = Some(entry_ptr);
257                        //modified = true;
258
259                        if let None = p.unwrap().right_field() {
260                            self.last = Some(entry_ptr);
261                        }
262
263                        let mut raw_p = p.unwrap();
264
265                        entry_ptr.set_left_field(Some(raw_p));
266
267                        entry_ptr.set_right_field(raw_p.right_field());
268                        raw_p.set_right_entry(Some(entry_ptr));
269                        break;
270                    }
271
272                    q = p;
273                    p = p.unwrap().right_field(); // What we do if unwrap() fails?
274                                                  //p = NonNull::new(unsafe { p.as_ref() }.raw_right()).unwrap();
275                } else {
276                    if p.unwrap().pred() {
277                        self.count += 1;
278                        let entry = Entry::new(key, value); // TODO: avoid cloning
279                        let entry = Box::new(entry);
280                        let raw_entry = Box::into_raw(entry);
281                        let mut entry_ptr = EntryPtr::new(NonNull::new(raw_entry).unwrap());
282                        e = Some(entry_ptr);
283                        //modified = true;
284
285                        if let None = p.unwrap().left_field() {
286                            self.first = Some(entry_ptr);
287                        }
288
289                        let mut raw_p = p.unwrap();
290
291                        entry_ptr.set_right_field(p);
292                        entry_ptr.set_left_field(raw_p.left_field());
293                        raw_p.set_left_entry(Some(entry_ptr));
294                        break;
295                    }
296
297                    q = p;
298                    p = p.unwrap().left_field(); // What we do if unwrap() fails?
299                }
300            }
301            p = y;
302            i = 0;
303            while let Some(c_entry) = e {
304                if p.is_none() || p.unwrap().as_ptr() == c_entry.as_ptr() {
305                    break;
306                }
307
308                if self.insert_path[i] {
309                    p.unwrap().increment_balance();
310                } else {
311                    p.unwrap().decrement_balance();
312                }
313
314                if self.insert_path[i] {
315                    p = p.unwrap().right_field();
316                    if p.is_none() {
317                        break;
318                    }
319                } else {
320                    p = p.unwrap().left_field();
321                    if p.is_none() {
322                        break;
323                    }
324                }
325
326                i += 1;
327            }
328
329            // Balance
330            if y.unwrap().balance() == -2 {
331                // x = y.left
332                let mut x = y.unwrap().left_field().unwrap();
333
334                let mut y = y.unwrap();
335
336                if x.balance() == -1 {
337                    // Ok(1)
338                    w = Some(x);
339
340                    if x.succ() {
341                        x.set_succ(false);
342                        y.set_pred_entry(Some(x));
343                    } else {
344                        y.set_left_field(x.right_field());
345                    }
346
347                    x.set_right_field(Some(y)); // Rotate
348                    x.set_balance(0);
349                    y.set_balance(0);
350                } else {
351                    debug_assert!(x.balance() == 1);
352                    w = x.raw_right(); // w = x.right
353                    let mut w = w.unwrap();
354
355                    {
356                        // Rotate
357                        x.set_right_field(w.left_field()); //x.right = w.left | x.right = x.right.left | sets the right side of X as left side of X.R
358                        w.set_left_field(Some(x)); // w.left = x | sets the left side of the right side of X as the X
359                        y.set_left_field(w.right_field()); // y.left = w.right
360                        w.set_right_field(Some(y));
361                        // w.right = y
362                    }
363
364                    if w.balance() == -1 {
365                        x.set_balance(0);
366                        y.set_balance(1);
367                    } else if w.balance() == 0 {
368                        x.set_balance(0);
369                        y.set_balance(0);
370                    } else {
371                        x.set_balance(-1);
372                        y.set_balance(0);
373                    }
374
375                    w.set_balance(0);
376
377                    if w.pred() {
378                        // Rotate
379                        x.set_succ_entry(Some(w)); // x.succ(w)
380                        w.set_pred(false);
381                    }
382
383                    if w.succ() {
384                        // Rotate
385                        y.set_pred_entry(Some(w)); // y.pred(w)
386                        w.set_succ(false);
387                    }
388                }
389            } else if y.unwrap().balance() == 2 {
390                let mut y = y.unwrap();
391                let mut x = y.right_field().unwrap(); // x = y.right
392
393                if x.balance() == 1 {
394                    // Rotate
395                    w = Some(x);
396
397                    if x.pred() {
398                        x.set_pred(false);
399                        y.set_succ_entry(Some(x)); // y.succ(x)
400                    } else {
401                        y.set_right_field(x.left_field());
402                        // y.right = x.left
403                    }
404
405                    x.set_left_field(Some(y)); // x.left = y
406                    x.set_balance(0);
407                    y.set_balance(0);
408                } else {
409                    debug_assert!(x.balance() == -1);
410                    w = x.left_field(); // w = x.left
411                    let mut w = w.unwrap();
412                    {
413                        // Rotate
414                        x.set_left_field(w.right_field()); // x.left = w.right
415                        w.set_right_field(Some(x)); // w.right = x
416                        y.set_right_field(w.left_field()); // y.right = w.left
417                        w.set_left_field(Some(y));
418                        // w.left = y
419                    }
420
421                    if w.balance() == 1 {
422                        x.set_balance(0);
423                        y.set_balance(-1);
424                    } else if w.balance() == 0 {
425                        x.set_balance(0);
426                        y.set_balance(0);
427                    } else {
428                        x.set_balance(1);
429                        y.set_balance(0);
430                    }
431
432                    w.set_balance(0);
433
434                    if w.pred() {
435                        y.set_succ_entry(Some(w)); // y.pred(w)
436                        w.set_pred(false);
437                    }
438
439                    if w.succ() {
440                        x.set_pred_entry(Some(w)); // x.pred(w)
441                        w.set_succ(false);
442                    }
443                }
444            } else {
445                return None;
446            };
447
448            if z.is_none() {
449                self.tree = w;
450            } else {
451                let mut z = z.unwrap();
452                if z.raw_left().map(|l| l.as_ptr()) == y.map(|y| y.as_ptr()) {
453                    z.set_left_field(w);
454                } else {
455                    z.set_right_field(w);
456                }
457            }
458        }
459
460        None
461    }
462
463    pub fn retain<F>(&mut self, mut f: F)
464        where
465            F: FnMut(&K, &mut V) -> bool,
466    {
467        for mut entry in self.entry_iter() {
468            let (key, value) = entry.as_mut();
469            if !f(key, value) {
470                self.remove(key);
471            }
472        }
473    }
474
475    pub fn remove<Q: ?Sized>(&mut self, key: &Q) -> Option<V> where K: Borrow<Q>, Q: PartialOrd + Ord {
476        self.remove_key(key).map(|v| v.take())
477    }
478
479    pub(crate) fn remove_key<Q: ?Sized>(&mut self, key: &Q) -> Option<Entry<K, V>> where K: Borrow<Q>, Q: PartialOrd + Ord {
480        self.modified = false;
481        let mut found_entry = self.tree?;
482        let mut cmp;
483        let mut parent_entry = None;
484        let mut greater = false;
485
486        let kk = key; // Inline
487
488        loop {
489            cmp = kk.cmp(found_entry.key().borrow());
490            if cmp == Ordering::Equal {
491                break;
492            }
493            greater = cmp == Ordering::Greater;
494
495            if greater {
496                // if greater
497                parent_entry = Some(found_entry);
498                found_entry = found_entry.right_field()?;
499            } else {
500                parent_entry = Some(found_entry);
501                found_entry = found_entry.left_field()?;
502            }
503        }
504        // Loop end
505
506        self.remove_entry(parent_entry, found_entry, greater)
507    }
508
509    #[allow(dead_code)]
510    pub(crate) fn remove_entry_<Q: ?Sized>(&mut self, prev: Option<EntryPtr<K, V>>, entry: EntryPtr<K, V>) -> Option<Entry<K, V>> where K: Borrow<Q>, Q: PartialOrd + Ord {
511        let mut greater = false;
512        if let Some(q) = prev {
513            greater = entry.key().cmp(q.key()) == Ordering::Greater;
514        }
515
516        self.remove_entry(prev, entry, greater)
517    }
518
519    pub(crate) fn remove_entry<Q: ?Sized>(&mut self, prev: Option<EntryPtr<K, V>>, entry: EntryPtr<K, V>, mut greater: bool) -> Option<Entry<K, V>> where K: Borrow<Q>, Q: PartialOrd + Ord {
520        self.modified = false;
521        let mut current_entry = entry;
522        let mut previous_entry = prev;
523
524        // Sets first and last nodes
525        if current_entry.left_field_ref().is_none() {
526            self.first = current_entry.next_as_ptr();
527        }
528
529        if current_entry.right_field_ref().is_none() {
530            self.last = current_entry.prev_as_ptr();
531        }
532        // /Sets first and last nodes
533
534        if current_entry.succ() {
535            if current_entry.pred() {
536                if let Some(mut q) = previous_entry {
537                    if greater {
538                        q.set_succ_entry(current_entry.right_field());
539                    } else {
540                        q.set_pred_entry(current_entry.left_field());
541                    }
542                } else {
543                    self.tree = if greater { current_entry.right_field() } else { current_entry.left_field() };
544                }
545            } else {
546                current_entry.prev_as_ptr().unwrap().set_right_entry(current_entry.right_field());
547
548                if let Some(mut q) = previous_entry {
549                    if greater {
550                        q.set_right_entry(current_entry.left_field())
551                    } else {
552                        q.set_left_entry(current_entry.left_field())
553                    }
554                } else {
555                    self.tree = current_entry.left_field();
556                }
557            }
558        } else {
559            let mut r = current_entry.right_field()?;
560
561            if r.pred() {
562                r.set_left_field(current_entry.left_field());
563                r.set_pred(current_entry.pred());
564
565                if !r.pred() {
566                    r.prev_as_ptr()?.set_right_field(Some(r));
567                }
568
569                if let Some(mut q) = previous_entry {
570                    if greater {
571                        q.set_right_field(Some(r));
572                    } else {
573                        q.set_left_field(Some(r));
574                    }
575                } else {
576                    self.tree = Some(r);
577                }
578
579                r.set_balance(current_entry.balance());
580                previous_entry = Some(r);
581                greater = true;
582            } else {
583                let mut s;
584                loop {
585                    s = r.left_field()?;
586                    if s.pred() {
587                        break;
588                    }
589                    r = s;
590                }
591
592                if s.succ() {
593                    r.set_pred_entry(Some(s));
594                } else {
595                    r.set_left_entry(s.right_field());
596                }
597
598                s.set_left_field(current_entry.left_field());
599
600                if !current_entry.pred() {
601                    current_entry.prev_as_ptr()?.set_right_field(Some(s));
602                    s.set_pred(false);
603                }
604
605                s.set_right_field(current_entry.right_field());
606                s.set_succ(true);
607
608                if let Some(mut q) = previous_entry {
609                    if greater {
610                        q.set_right_field(Some(s));
611                    } else {
612                        q.set_left_field(Some(s));
613                    }
614                } else {
615                    self.tree = Some(s);
616                }
617
618                s.set_balance(current_entry.balance());
619                previous_entry = Some(r);
620                greater = false;
621            }
622        }
623
624        let mut y;
625        while let Some(q2) = previous_entry {
626            y = q2;
627            previous_entry = self.parent(y);
628
629            if !greater {
630                greater = if let Some(q) = previous_entry {
631                    q.left_field_as_ptr() != Some(y.as_ptr())
632                } else {
633                    false
634                };
635                y.increment_balance();
636                if y.balance() == 1 {
637                    break;
638                } else if y.balance() == 2 {
639                    let mut x = y.right_field().unwrap();
640                    if x.balance() == -1 {
641                        let mut w = x.left_field().unwrap();
642                        x.set_left_field(w.right_field());
643                        w.set_right_field(Some(x));
644                        y.set_right_field(w.left_field());
645                        w.set_left_field(Some(y));
646
647                        if w.balance() == 1 {
648                            x.set_balance(0);
649                            y.set_balance(-1)
650                        } else if w.balance() == 0 {
651                            x.set_balance(0);
652                            y.set_balance(0)
653                        } else {
654                            assert_eq!(w.balance(), -1);
655                            x.set_balance(1);
656                            y.set_balance(0)
657                        }
658
659                        w.set_balance(0);
660                        if w.pred() {
661                            y.set_succ_entry(Some(w));
662                            w.set_pred(false);
663                        }
664                        if w.succ() {
665                            x.set_pred_entry(Some(w));
666                            w.set_succ(false);
667                        }
668
669                        if let Some(mut q) = previous_entry {
670                            if greater {
671                                q.set_right_field(Some(w));
672                            } else {
673                                q.set_left_field(Some(w));
674                            }
675                        } else {
676                            self.tree = Some(w);
677                        }
678                    } else {
679                        if let Some(mut q) = previous_entry {
680                            if greater {
681                                q.set_right_entry(Some(x));
682                            } else {
683                                q.set_left_entry(Some(x));
684                            }
685                        } else {
686                            self.tree = Some(x);
687                        }
688
689                        if x.balance() == 0 {
690                            y.set_right_field(x.left_field());
691                            x.set_left_field(Some(y));
692                            x.set_balance(-1);
693                            y.set_balance(1);
694                            break;
695                        }
696                        assert_eq!(x.balance(), 1);
697                        if x.pred() {
698                            y.set_succ(true);
699                            x.set_pred(false);
700                        } else {
701                            y.set_right_entry(x.left_field());
702                        }
703                        x.set_left_field(Some(y));
704                        y.set_balance(0);
705                        x.set_balance(0);
706                    }
707                }
708            } else {
709                greater = if let Some(q) = previous_entry {
710                    q.left_field_as_ptr() != Some(y.as_ptr())
711                } else {
712                    false
713                };
714                y.decrement_balance();
715                if y.balance() == -1 {
716                    break;
717                } else if y.balance() == -2 {
718                    let mut x = y.left_field().unwrap();
719                    if x.balance() == 1 {
720                        let mut w = x.right_field().unwrap();
721                        x.set_right_field(w.left_field());
722                        w.set_left_field(Some(x));
723                        y.set_left_field(w.right_field());
724                        w.set_right_field(Some(y));
725
726                        if w.balance() == -1 {
727                            x.set_balance(0);
728                            y.set_balance(1);
729                        } else if w.balance() == 0 {
730                            x.set_balance(0);
731                            y.set_balance(0);
732                        } else {
733                            assert_eq!(w.balance(), 1);
734                            x.set_balance(-1);
735                            y.set_balance(0);
736                        }
737                        w.set_balance(0);
738
739                        if w.pred() {
740                            x.set_succ_entry(Some(w));
741                            w.set_pred(false);
742                        }
743                        if w.succ() {
744                            y.set_pred_entry(Some(w));
745                            w.set_succ(false);
746                        }
747                        if let Some(mut q) = previous_entry {
748                            if greater {
749                                q.set_right_field(Some(w));
750                            } else {
751                                q.set_left_field(Some(w));
752                            }
753                        } else {
754                            self.tree = Some(w);
755                        }
756                    } else {
757                        if let Some(mut q) = previous_entry {
758                            if greater {
759                                q.set_right_field(Some(x));
760                            } else {
761                                q.set_left_field(Some(x));
762                            }
763                        } else {
764                            self.tree = Some(x);
765                        }
766
767                        if x.balance() == 0 {
768                            y.set_left_field(x.right_field());
769                            x.set_right_field(Some(y));
770                            x.set_balance(1);
771                            y.set_balance(-1);
772                            break;
773                        }
774
775                        assert_eq!(x.balance(), -1);
776
777                        if x.succ() {
778                            y.set_pred(true);
779                            x.set_succ(false);
780                        } else {
781                            y.set_left_field(x.right_field());
782                        }
783
784                        x.set_right_field(Some(y));
785                        y.set_balance(0);
786                        x.set_balance(0);
787                    }
788                }
789            }
790        }
791
792        self.modified = true;
793        self.count -= 1;
794
795        if let Some(_) = current_entry.left_field() {
796            current_entry.set_left_field(None);
797        }
798
799        if let Some(_) = current_entry.right_field() {
800            current_entry.set_right_field(None);
801        }
802
803        Some(current_entry.take_entry())
804    }
805
806    fn parent(&self, entry: EntryPtr<K, V>) -> Option<EntryPtr<K, V>> {
807        let _ = self.tree?;
808        let mut x;
809        let mut y;
810        let mut p;
811        y = Some(entry);
812        x = y;
813        loop {
814            if y?.succ() {
815                p = y?.right_field();
816                if p.is_none() || p.unwrap().left_field_as_ptr() != Some(entry.as_ptr()) {
817                    while !x?.pred() {
818                        x = x?.left_field();
819                    }
820
821                    p = x?.left_field();
822                }
823                return p;
824            } else if x?.pred() {
825                p = x?.left_field();
826                if p.is_none() || p.unwrap().right_field_as_ptr() != Some(entry.as_ptr()) {
827                    while !y?.succ() {
828                        y = y?.right_field();
829                    }
830                    p = y?.right_field();
831                }
832                return p;
833            }
834            x = x.map(|v| v.left_field()).flatten();
835            y = y.map(|v| v.right_field()).flatten();
836        }
837    }
838
839    pub fn iter(&self) -> AvlTreeIter<'_, K, V> {
840        AvlTreeIter::new(self)
841    }
842
843    pub(crate) fn entry_iter(&self) -> AvlTreeEntryIter<K, V> {
844        AvlTreeEntryIter::new(self)
845    }
846}
847
848impl<K, V> Drop for AvlTreeMap<K, V> {
849    fn drop(&mut self) {
850        let mut start = self.tree;
851        if let Some(start) = start.take() {
852            let _ = unsafe { Box::from_raw(start.as_ptr()) };
853        }
854    }
855}
856
857pub struct AvlTreeIter<'a, K, V> {
858    next: Option<EntryPtr<K, V>>,
859    prev: Option<EntryPtr<K, V>>,
860    _marker: PhantomData<(&'a K, &'a V)>,
861}
862
863impl<'a, K, V> AvlTreeIter<'a, K, V> {
864    fn new(tree: &'a AvlTreeMap<K, V>) -> Self {
865        Self {
866            next: tree.first,
867            prev: tree.last,
868            _marker: PhantomData,
869        }
870    }
871}
872
873impl<'a, K, V> Iterator for AvlTreeIter<'a, K, V>
874where
875    K: PartialOrd,
876{
877    type Item = (&'a K, &'a V);
878
879    fn next(&mut self) -> Option<Self::Item> {
880        if let Some(e) = self.next {
881            if let Some(p) = self.prev {
882                if e.key().partial_cmp(p.key()) == Some(Ordering::Greater) {
883                    return None;
884                }
885            }
886
887            let key = unsafe { (*e.ptr.as_ptr()).key() };
888            let value = unsafe { (*e.ptr.as_ptr()).value() };
889            let kv = (key, value);
890            self.next = e.next_as_ptr();
891            Some(kv)
892        } else {
893            None
894        }
895    }
896}
897
898impl<'a, K, V> DoubleEndedIterator for AvlTreeIter<'a, K, V>
899where
900    K: PartialOrd,
901{
902    fn next_back(&mut self) -> Option<Self::Item> {
903        if let Some(e) = self.prev {
904            if let Some(p) = self.next {
905                if e.key().partial_cmp(p.key()) == Some(Ordering::Less) {
906                    return None;
907                }
908            }
909
910            let key = unsafe { (*e.ptr.as_ptr()).key() };
911            let value = unsafe { (*e.ptr.as_ptr()).value() };
912            let kv = (key, value);
913            self.prev = e.prev_as_ptr();
914            Some(kv)
915        } else {
916            None
917        }
918    }
919}
920
921pub(crate) struct AvlTreeEntryIter<K, V> {
922    next: Option<EntryPtr<K, V>>,
923    prev: Option<EntryPtr<K, V>>,
924}
925
926impl<K, V> AvlTreeEntryIter<K, V> {
927    fn new(tree: &AvlTreeMap<K, V>) -> Self {
928        Self {
929            next: tree.first,
930            prev: tree.last
931        }
932    }
933}
934
935impl<K, V> Iterator for AvlTreeEntryIter<K, V>
936    where
937        K: PartialOrd,
938{
939    type Item = EntryPtr<K, V>;
940
941    fn next(&mut self) -> Option<Self::Item> {
942        if let Some(e) = self.next {
943            if let Some(p) = self.prev {
944                if e.key().partial_cmp(p.key()) == Some(Ordering::Greater) {
945                    return None;
946                }
947            }
948
949            self.next = e.next_as_ptr();
950            Some(e)
951        } else {
952            None
953        }
954    }
955}
956
957impl<K, V> DoubleEndedIterator for AvlTreeEntryIter<K, V>
958    where
959        K: PartialOrd,
960{
961    fn next_back(&mut self) -> Option<Self::Item> {
962        if let Some(e) = self.prev {
963            if let Some(p) = self.next {
964                if e.key().partial_cmp(p.key()) == Some(Ordering::Less) {
965                    return None;
966                }
967            }
968
969            self.prev = e.prev_as_ptr();
970            Some(e)
971        } else {
972            None
973        }
974    }
975}
976
977
978#[cfg(not(feature = "print-tree"))]
979impl<K, V> Debug for AvlTreeMap<K, V>
980where
981    K: Debug + Ord,
982    V: Debug,
983{
984    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
985        let mut d = f.debug_map();
986
987        for (k, v) in self.iter() {
988            d.entry(k, v);
989        }
990
991        d.finish()
992    }
993}
994
995#[cfg(feature = "print-tree")]
996impl<K, V> Debug for AvlTreeMap<K, V>
997where
998    K: Debug,
999    V: Debug,
1000{
1001    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
1002        write!(f, "AVLTree {{\n")?;
1003
1004        let current_node = self.tree;
1005        if let Some(current_node) = current_node {
1006            let tree = TreeDebug {
1007                root: current_node.ptr,
1008            };
1009            write!(f, "{:?}", tree)?;
1010        }
1011        write!(f, "}}")
1012    }
1013}
1014
1015#[cfg(feature = "print-tree")]
1016struct TreeDebug<K, V> {
1017    root: NonNull<Entry<K, V>>,
1018}
1019
1020#[cfg(feature = "print-tree")]
1021impl<K, V> Debug for TreeDebug<K, V>
1022where
1023    K: Debug,
1024    V: Debug,
1025{
1026    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1027        let root_node = EntryPtr::new(self.root);
1028
1029        fn write_child<K, V>(
1030            f: &mut Formatter<'_>,
1031            node: EntryPtr<K, V>,
1032            pref: &str,
1033            child_pref: &str,
1034            postf: &str,
1035        ) -> std::fmt::Result
1036        where
1037            K: Debug,
1038            V: Debug,
1039        {
1040            write!(f, "{}{:?}{}\n", pref, unsafe { node.ptr.as_ref() }, postf)?;
1041            let left = node.raw_left();
1042            let right = node.raw_right();
1043            let has_right = right.is_some();
1044
1045            if let Some(left) = left {
1046                if has_right {
1047                    write_child(
1048                        f,
1049                        left,
1050                        &format!("{}├── L(", child_pref),
1051                        &format!("{}│   ", child_pref),
1052                        ")",
1053                    )?;
1054                } else {
1055                    write_child(
1056                        f,
1057                        left,
1058                        &format!("{}└── L(", child_pref),
1059                        &format!("{}    ", child_pref),
1060                        ")",
1061                    )?;
1062                }
1063            }
1064
1065            if let Some(right) = right {
1066                write_child(
1067                    f,
1068                    right,
1069                    &format!("{}└── R(", child_pref),
1070                    &format!("{}    ", child_pref),
1071                    ")",
1072                )?;
1073            }
1074
1075            Ok(())
1076        }
1077
1078        write_child(f, root_node, "    Root(", "    ", ")")
1079    }
1080}
kravltree/avltree.rs

kravltree/
avltree.rs
