library/alloc/src/collections/btree/navigate.rs

   1 use core::borrow::Borrow;
   2 use core::cmp::Ordering;
   3 use core::intrinsics;
   4 use core::mem;
   5 use core::ops::Bound::{Excluded, Included, Unbounded};
   6 use core::ops::RangeBounds;
   7 use core::ptr;
   8
   9 use super::node::{marker, ForceResult::*, Handle, NodeRef};
  10 use super::search::{self, SearchResult};
  11 use super::unwrap_unchecked;
  12
  13 /// Finds the leaf edges delimiting a specified range in or underneath a node.
  14 fn range_search<BorrowType, K, V, Q, R>(
  15     root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
  16     root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
  17     range: R,
  18 ) -> (
  19     Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
  20     Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
  21 )
  22 where
  23     Q: ?Sized + Ord,
  24     K: Borrow<Q>,
  25     R: RangeBounds<Q>,
  26 {
  27     match (range.start_bound(), range.end_bound()) {
  28         (Excluded(s), Excluded(e)) if s == e => {
  29             panic!("range start and end are equal and excluded in BTreeMap")
  30         }
  31         (Included(s) | Excluded(s), Included(e) | Excluded(e)) if s > e => {
  32             panic!("range start is greater than range end in BTreeMap")
  33         }
  34         _ => {}
  35     };
  36
  37     let mut min_node = root1;
  38     let mut max_node = root2;
  39     let mut min_found = false;
  40     let mut max_found = false;
  41
  42     loop {
  43         let front = match (min_found, range.start_bound()) {
  44             (false, Included(key)) => match search::search_node(min_node, key) {
  45                 SearchResult::Found(kv) => {
  46                     min_found = true;
  47                     kv.left_edge()
  48                 }
  49                 SearchResult::GoDown(edge) => edge,
  50             },
  51             (false, Excluded(key)) => match search::search_node(min_node, key) {
  52                 SearchResult::Found(kv) => {
  53                     min_found = true;
  54                     kv.right_edge()
  55                 }
  56                 SearchResult::GoDown(edge) => edge,
  57             },
  58             (true, Included(_)) => min_node.last_edge(),
  59             (true, Excluded(_)) => min_node.first_edge(),
  60             (_, Unbounded) => min_node.first_edge(),
  61         };
  62
  63         let back = match (max_found, range.end_bound()) {
  64             (false, Included(key)) => match search::search_node(max_node, key) {
  65                 SearchResult::Found(kv) => {
  66                     max_found = true;
  67                     kv.right_edge()
  68                 }
  69                 SearchResult::GoDown(edge) => edge,
  70             },
  71             (false, Excluded(key)) => match search::search_node(max_node, key) {
  72                 SearchResult::Found(kv) => {
  73                     max_found = true;
  74                     kv.left_edge()
  75                 }
  76                 SearchResult::GoDown(edge) => edge,
  77             },
  78             (true, Included(_)) => max_node.first_edge(),
  79             (true, Excluded(_)) => max_node.last_edge(),
  80             (_, Unbounded) => max_node.last_edge(),
  81         };
  82
  83         if front.partial_cmp(&back) == Some(Ordering::Greater) {
  84             panic!("Ord is ill-defined in BTreeMap range");
  85         }
  86         match (front.force(), back.force()) {
  87             (Leaf(f), Leaf(b)) => {
  88                 return (f, b);
  89             }
  90             (Internal(min_int), Internal(max_int)) => {
  91                 min_node = min_int.descend();
  92                 max_node = max_int.descend();
  93             }
  94             _ => unreachable!("BTreeMap has different depths"),
  95         };
  96     }
  97 }
  98
  99 /// Equivalent to `range_search(k, v, ..)` but without the `Ord` bound.
 100 fn full_range<BorrowType, K, V>(
 101     root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
 102     root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
 103 ) -> (
 104     Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
 105     Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
 106 ) {
 107     let mut min_node = root1;
 108     let mut max_node = root2;
 109     loop {
 110         let front = min_node.first_edge();
 111         let back = max_node.last_edge();
 112         match (front.force(), back.force()) {
 113             (Leaf(f), Leaf(b)) => {
 114                 return (f, b);
 115             }
 116             (Internal(min_int), Internal(max_int)) => {
 117                 min_node = min_int.descend();
 118                 max_node = max_int.descend();
 119             }
 120             _ => unreachable!("BTreeMap has different depths"),
 121         };
 122     }
 123 }
 124
 125 impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal> {
 126     /// Creates a pair of leaf edges delimiting a specified range in or underneath a node.
 127     pub fn range_search<Q, R>(
 128         self,
 129         range: R,
 130     ) -> (
 131         Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
 132         Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
 133     )
 134     where
 135         Q: ?Sized + Ord,
 136         K: Borrow<Q>,
 137         R: RangeBounds<Q>,
 138     {
 139         range_search(self, self, range)
 140     }
 141
 142     /// Returns (self.first_leaf_edge(), self.last_leaf_edge()), but more efficiently.
 143     pub fn full_range(
 144         self,
 145     ) -> (
 146         Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
 147         Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
 148     ) {
 149         full_range(self, self)
 150     }
 151 }
 152
 153 impl<'a, K: 'a, V: 'a> NodeRef<marker::ValMut<'a>, K, V, marker::LeafOrInternal> {
 154     /// Splits a unique reference into a pair of leaf edges delimiting a specified range.
 155     /// The result are non-unique references allowing (some) mutation, which must be used
 156     /// carefully.
 157     pub fn range_search<Q, R>(
 158         self,
 159         range: R,
 160     ) -> (
 161         Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
 162         Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
 163     )
 164     where
 165         Q: ?Sized + Ord,
 166         K: Borrow<Q>,
 167         R: RangeBounds<Q>,
 168     {
 169         // We duplicate the root NodeRef here -- we will never visit the same KV
 170         // twice, and never end up with overlapping value references.
 171         let self2 = unsafe { ptr::read(&self) };
 172         range_search(self, self2, range)
 173     }
 174
 175     /// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree.
 176     /// The results are non-unique references allowing mutation (of values only), so must be used
 177     /// with care.
 178     pub fn full_range(
 179         self,
 180     ) -> (
 181         Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
 182         Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
 183     ) {
 184         // We duplicate the root NodeRef here -- we will never visit the same KV
 185         // twice, and never end up with overlapping value references.
 186         let self2 = unsafe { ptr::read(&self) };
 187         full_range(self, self2)
 188     }
 189 }
 190
 191 impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
 192     /// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree.
 193     /// The results are non-unique references allowing massively destructive mutation, so must be
 194     /// used with the utmost care.
 195     pub fn full_range(
 196         self,
 197     ) -> (
 198         Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
 199         Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
 200     ) {
 201         // We duplicate the root NodeRef here -- we will never access it in a way
 202         // that overlaps references obtained from the root.
 203         let self2 = unsafe { ptr::read(&self) };
 204         full_range(self, self2)
 205     }
 206 }
 207
 208 impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
 209     /// Given a leaf edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
 210     /// on the right side, which is either in the same leaf node or in an ancestor node.
 211     /// If the leaf edge is the last one in the tree, returns [`Result::Err`] with the root node.
 212     pub fn next_kv(
 213         self,
 214     ) -> Result<
 215         Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>,
 216         NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
 217     > {
 218         let mut edge = self.forget_node_type();
 219         loop {
 220             edge = match edge.right_kv() {
 221                 Ok(kv) => return Ok(kv),
 222                 Err(last_edge) => match last_edge.into_node().ascend() {
 223                     Ok(parent_edge) => parent_edge.forget_node_type(),
 224                     Err(root) => return Err(root),
 225                 },
 226             }
 227         }
 228     }
 229
 230     /// Given a leaf edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
 231     /// on the left side, which is either in the same leaf node or in an ancestor node.
 232     /// If the leaf edge is the first one in the tree, returns [`Result::Err`] with the root node.
 233     pub fn next_back_kv(
 234         self,
 235     ) -> Result<
 236         Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>,
 237         NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
 238     > {
 239         let mut edge = self.forget_node_type();
 240         loop {
 241             edge = match edge.left_kv() {
 242                 Ok(kv) => return Ok(kv),
 243                 Err(last_edge) => match last_edge.into_node().ascend() {
 244                     Ok(parent_edge) => parent_edge.forget_node_type(),
 245                     Err(root) => return Err(root),
 246                 },
 247             }
 248         }
 249     }
 250 }
 251
 252 impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge> {
 253     /// Given an internal edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
 254     /// on the right side, which is either in the same internal node or in an ancestor node.
 255     /// If the internal edge is the last one in the tree, returns [`Result::Err`] with the root node.
 256     pub fn next_kv(
 257         self,
 258     ) -> Result<
 259         Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV>,
 260         NodeRef<BorrowType, K, V, marker::Internal>,
 261     > {
 262         let mut edge = self;
 263         loop {
 264             edge = match edge.right_kv() {
 265                 Ok(internal_kv) => return Ok(internal_kv),
 266                 Err(last_edge) => match last_edge.into_node().ascend() {
 267                     Ok(parent_edge) => parent_edge,
 268                     Err(root) => return Err(root),
 269                 },
 270             }
 271         }
 272     }
 273 }
 274
 275 macro_rules! def_next_kv_uncheched_dealloc {
 276     { unsafe fn $name:ident : $adjacent_kv:ident } => {
 277         /// Given a leaf edge handle into an owned tree, returns a handle to the next KV,
 278         /// while deallocating any node left behind yet leaving the corresponding edge
 279         /// in its parent node dangling.
 280         ///
 281         /// # Safety
 282         /// - The leaf edge must not be the last one in the direction travelled.
 283         /// - The node carrying the next KV returned must not have been deallocated by a
 284         ///   previous call on any handle obtained for this tree.
 285         unsafe fn $name <K, V>(
 286             leaf_edge: Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
 287         ) -> Handle<NodeRef<marker::Owned, K, V, marker::LeafOrInternal>, marker::KV> {
 288             let mut edge = leaf_edge.forget_node_type();
 289             loop {
 290                 edge = match edge.$adjacent_kv() {
 291                     Ok(internal_kv) => return internal_kv,
 292                     Err(last_edge) => {
 293                         unsafe {
 294                             let parent_edge = last_edge.into_node().deallocate_and_ascend();
 295                             unwrap_unchecked(parent_edge).forget_node_type()
 296                         }
 297                     }
 298                 }
 299             }
 300         }
 301     };
 302 }
 303
 304 def_next_kv_uncheched_dealloc! {unsafe fn next_kv_unchecked_dealloc: right_kv}
 305 def_next_kv_uncheched_dealloc! {unsafe fn next_back_kv_unchecked_dealloc: left_kv}
 306
 307 /// This replaces the value behind the `v` unique reference by calling the
 308 /// relevant function.
 309 ///
 310 /// If a panic occurs in the `change` closure, the entire process will be aborted.
 311 #[inline]
 312 fn take_mut<T>(v: &mut T, change: impl FnOnce(T) -> T) {
 313     replace(v, |value| (change(value), ()))
 314 }
 315
 316 /// This replaces the value behind the `v` unique reference by calling the
 317 /// relevant function, and returns a result obtained along the way.
 318 ///
 319 /// If a panic occurs in the `change` closure, the entire process will be aborted.
 320 #[inline]
 321 fn replace<T, R>(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R {
 322     struct PanicGuard;
 323     impl Drop for PanicGuard {
 324         fn drop(&mut self) {
 325             intrinsics::abort()
 326         }
 327     }
 328     let guard = PanicGuard;
 329     let value = unsafe { ptr::read(v) };
 330     let (new_value, ret) = change(value);
 331     unsafe {
 332         ptr::write(v, new_value);
 333     }
 334     mem::forget(guard);
 335     ret
 336 }
 337
 338 impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge> {
 339     /// Moves the leaf edge handle to the next leaf edge and returns references to the
 340     /// key and value in between.
 341     ///
 342     /// # Safety
 343     /// There must be another KV in the direction travelled.
 344     pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
 345         replace(self, |leaf_edge| {
 346             let kv = leaf_edge.next_kv();
 347             let kv = unsafe { unwrap_unchecked(kv.ok()) };
 348             (kv.next_leaf_edge(), kv.into_kv())
 349         })
 350     }
 351
 352     /// Moves the leaf edge handle to the previous leaf edge and returns references to the
 353     /// key and value in between.
 354     ///
 355     /// # Safety
 356     /// There must be another KV in the direction travelled.
 357     pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
 358         replace(self, |leaf_edge| {
 359             let kv = leaf_edge.next_back_kv();
 360             let kv = unsafe { unwrap_unchecked(kv.ok()) };
 361             (kv.next_back_leaf_edge(), kv.into_kv())
 362         })
 363     }
 364 }
 365
 366 impl<'a, K, V> Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge> {
 367     /// Moves the leaf edge handle to the next leaf edge and returns references to the
 368     /// key and value in between.
 369     ///
 370     /// # Safety
 371     /// There must be another KV in the direction travelled.
 372     pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
 373         let kv = replace(self, |leaf_edge| {
 374             let kv = leaf_edge.next_kv();
 375             let kv = unsafe { unwrap_unchecked(kv.ok()) };
 376             (unsafe { ptr::read(&kv) }.next_leaf_edge(), kv)
 377         });
 378         // Doing this last is faster, according to benchmarks.
 379         kv.into_kv_valmut()
 380     }
 381
 382     /// Moves the leaf edge handle to the previous leaf and returns references to the
 383     /// key and value in between.
 384     ///
 385     /// # Safety
 386     /// There must be another KV in the direction travelled.
 387     pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
 388         let kv = replace(self, |leaf_edge| {
 389             let kv = leaf_edge.next_back_kv();
 390             let kv = unsafe { unwrap_unchecked(kv.ok()) };
 391             (unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv)
 392         });
 393         // Doing this last is faster, according to benchmarks.
 394         kv.into_kv_valmut()
 395     }
 396 }
 397
 398 impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge> {
 399     /// Moves the leaf edge handle to the next leaf edge.
 400     ///
 401     /// # Safety
 402     /// There must be another KV in the direction travelled.
 403     pub unsafe fn move_next_unchecked(&mut self) {
 404         take_mut(self, |leaf_edge| {
 405             let kv = leaf_edge.next_kv();
 406             let kv = unsafe { unwrap_unchecked(kv.ok()) };
 407             kv.next_leaf_edge()
 408         })
 409     }
 410 }
 411
 412 impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
 413     /// Moves the leaf edge handle to the next leaf edge and returns the key and value
 414     /// in between, deallocating any node left behind while leaving the corresponding
 415     /// edge in its parent node dangling.
 416     ///
 417     /// # Safety
 418     /// - There must be another KV in the direction travelled.
 419     /// - That KV was not previously returned by counterpart `next_back_unchecked`
 420     ///   on any copy of the handles being used to traverse the tree.
 421     ///
 422     /// The only safe way to proceed with the updated handle is to compare it, drop it,
 423     /// call this method again subject to its safety conditions, or call counterpart
 424     /// `next_back_unchecked` subject to its safety conditions.
 425     pub unsafe fn next_unchecked(&mut self) -> (K, V) {
 426         replace(self, |leaf_edge| {
 427             let kv = unsafe { next_kv_unchecked_dealloc(leaf_edge) };
 428             let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
 429             let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
 430             (kv.next_leaf_edge(), (k, v))
 431         })
 432     }
 433
 434     /// Moves the leaf edge handle to the previous leaf edge and returns the key and value
 435     /// in between, deallocating any node left behind while leaving the corresponding
 436     /// edge in its parent node dangling.
 437     ///
 438     /// # Safety
 439     /// - There must be another KV in the direction travelled.
 440     /// - That leaf edge was not previously returned by counterpart `next_unchecked`
 441     ///   on any copy of the handles being used to traverse the tree.
 442     ///
 443     /// The only safe way to proceed with the updated handle is to compare it, drop it,
 444     /// call this method again subject to its safety conditions, or call counterpart
 445     /// `next_unchecked` subject to its safety conditions.
 446     pub unsafe fn next_back_unchecked(&mut self) -> (K, V) {
 447         replace(self, |leaf_edge| {
 448             let kv = unsafe { next_back_kv_unchecked_dealloc(leaf_edge) };
 449             let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
 450             let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
 451             (kv.next_back_leaf_edge(), (k, v))
 452         })
 453     }
 454 }
 455
 456 impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
 457     /// Returns the leftmost leaf edge in or underneath a node - in other words, the edge
 458     /// you need first when navigating forward (or last when navigating backward).
 459     #[inline]
 460     pub fn first_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
 461         let mut node = self;
 462         loop {
 463             match node.force() {
 464                 Leaf(leaf) => return leaf.first_edge(),
 465                 Internal(internal) => node = internal.first_edge().descend(),
 466             }
 467         }
 468     }
 469
 470     /// Returns the rightmost leaf edge in or underneath a node - in other words, the edge
 471     /// you need last when navigating forward (or first when navigating backward).
 472     #[inline]
 473     pub fn last_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
 474         let mut node = self;
 475         loop {
 476             match node.force() {
 477                 Leaf(leaf) => return leaf.last_edge(),
 478                 Internal(internal) => node = internal.last_edge().descend(),
 479             }
 480         }
 481     }
 482 }
 483
 484 pub enum Position<BorrowType, K, V> {
 485     Leaf(NodeRef<BorrowType, K, V, marker::Leaf>),
 486     Internal(NodeRef<BorrowType, K, V, marker::Internal>),
 487     InternalKV(Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV>),
 488 }
 489
 490 impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal> {
 491     /// Visits leaf nodes and internal KVs in order of ascending keys, and also
 492     /// visits internal nodes as a whole in a depth first order, meaning that
 493     /// internal nodes precede their individual KVs and their child nodes.
 494     pub fn visit_nodes_in_order<F>(self, mut visit: F)
 495     where
 496         F: FnMut(Position<marker::Immut<'a>, K, V>),
 497     {
 498         match self.force() {
 499             Leaf(leaf) => visit(Position::Leaf(leaf)),
 500             Internal(internal) => {
 501                 visit(Position::Internal(internal));
 502                 let mut edge = internal.first_edge();
 503                 loop {
 504                     edge = match edge.descend().force() {
 505                         Leaf(leaf) => {
 506                             visit(Position::Leaf(leaf));
 507                             match edge.next_kv() {
 508                                 Ok(kv) => {
 509                                     visit(Position::InternalKV(kv));
 510                                     kv.right_edge()
 511                                 }
 512                                 Err(_) => return,
 513                             }
 514                         }
 515                         Internal(internal) => {
 516                             visit(Position::Internal(internal));
 517                             internal.first_edge()
 518                         }
 519                     }
 520                 }
 521             }
 522         }
 523     }
 524
 525     /// Calculates the number of elements in a (sub)tree.
 526     pub fn calc_length(self) -> usize {
 527         let mut result = 0;
 528         self.visit_nodes_in_order(|pos| match pos {
 529             Position::Leaf(node) => result += node.len(),
 530             Position::Internal(node) => result += node.len(),
 531             Position::InternalKV(_) => (),
 532         });
 533         result
 534     }
 535 }
 536
 537 impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
 538     /// Returns the leaf edge closest to a KV for forward navigation.
 539     pub fn next_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
 540         match self.force() {
 541             Leaf(leaf_kv) => leaf_kv.right_edge(),
 542             Internal(internal_kv) => {
 543                 let next_internal_edge = internal_kv.right_edge();
 544                 next_internal_edge.descend().first_leaf_edge()
 545             }
 546         }
 547     }
 548
 549     /// Returns the leaf edge closest to a KV for backward navigation.
 550     pub fn next_back_leaf_edge(
 551         self,
 552     ) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
 553         match self.force() {
 554             Leaf(leaf_kv) => leaf_kv.left_edge(),
 555             Internal(internal_kv) => {
 556                 let next_internal_edge = internal_kv.left_edge();
 557                 next_internal_edge.descend().last_leaf_edge()
 558             }
 559         }
 560     }
 561 }