sp_im/nodes/
btree.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4
5use alloc::{
6  borrow::Borrow,
7  vec::Vec,
8};
9
10use core::{
11  cmp::Ordering,
12  mem,
13  ops::{
14    Bound,
15    RangeBounds,
16  },
17};
18use sp_sized_chunks::Chunk;
19
20use crate::{
21  config::ORD_CHUNK_SIZE,
22  util::{
23    Pool,
24    PoolClone,
25    PoolDefault,
26    PoolRef,
27  },
28};
29
30use self::{
31  Insert::*,
32  InsertAction::*,
33};
34
35const NODE_SIZE: usize = ORD_CHUNK_SIZE;
36const MEDIAN: usize = (NODE_SIZE + 1) >> 1;
37
38pub trait BTreeValue {
39  type Key;
40  fn ptr_eq(&self, other: &Self) -> bool;
41  fn search_key<BK>(slice: &[Self], key: &BK) -> Result<usize, usize>
42  where
43    BK: Ord + ?Sized,
44    Self: Sized,
45    Self::Key: Borrow<BK>;
46  fn search_value(slice: &[Self], value: &Self) -> Result<usize, usize>
47  where Self: Sized;
48  fn cmp_keys<BK>(&self, other: &BK) -> Ordering
49  where
50    BK: Ord + ?Sized,
51    Self::Key: Borrow<BK>;
52  fn cmp_values(&self, other: &Self) -> Ordering;
53}
54
55pub(crate) struct Node<A> {
56  keys: Chunk<A, 64>,
57  children: Chunk<Option<PoolRef<Node<A>>>, 65>,
58}
59
60#[cfg(feature = "pool")]
61#[allow(unsafe_code)]
62unsafe fn cast_uninit<A>(target: &mut A) -> &mut mem::MaybeUninit<A> {
63  &mut *(target as *mut A as *mut mem::MaybeUninit<A>)
64}
65
66#[allow(unsafe_code)]
67impl<A> PoolDefault for Node<A> {
68  #[cfg(feature = "pool")]
69  unsafe fn default_uninit(target: &mut mem::MaybeUninit<Self>) {
70    let ptr: *mut Self = target.as_mut_ptr();
71    Chunk::default_uninit(cast_uninit(&mut (*ptr).keys));
72    Chunk::default_uninit(cast_uninit(&mut (*ptr).children));
73    (*ptr).children.push_back(None);
74  }
75}
76
77#[allow(unsafe_code)]
78impl<A> PoolClone for Node<A>
79where A: Clone
80{
81  #[cfg(feature = "pool")]
82  unsafe fn clone_uninit(&self, target: &mut mem::MaybeUninit<Self>) {
83    self.keys.clone_uninit(cast_uninit(&mut (*target.as_mut_ptr()).keys));
84    self
85      .children
86      .clone_uninit(cast_uninit(&mut (*target.as_mut_ptr()).children));
87  }
88}
89
90pub(crate) enum Insert<A> {
91  Added,
92  Replaced(A),
93  Split(Node<A>, A, Node<A>),
94}
95
96enum InsertAction<A> {
97  AddedAction,
98  ReplacedAction(A),
99  InsertAt,
100  InsertSplit(Node<A>, A, Node<A>),
101}
102
103pub(crate) enum Remove<A> {
104  NoChange,
105  Removed(A),
106  Update(A, Node<A>),
107}
108
109enum RemoveAction {
110  DeleteAt(usize),
111  PullUp(usize, usize, usize),
112  Merge(usize),
113  StealFromLeft(usize),
114  StealFromRight(usize),
115  MergeFirst(usize),
116  ContinueDown(usize),
117}
118
119impl<A> Clone for Node<A>
120where A: Clone
121{
122  fn clone(&self) -> Self {
123    Node { keys: self.keys.clone(), children: self.children.clone() }
124  }
125}
126
127impl<A> Default for Node<A> {
128  fn default() -> Self {
129    Node { keys: Chunk::new(), children: Chunk::unit(None) }
130  }
131}
132
133impl<A> Node<A> {
134  #[inline]
135  fn has_room(&self) -> bool { self.keys.len() < NODE_SIZE }
136
137  #[inline]
138  fn too_small(&self) -> bool { self.keys.len() < MEDIAN }
139
140  #[inline]
141  pub(crate) fn unit(value: A) -> Self {
142    Node { keys: Chunk::unit(value), children: Chunk::pair(None, None) }
143  }
144
145  #[inline]
146  pub(crate) fn new_from_split(
147    pool: &Pool<Node<A>>,
148    left: Node<A>,
149    median: A,
150    right: Node<A>,
151  ) -> Self {
152    Node {
153      keys: Chunk::unit(median),
154      children: Chunk::pair(
155        Some(PoolRef::new(pool, left)),
156        Some(PoolRef::new(pool, right)),
157      ),
158    }
159  }
160
161  pub(crate) fn min(&self) -> Option<&A> {
162    match self.children.first().unwrap() {
163      None => self.keys.first(),
164      Some(ref child) => child.min(),
165    }
166  }
167
168  pub(crate) fn max(&self) -> Option<&A> {
169    match self.children.last().unwrap() {
170      None => self.keys.last(),
171      Some(ref child) => child.max(),
172    }
173  }
174}
175
176impl<A: BTreeValue> Node<A> {
177  fn child_contains<BK>(&self, index: usize, key: &BK) -> bool
178  where
179    BK: Ord + ?Sized,
180    A::Key: Borrow<BK>, {
181    if let Some(Some(ref child)) = self.children.get(index) {
182      child.lookup(key).is_some()
183    }
184    else {
185      false
186    }
187  }
188
189  pub(crate) fn lookup<BK>(&self, key: &BK) -> Option<&A>
190  where
191    BK: Ord + ?Sized,
192    A::Key: Borrow<BK>, {
193    if self.keys.is_empty() {
194      return None;
195    }
196    // Perform a binary search, resulting in either a match or
197    // the index of the first higher key, meaning we search the
198    // child to the left of it.
199    match A::search_key(&self.keys, key) {
200      Ok(index) => Some(&self.keys[index]),
201      Err(index) => match self.children[index] {
202        None => None,
203        Some(ref node) => node.lookup(key),
204      },
205    }
206  }
207
208  pub(crate) fn lookup_mut<BK>(
209    &mut self,
210    pool: &Pool<Node<A>>,
211    key: &BK,
212  ) -> Option<&mut A>
213  where
214    A: Clone,
215    BK: Ord + ?Sized,
216    A::Key: Borrow<BK>,
217  {
218    if self.keys.is_empty() {
219      return None;
220    }
221    // Perform a binary search, resulting in either a match or
222    // the index of the first higher key, meaning we search the
223    // child to the left of it.
224    match A::search_key(&self.keys, key) {
225      Ok(index) => Some(&mut self.keys[index]),
226      Err(index) => match self.children[index] {
227        None => None,
228        Some(ref mut child_ref) => {
229          let child = PoolRef::make_mut(pool, child_ref);
230          child.lookup_mut(pool, key)
231        }
232      },
233    }
234  }
235
236  pub(crate) fn lookup_prev<'a, BK>(&'a self, key: &BK) -> Option<&A>
237  where
238    BK: Ord + ?Sized,
239    A::Key: Borrow<BK>, {
240    if self.keys.is_empty() {
241      return None;
242    }
243    match A::search_key(&self.keys, key) {
244      Ok(index) => Some(&self.keys[index]),
245      Err(index) => match self.children[index] {
246        None if index == 0 => None,
247        None => match self.keys.get(index - 1) {
248          Some(_) => Some(&self.keys[index - 1]),
249          None => None,
250        },
251        Some(ref node) => node.lookup_prev(key),
252      },
253    }
254  }
255
256  pub(crate) fn lookup_next<'a, BK>(&'a self, key: &BK) -> Option<&A>
257  where
258    BK: Ord + ?Sized,
259    A::Key: Borrow<BK>, {
260    if self.keys.is_empty() {
261      return None;
262    }
263    match A::search_key(&self.keys, key) {
264      Ok(index) => Some(&self.keys[index]),
265      Err(index) => match self.children[index] {
266        None => match self.keys.get(index) {
267          Some(_) => Some(&self.keys[index]),
268          None => None,
269        },
270        Some(ref node) => node.lookup_next(key),
271      },
272    }
273  }
274
275  pub(crate) fn lookup_prev_mut<'a, BK>(
276    &'a mut self,
277    pool: &Pool<Node<A>>,
278    key: &BK,
279  ) -> Option<&mut A>
280  where
281    A: Clone,
282    BK: Ord + ?Sized,
283    A::Key: Borrow<BK>,
284  {
285    if self.keys.is_empty() {
286      return None;
287    }
288    match A::search_key(&self.keys, key) {
289      Ok(index) => Some(&mut self.keys[index]),
290      Err(index) => match self.children[index] {
291        None if index == 0 => None,
292        None => match self.keys.get(index - 1) {
293          Some(_) => Some(&mut self.keys[index - 1]),
294          None => None,
295        },
296        Some(ref mut node) => {
297          PoolRef::make_mut(pool, node).lookup_prev_mut(pool, key)
298        }
299      },
300    }
301  }
302
303  pub(crate) fn lookup_next_mut<'a, BK>(
304    &'a mut self,
305    pool: &Pool<Node<A>>,
306    key: &BK,
307  ) -> Option<&mut A>
308  where
309    A: Clone,
310    BK: Ord + ?Sized,
311    A::Key: Borrow<BK>,
312  {
313    if self.keys.is_empty() {
314      return None;
315    }
316    match A::search_key(&self.keys, key) {
317      Ok(index) => Some(&mut self.keys[index]),
318      Err(index) => match self.children[index] {
319        None => match self.keys.get(index) {
320          Some(_) => Some(&mut self.keys[index]),
321          None => None,
322        },
323        Some(ref mut node) => {
324          PoolRef::make_mut(pool, node).lookup_next_mut(pool, key)
325        }
326      },
327    }
328  }
329
330  pub(crate) fn path_first<'a, BK>(
331    &'a self,
332    mut path: Vec<(&'a Node<A>, usize)>,
333  ) -> Vec<(&'a Node<A>, usize)>
334  where
335    A: 'a,
336    BK: Ord + ?Sized,
337    A::Key: Borrow<BK>,
338  {
339    if self.keys.is_empty() {
340      return Vec::new();
341    }
342    match self.children[0] {
343      None => {
344        path.push((self, 0));
345        path
346      }
347      Some(ref node) => {
348        path.push((self, 0));
349        node.path_first(path)
350      }
351    }
352  }
353
354  pub(crate) fn path_last<'a, BK>(
355    &'a self,
356    mut path: Vec<(&'a Node<A>, usize)>,
357  ) -> Vec<(&'a Node<A>, usize)>
358  where
359    A: 'a,
360    BK: Ord + ?Sized,
361    A::Key: Borrow<BK>,
362  {
363    if self.keys.is_empty() {
364      return Vec::new();
365    }
366    let end = self.children.len() - 1;
367    match self.children[end] {
368      None => {
369        path.push((self, end - 1));
370        path
371      }
372      Some(ref node) => {
373        path.push((self, end));
374        node.path_last(path)
375      }
376    }
377  }
378
379  pub(crate) fn path_next<'a, BK>(
380    &'a self,
381    key: &BK,
382    mut path: Vec<(&'a Node<A>, usize)>,
383  ) -> Vec<(&'a Node<A>, usize)>
384  where
385    A: 'a,
386    BK: Ord + ?Sized,
387    A::Key: Borrow<BK>,
388  {
389    if self.keys.is_empty() {
390      return Vec::new();
391    }
392    match A::search_key(&self.keys, key) {
393      Ok(index) => {
394        path.push((self, index));
395        path
396      }
397      Err(index) => match self.children[index] {
398        None => match self.keys.get(index) {
399          Some(_) => {
400            path.push((self, index));
401            path
402          }
403          None => Vec::new(),
404        },
405        Some(ref node) => {
406          path.push((self, index));
407          node.path_next(key, path)
408        }
409      },
410    }
411  }
412
413  pub(crate) fn path_prev<'a, BK>(
414    &'a self,
415    key: &BK,
416    mut path: Vec<(&'a Node<A>, usize)>,
417  ) -> Vec<(&'a Node<A>, usize)>
418  where
419    A: 'a,
420    BK: Ord + ?Sized,
421    A::Key: Borrow<BK>,
422  {
423    if self.keys.is_empty() {
424      return Vec::new();
425    }
426    match A::search_key(&self.keys, key) {
427      Ok(index) => {
428        path.push((self, index));
429        path
430      }
431      Err(index) => match self.children[index] {
432        None if index == 0 => Vec::new(),
433        None => match self.keys.get(index - 1) {
434          Some(_) => {
435            path.push((self, index - 1));
436            path
437          }
438          None => Vec::new(),
439        },
440        Some(ref node) => {
441          path.push((self, index));
442          node.path_prev(key, path)
443        }
444      },
445    }
446  }
447
448  fn split(
449    &mut self,
450    pool: &Pool<Node<A>>,
451    value: A,
452    ins_left: Option<Node<A>>,
453    ins_right: Option<Node<A>>,
454  ) -> Insert<A> {
455    let left_child = ins_left.map(|node| PoolRef::new(pool, node));
456    let right_child = ins_right.map(|node| PoolRef::new(pool, node));
457    let index = A::search_value(&self.keys, &value).unwrap_err();
458    let mut left_keys;
459    let mut left_children;
460    let mut right_keys;
461    let mut right_children;
462    let median;
463    match index.cmp(&MEDIAN) {
464      Ordering::Less => {
465        self.children[index] = left_child;
466
467        left_keys = Chunk::from_front(&mut self.keys, index);
468        left_keys.push_back(value);
469        left_keys.drain_from_front(&mut self.keys, MEDIAN - index - 1);
470
471        left_children = Chunk::from_front(&mut self.children, index + 1);
472        left_children.push_back(right_child);
473        left_children.drain_from_front(&mut self.children, MEDIAN - index - 1);
474
475        median = self.keys.pop_front();
476
477        right_keys = Chunk::drain_from(&mut self.keys);
478        right_children = Chunk::drain_from(&mut self.children);
479      }
480      Ordering::Greater => {
481        self.children[index] = left_child;
482
483        left_keys = Chunk::from_front(&mut self.keys, MEDIAN);
484        left_children = Chunk::from_front(&mut self.children, MEDIAN + 1);
485
486        median = self.keys.pop_front();
487
488        right_keys = Chunk::from_front(&mut self.keys, index - MEDIAN - 1);
489        right_keys.push_back(value);
490        right_keys.append(&mut self.keys);
491
492        right_children = Chunk::from_front(&mut self.children, index - MEDIAN);
493        right_children.push_back(right_child);
494        right_children.append(&mut self.children);
495      }
496      Ordering::Equal => {
497        left_keys = Chunk::from_front(&mut self.keys, MEDIAN);
498        left_children = Chunk::from_front(&mut self.children, MEDIAN);
499        left_children.push_back(left_child);
500
501        median = value;
502
503        right_keys = Chunk::drain_from(&mut self.keys);
504        right_children = Chunk::drain_from(&mut self.children);
505        right_children[0] = right_child;
506      }
507    }
508
509    debug_assert!(left_keys.len() == MEDIAN);
510    debug_assert!(left_children.len() == MEDIAN + 1);
511    debug_assert!(right_keys.len() == MEDIAN);
512    debug_assert!(right_children.len() == MEDIAN + 1);
513
514    Split(Node { keys: left_keys, children: left_children }, median, Node {
515      keys: right_keys,
516      children: right_children,
517    })
518  }
519
520  fn merge(middle: A, left: Node<A>, mut right: Node<A>) -> Node<A> {
521    let mut keys = left.keys;
522    keys.push_back(middle);
523    keys.append(&mut right.keys);
524    let mut children = left.children;
525    children.append(&mut right.children);
526    Node { keys, children }
527  }
528
529  fn pop_min(&mut self) -> (A, Option<PoolRef<Node<A>>>) {
530    let value = self.keys.pop_front();
531    let child = self.children.pop_front();
532    (value, child)
533  }
534
535  fn pop_max(&mut self) -> (A, Option<PoolRef<Node<A>>>) {
536    let value = self.keys.pop_back();
537    let child = self.children.pop_back();
538    (value, child)
539  }
540
541  fn push_min(&mut self, child: Option<PoolRef<Node<A>>>, value: A) {
542    self.keys.push_front(value);
543    self.children.push_front(child);
544  }
545
546  fn push_max(&mut self, child: Option<PoolRef<Node<A>>>, value: A) {
547    self.keys.push_back(value);
548    self.children.push_back(child);
549  }
550
551  pub(crate) fn insert(&mut self, pool: &Pool<Node<A>>, value: A) -> Insert<A>
552  where A: Clone {
553    if self.keys.is_empty() {
554      self.keys.push_back(value);
555      self.children.push_back(None);
556      return Insert::Added;
557    }
558    let (median, left, right) = match A::search_value(&self.keys, &value) {
559      // Key exists in node
560      Ok(index) => {
561        return Insert::Replaced(mem::replace(&mut self.keys[index], value));
562      }
563      // Key is adjacent to some key in node
564      Err(index) => {
565        let has_room = self.has_room();
566        let action = match self.children[index] {
567          // No child at location, this is the target node.
568          None => InsertAt,
569          // Child at location, pass it on.
570          Some(ref mut child_ref) => {
571            let child = PoolRef::make_mut(pool, child_ref);
572            match child.insert(pool, value.clone()) {
573              Insert::Added => AddedAction,
574              Insert::Replaced(value) => ReplacedAction(value),
575              Insert::Split(left, median, right) => {
576                InsertSplit(left, median, right)
577              }
578            }
579          }
580        };
581        match action {
582          ReplacedAction(value) => return Insert::Replaced(value),
583          AddedAction => {
584            return Insert::Added;
585          }
586          InsertAt => {
587            if has_room {
588              self.keys.insert(index, value);
589              self.children.insert(index + 1, None);
590              return Insert::Added;
591            }
592            else {
593              (value, None, None)
594            }
595          }
596          InsertSplit(left, median, right) => {
597            if has_room {
598              self.children[index] = Some(PoolRef::new(pool, left));
599              self.keys.insert(index, median);
600              self.children.insert(index + 1, Some(PoolRef::new(pool, right)));
601              return Insert::Added;
602            }
603            else {
604              (median, Some(left), Some(right))
605            }
606          }
607        }
608      }
609    };
610    self.split(pool, median, left, right)
611  }
612
613  pub(crate) fn remove<BK>(
614    &mut self,
615    pool: &Pool<Node<A>>,
616    key: &BK,
617  ) -> Remove<A>
618  where
619    A: Clone,
620    BK: Ord + ?Sized,
621    A::Key: Borrow<BK>,
622  {
623    let index = A::search_key(&self.keys, key);
624    self.remove_index(pool, index, key)
625  }
626
627  fn remove_index<BK>(
628    &mut self,
629    pool: &Pool<Node<A>>,
630    index: Result<usize, usize>,
631    key: &BK,
632  ) -> Remove<A>
633  where
634    A: Clone,
635    BK: Ord + ?Sized,
636    A::Key: Borrow<BK>,
637  {
638    let action = match index {
639      // Key exists in node, remove it.
640      Ok(index) => {
641        match (&self.children[index], &self.children[index + 1]) {
642          // If we're a leaf, just delete the entry.
643          (&None, &None) => RemoveAction::DeleteAt(index),
644          // Right is empty. Attempt to steal from left if enough capacity,
645          // otherwise pull the predecessor up.
646          (&Some(ref left), &None) => {
647            if !left.too_small() {
648              RemoveAction::StealFromLeft(index + 1)
649            }
650            else {
651              RemoveAction::PullUp(left.keys.len() - 1, index, index)
652            }
653          }
654          // Left is empty. Attempt to steal from right if enough capacity,
655          // otherwise pull the successor up.
656          (&None, &Some(ref right)) => {
657            if !right.too_small() {
658              RemoveAction::StealFromRight(index)
659            }
660            else {
661              RemoveAction::PullUp(0, index, index + 1)
662            }
663          }
664          // Both left and right are non empty. Attempt to steal from left or
665          // right if enough capacity, otherwise just merge the children.
666          (&Some(ref left), &Some(ref right)) => {
667            if left.has_room() && !right.too_small() {
668              RemoveAction::StealFromRight(index)
669            }
670            else if right.has_room() && !left.too_small() {
671              RemoveAction::StealFromLeft(index + 1)
672            }
673            else {
674              RemoveAction::Merge(index)
675            }
676          }
677        }
678      }
679      // Key is adjacent to some key in node
680      Err(index) => match self.children[index] {
681        // No child at location means key isn't in map.
682        None => return Remove::NoChange,
683        // Child at location, but it's at minimum capacity.
684        Some(ref child) if child.too_small() => {
685          let left =
686            if index > 0 { self.children.get(index - 1) } else { None }; // index is usize and can't be negative, best make sure it never is.
687          match (left, self.children.get(index + 1)) {
688            // If it has a left sibling with capacity, steal a key from it.
689            (Some(&Some(ref old_left)), _) if !old_left.too_small() => {
690              RemoveAction::StealFromLeft(index)
691            }
692            // If it has a right sibling with capacity, same as above.
693            (_, Some(&Some(ref old_right))) if !old_right.too_small() => {
694              RemoveAction::StealFromRight(index)
695            }
696            // If it has neither, we'll have to merge it with a sibling.
697            // If we have a right sibling, we'll merge with that.
698            (_, Some(&Some(_))) => RemoveAction::MergeFirst(index),
699            // If we have a left sibling, we'll merge with that.
700            (Some(&Some(_)), _) => RemoveAction::MergeFirst(index - 1),
701            // If none of the above, we're in a bad state.
702            _ => unreachable!(),
703          }
704        }
705        // Child at location, and it's big enough, we can recurse down.
706        Some(_) => RemoveAction::ContinueDown(index),
707      },
708    };
709    match action {
710      RemoveAction::DeleteAt(index) => {
711        let pair = self.keys.remove(index);
712        self.children.remove(index);
713        Remove::Removed(pair)
714      }
715      RemoveAction::PullUp(target_index, pull_to, child_index) => {
716        let children = &mut self.children;
717        let mut update = None;
718        let value;
719        if let Some(&mut Some(ref mut child_ref)) =
720          children.get_mut(child_index)
721        {
722          let child = PoolRef::make_mut(pool, child_ref);
723          match child.remove_index(pool, Ok(target_index), key) {
724            Remove::NoChange => unreachable!(),
725            Remove::Removed(pulled_value) => {
726              value = self.keys.set(pull_to, pulled_value);
727            }
728            Remove::Update(pulled_value, new_child) => {
729              value = self.keys.set(pull_to, pulled_value);
730              update = Some(new_child);
731            }
732          }
733        }
734        else {
735          unreachable!()
736        }
737        if let Some(new_child) = update {
738          children[child_index] = Some(PoolRef::new(pool, new_child));
739        }
740        Remove::Removed(value)
741      }
742      RemoveAction::Merge(index) => {
743        let left = self.children.remove(index).unwrap();
744        let right = mem::replace(&mut self.children[index], None).unwrap();
745        let value = self.keys.remove(index);
746        let mut merged_child = Node::merge(
747          value,
748          PoolRef::unwrap_or_clone(left),
749          PoolRef::unwrap_or_clone(right),
750        );
751        let (removed, new_child) = match merged_child.remove(pool, key) {
752          Remove::NoChange => unreachable!(),
753          Remove::Removed(removed) => (removed, merged_child),
754          Remove::Update(removed, updated_child) => (removed, updated_child),
755        };
756        if self.keys.is_empty() {
757          // If we've depleted the root node, the merged child becomes the root.
758          Remove::Update(removed, new_child)
759        }
760        else {
761          self.children[index] = Some(PoolRef::new(pool, new_child));
762          Remove::Removed(removed)
763        }
764      }
765      RemoveAction::StealFromLeft(index) => {
766        let mut update = None;
767        let out_value;
768        {
769          let mut children = self.children.as_mut_slice()[index - 1..=index]
770            .iter_mut()
771            .map(|n| if let Some(ref mut o) = *n { o } else { unreachable!() });
772          let left = PoolRef::make_mut(pool, children.next().unwrap());
773          let child = PoolRef::make_mut(pool, children.next().unwrap());
774          // Prepare the rebalanced node.
775          child.push_min(
776            left.children.last().unwrap().clone(),
777            self.keys[index - 1].clone(),
778          );
779          match child.remove(pool, key) {
780            Remove::NoChange => {
781              // Key wasn't there, we need to revert the steal.
782              child.pop_min();
783              return Remove::NoChange;
784            }
785            Remove::Removed(value) => {
786              // If we did remove something, we complete the rebalancing.
787              let (left_value, _) = left.pop_max();
788              self.keys[index - 1] = left_value;
789              out_value = value;
790            }
791            Remove::Update(value, new_child) => {
792              // If we did remove something, we complete the rebalancing.
793              let (left_value, _) = left.pop_max();
794              self.keys[index - 1] = left_value;
795              update = Some(new_child);
796              out_value = value;
797            }
798          }
799        }
800        if let Some(new_child) = update {
801          self.children[index] = Some(PoolRef::new(pool, new_child));
802        }
803        Remove::Removed(out_value)
804      }
805      RemoveAction::StealFromRight(index) => {
806        let mut update = None;
807        let out_value;
808        {
809          let mut children = self.children.as_mut_slice()[index..index + 2]
810            .iter_mut()
811            .map(|n| if let Some(ref mut o) = *n { o } else { unreachable!() });
812          let child = PoolRef::make_mut(pool, children.next().unwrap());
813          let right = PoolRef::make_mut(pool, children.next().unwrap());
814          // Prepare the rebalanced node.
815          child.push_max(right.children[0].clone(), self.keys[index].clone());
816          match child.remove(pool, key) {
817            Remove::NoChange => {
818              // Key wasn't there, we need to revert the steal.
819              child.pop_max();
820              return Remove::NoChange;
821            }
822            Remove::Removed(value) => {
823              // If we did remove something, we complete the rebalancing.
824              let (right_value, _) = right.pop_min();
825              self.keys[index] = right_value;
826              out_value = value;
827            }
828            Remove::Update(value, new_child) => {
829              // If we did remove something, we complete the rebalancing.
830              let (right_value, _) = right.pop_min();
831              self.keys[index] = right_value;
832              update = Some(new_child);
833              out_value = value;
834            }
835          }
836        }
837        if let Some(new_child) = update {
838          self.children[index] = Some(PoolRef::new(pool, new_child));
839        }
840        Remove::Removed(out_value)
841      }
842      RemoveAction::MergeFirst(index) => {
843        if self.keys[index].cmp_keys(key) != Ordering::Equal
844          && !self.child_contains(index, key)
845          && !self.child_contains(index + 1, key)
846        {
847          return Remove::NoChange;
848        }
849        let left = self.children.remove(index).unwrap();
850        let right = mem::replace(&mut self.children[index], None).unwrap();
851        let middle = self.keys.remove(index);
852        let mut merged = Node::merge(
853          middle,
854          PoolRef::unwrap_or_clone(left),
855          PoolRef::unwrap_or_clone(right),
856        );
857        let update;
858        let out_value;
859        match merged.remove(pool, key) {
860          Remove::NoChange => {
861            panic!(
862              "nodes::btree::Node::remove: caught an absent key too late \
863               while merging"
864            );
865          }
866          Remove::Removed(value) => {
867            if self.keys.is_empty() {
868              return Remove::Update(value, merged);
869            }
870            update = merged;
871            out_value = value;
872          }
873          Remove::Update(value, new_child) => {
874            if self.keys.is_empty() {
875              return Remove::Update(value, new_child);
876            }
877            update = new_child;
878            out_value = value;
879          }
880        }
881        self.children[index] = Some(PoolRef::new(pool, update));
882        Remove::Removed(out_value)
883      }
884      RemoveAction::ContinueDown(index) => {
885        let mut update = None;
886        let out_value;
887        if let Some(&mut Some(ref mut child_ref)) = self.children.get_mut(index)
888        {
889          let child = PoolRef::make_mut(pool, child_ref);
890          match child.remove(pool, key) {
891            Remove::NoChange => return Remove::NoChange,
892            Remove::Removed(value) => {
893              out_value = value;
894            }
895            Remove::Update(value, new_child) => {
896              update = Some(new_child);
897              out_value = value;
898            }
899          }
900        }
901        else {
902          unreachable!()
903        }
904        if let Some(new_child) = update {
905          self.children[index] = Some(PoolRef::new(pool, new_child));
906        }
907        Remove::Removed(out_value)
908      }
909    }
910  }
911}
912
913// Iterator
914
915/// An iterator over an ordered set.
916pub struct Iter<'a, A> {
917  fwd_path: Vec<(&'a Node<A>, usize)>,
918  back_path: Vec<(&'a Node<A>, usize)>,
919  pub(crate) remaining: usize,
920}
921
922impl<'a, A: BTreeValue> Iter<'a, A> {
923  pub(crate) fn new<R, BK>(root: &'a Node<A>, size: usize, range: R) -> Self
924  where
925    R: RangeBounds<BK>,
926    A::Key: Borrow<BK>,
927    BK: Ord + ?Sized, {
928    let fwd_path = match range.start_bound() {
929      Bound::Included(key) => root.path_next(key, Vec::new()),
930      Bound::Excluded(key) => {
931        let mut path = root.path_next(key, Vec::new());
932        if let Some(value) = Self::get(&path) {
933          if value.cmp_keys(key) == Ordering::Equal {
934            Self::step_forward(&mut path);
935          }
936        }
937        path
938      }
939      Bound::Unbounded => root.path_first(Vec::new()),
940    };
941    let back_path = match range.end_bound() {
942      Bound::Included(key) => root.path_prev(key, Vec::new()),
943      Bound::Excluded(key) => {
944        let mut path = root.path_prev(key, Vec::new());
945        if let Some(value) = Self::get(&path) {
946          if value.cmp_keys(key) == Ordering::Equal {
947            Self::step_back(&mut path);
948          }
949        }
950        path
951      }
952      Bound::Unbounded => root.path_last(Vec::new()),
953    };
954    Iter { fwd_path, back_path, remaining: size }
955  }
956
957  fn get(path: &[(&'a Node<A>, usize)]) -> Option<&'a A> {
958    match path.last() {
959      Some((node, index)) => Some(&node.keys[*index]),
960      None => None,
961    }
962  }
963
964  fn step_forward(path: &mut Vec<(&'a Node<A>, usize)>) -> Option<&'a A> {
965    match path.pop() {
966      Some((node, index)) => {
967        let index = index + 1;
968        match node.children[index] {
969          // Child between current and next key -> step down
970          Some(ref child) => {
971            path.push((node, index));
972            path.push((child, 0));
973            let mut node = child;
974            while let Some(ref left_child) = node.children[0] {
975              path.push((left_child, 0));
976              node = left_child;
977            }
978            Some(&node.keys[0])
979          }
980          None => match node.keys.get(index) {
981            // Yield next key
982            value @ Some(_) => {
983              path.push((node, index));
984              value
985            }
986            // No more keys -> exhausted level, step up and yield
987            None => loop {
988              match path.pop() {
989                None => {
990                  return None;
991                }
992                Some((node, index)) => {
993                  if let value @ Some(_) = node.keys.get(index) {
994                    path.push((node, index));
995                    return value;
996                  }
997                }
998              }
999            },
1000          },
1001        }
1002      }
1003      None => None,
1004    }
1005  }
1006
1007  fn step_back(path: &mut Vec<(&'a Node<A>, usize)>) -> Option<&'a A> {
1008    match path.pop() {
1009      Some((node, index)) => match node.children[index] {
1010        Some(ref child) => {
1011          path.push((node, index));
1012          let mut end = child.keys.len() - 1;
1013          path.push((child, end));
1014          let mut node = child;
1015          while let Some(ref right_child) = node.children[end + 1] {
1016            end = right_child.keys.len() - 1;
1017            path.push((right_child, end));
1018            node = right_child;
1019          }
1020          Some(&node.keys[end])
1021        }
1022        None => {
1023          if index == 0 {
1024            loop {
1025              match path.pop() {
1026                None => {
1027                  return None;
1028                }
1029                Some((node, index)) => {
1030                  if index > 0 {
1031                    let index = index - 1;
1032                    path.push((node, index));
1033                    return Some(&node.keys[index]);
1034                  }
1035                }
1036              }
1037            }
1038          }
1039          else {
1040            let index = index - 1;
1041            path.push((node, index));
1042            Some(&node.keys[index])
1043          }
1044        }
1045      },
1046      None => None,
1047    }
1048  }
1049}
1050
1051impl<'a, A: 'a + BTreeValue> Iterator for Iter<'a, A> {
1052  type Item = &'a A;
1053
1054  fn next(&mut self) -> Option<Self::Item> {
1055    match Iter::get(&self.fwd_path) {
1056      None => None,
1057      Some(value) => match Iter::get(&self.back_path) {
1058        Some(last_value)
1059          if value.cmp_values(last_value) == Ordering::Greater =>
1060        {
1061          None
1062        }
1063        None => None,
1064        Some(_) => {
1065          Iter::step_forward(&mut self.fwd_path);
1066          self.remaining -= 1;
1067          Some(value)
1068        }
1069      },
1070    }
1071  }
1072
1073  fn size_hint(&self) -> (usize, Option<usize>) {
1074    // (0, Some(self.remaining))
1075    (0, None)
1076  }
1077}
1078
1079impl<'a, A: 'a + BTreeValue> DoubleEndedIterator for Iter<'a, A> {
1080  fn next_back(&mut self) -> Option<Self::Item> {
1081    match Iter::get(&self.back_path) {
1082      None => None,
1083      Some(value) => match Iter::get(&self.fwd_path) {
1084        Some(last_value) if value.cmp_values(last_value) == Ordering::Less => {
1085          None
1086        }
1087        None => None,
1088        Some(_) => {
1089          Iter::step_back(&mut self.back_path);
1090          self.remaining -= 1;
1091          Some(value)
1092        }
1093      },
1094    }
1095  }
1096}
1097
1098// Consuming iterator
1099
1100enum ConsumingIterItem<A> {
1101  Consider(Node<A>),
1102  Yield(A),
1103}
1104
1105/// A consuming iterator over an ordered set.
1106pub struct ConsumingIter<A> {
1107  fwd_last: Option<A>,
1108  fwd_stack: Vec<ConsumingIterItem<A>>,
1109  back_last: Option<A>,
1110  back_stack: Vec<ConsumingIterItem<A>>,
1111  remaining: usize,
1112}
1113
1114impl<A: Clone> ConsumingIter<A> {
1115  pub(crate) fn new(root: &Node<A>, total: usize) -> Self {
1116    ConsumingIter {
1117      fwd_last: None,
1118      fwd_stack: vec![ConsumingIterItem::Consider(root.clone())],
1119      back_last: None,
1120      back_stack: vec![ConsumingIterItem::Consider(root.clone())],
1121      remaining: total,
1122    }
1123  }
1124
1125  fn push_node(
1126    stack: &mut Vec<ConsumingIterItem<A>>,
1127    maybe_node: Option<PoolRef<Node<A>>>,
1128  ) {
1129    if let Some(node) = maybe_node {
1130      stack.push(ConsumingIterItem::Consider(PoolRef::unwrap_or_clone(node)))
1131    }
1132  }
1133
1134  fn push(stack: &mut Vec<ConsumingIterItem<A>>, mut node: Node<A>) {
1135    for _n in 0..node.keys.len() {
1136      ConsumingIter::push_node(stack, node.children.pop_back());
1137      stack.push(ConsumingIterItem::Yield(node.keys.pop_back()));
1138    }
1139    ConsumingIter::push_node(stack, node.children.pop_back());
1140  }
1141
1142  fn push_fwd(&mut self, node: Node<A>) {
1143    ConsumingIter::push(&mut self.fwd_stack, node)
1144  }
1145
1146  fn push_node_back(&mut self, maybe_node: Option<PoolRef<Node<A>>>) {
1147    if let Some(node) = maybe_node {
1148      self
1149        .back_stack
1150        .push(ConsumingIterItem::Consider(PoolRef::unwrap_or_clone(node)))
1151    }
1152  }
1153
1154  fn push_back(&mut self, mut node: Node<A>) {
1155    for _i in 0..node.keys.len() {
1156      self.push_node_back(node.children.pop_front());
1157      self.back_stack.push(ConsumingIterItem::Yield(node.keys.pop_front()));
1158    }
1159    self.push_node_back(node.children.pop_back());
1160  }
1161}
1162
1163impl<A> Iterator for ConsumingIter<A>
1164where A: BTreeValue + Clone
1165{
1166  type Item = A;
1167
1168  fn next(&mut self) -> Option<Self::Item> {
1169    loop {
1170      match self.fwd_stack.pop() {
1171        None => {
1172          self.remaining = 0;
1173          return None;
1174        }
1175        Some(ConsumingIterItem::Consider(node)) => self.push_fwd(node),
1176        Some(ConsumingIterItem::Yield(value)) => {
1177          if let Some(ref last) = self.back_last {
1178            if value.cmp_values(last) != Ordering::Less {
1179              self.fwd_stack.clear();
1180              self.back_stack.clear();
1181              self.remaining = 0;
1182              return None;
1183            }
1184          }
1185          self.remaining -= 1;
1186          self.fwd_last = Some(value.clone());
1187          return Some(value);
1188        }
1189      }
1190    }
1191  }
1192
1193  fn size_hint(&self) -> (usize, Option<usize>) {
1194    (self.remaining, Some(self.remaining))
1195  }
1196}
1197
1198impl<A> DoubleEndedIterator for ConsumingIter<A>
1199where A: BTreeValue + Clone
1200{
1201  fn next_back(&mut self) -> Option<Self::Item> {
1202    loop {
1203      match self.back_stack.pop() {
1204        None => {
1205          self.remaining = 0;
1206          return None;
1207        }
1208        Some(ConsumingIterItem::Consider(node)) => self.push_back(node),
1209        Some(ConsumingIterItem::Yield(value)) => {
1210          if let Some(ref last) = self.fwd_last {
1211            if value.cmp_values(last) != Ordering::Greater {
1212              self.fwd_stack.clear();
1213              self.back_stack.clear();
1214              self.remaining = 0;
1215              return None;
1216            }
1217          }
1218          self.remaining -= 1;
1219          self.back_last = Some(value.clone());
1220          return Some(value);
1221        }
1222      }
1223    }
1224  }
1225}
1226
1227impl<A: BTreeValue + Clone> ExactSizeIterator for ConsumingIter<A> {}
1228
1229// DiffIter
1230
1231/// An iterator over the differences between two ordered sets.
1232pub struct DiffIter<'a, A> {
1233  old_stack: Vec<IterItem<'a, A>>,
1234  new_stack: Vec<IterItem<'a, A>>,
1235}
1236
1237/// A description of a difference between two ordered sets.
1238#[derive(PartialEq, Eq, Debug)]
1239pub enum DiffItem<'a, A> {
1240  /// This value has been added to the new set.
1241  Add(&'a A),
1242  /// This value has been changed between the two sets.
1243  Update {
1244    /// The old value.
1245    old: &'a A,
1246    /// The new value.
1247    new: &'a A,
1248  },
1249  /// This value has been removed from the new set.
1250  Remove(&'a A),
1251}
1252
1253enum IterItem<'a, A> {
1254  Consider(&'a Node<A>),
1255  Yield(&'a A),
1256}
1257
1258impl<'a, A: 'a> DiffIter<'a, A> {
1259  pub(crate) fn new(old: &'a Node<A>, new: &'a Node<A>) -> Self {
1260    DiffIter {
1261      old_stack: if old.keys.is_empty() {
1262        Vec::new()
1263      }
1264      else {
1265        vec![IterItem::Consider(old)]
1266      },
1267      new_stack: if new.keys.is_empty() {
1268        Vec::new()
1269      }
1270      else {
1271        vec![IterItem::Consider(new)]
1272      },
1273    }
1274  }
1275
1276  fn push_node(
1277    stack: &mut Vec<IterItem<'a, A>>,
1278    maybe_node: &'a Option<PoolRef<Node<A>>>,
1279  ) {
1280    if let Some(ref node) = *maybe_node {
1281      stack.push(IterItem::Consider(&node))
1282    }
1283  }
1284
1285  fn push(stack: &mut Vec<IterItem<'a, A>>, node: &'a Node<A>) {
1286    for n in 0..node.keys.len() {
1287      let i = node.keys.len() - n;
1288      Self::push_node(stack, &node.children[i]);
1289      stack.push(IterItem::Yield(&node.keys[i - 1]));
1290    }
1291    Self::push_node(stack, &node.children[0]);
1292  }
1293}
1294
1295impl<'a, A> Iterator for DiffIter<'a, A>
1296where A: 'a + BTreeValue + PartialEq
1297{
1298  type Item = DiffItem<'a, A>;
1299
1300  fn next(&mut self) -> Option<Self::Item> {
1301    loop {
1302      match (self.old_stack.pop(), self.new_stack.pop()) {
1303        (None, None) => return None,
1304        (None, Some(new)) => match new {
1305          IterItem::Consider(new) => Self::push(&mut self.new_stack, &new),
1306          IterItem::Yield(new) => return Some(DiffItem::Add(new)),
1307        },
1308        (Some(old), None) => match old {
1309          IterItem::Consider(old) => Self::push(&mut self.old_stack, &old),
1310          IterItem::Yield(old) => return Some(DiffItem::Remove(old)),
1311        },
1312        (Some(old), Some(new)) => match (old, new) {
1313          (IterItem::Consider(old), IterItem::Consider(new)) => {
1314            if !core::ptr::eq(old, new) {
1315              match old.keys[0].cmp_values(&new.keys[0]) {
1316                Ordering::Less => {
1317                  Self::push(&mut self.old_stack, &old);
1318                  self.new_stack.push(IterItem::Consider(new));
1319                }
1320                Ordering::Greater => {
1321                  self.old_stack.push(IterItem::Consider(old));
1322                  Self::push(&mut self.new_stack, &new);
1323                }
1324                Ordering::Equal => {
1325                  Self::push(&mut self.old_stack, &old);
1326                  Self::push(&mut self.new_stack, &new);
1327                }
1328              }
1329            }
1330          }
1331          (IterItem::Consider(old), IterItem::Yield(new)) => {
1332            Self::push(&mut self.old_stack, &old);
1333            self.new_stack.push(IterItem::Yield(new));
1334          }
1335          (IterItem::Yield(old), IterItem::Consider(new)) => {
1336            self.old_stack.push(IterItem::Yield(old));
1337            Self::push(&mut self.new_stack, &new);
1338          }
1339          (IterItem::Yield(old), IterItem::Yield(new)) => {
1340            match old.cmp_values(&new) {
1341              Ordering::Less => {
1342                self.new_stack.push(IterItem::Yield(new));
1343                return Some(DiffItem::Remove(old));
1344              }
1345              Ordering::Equal => {
1346                if old != new {
1347                  return Some(DiffItem::Update { old, new });
1348                }
1349              }
1350              Ordering::Greater => {
1351                self.old_stack.push(IterItem::Yield(old));
1352                return Some(DiffItem::Add(new));
1353              }
1354            }
1355          }
1356        },
1357      }
1358    }
1359  }
1360}
sp_im/nodes/btree.rs

sp_im/nodes/
btree.rs
