src/libcollections/btree/node.rs

   1 // Copyright 2014 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 // This is an attempt at an implementation following the ideal
  12 //
  13 // ```
  14 // struct BTreeMap<K, V> {
  15 //     height: usize,
  16 //     root: Option<Box<Node<K, V, height>>>
  17 // }
  18 //
  19 // struct Node<K, V, height: usize> {
  20 //     keys: [K; 2 * B - 1],
  21 //     vals: [V; 2 * B - 1],
  22 //     edges: if height > 0 {
  23 //         [Box<Node<K, V, height - 1>>; 2 * B]
  24 //     } else { () },
  25 //     parent: *const Node<K, V, height + 1>,
  26 //     parent_idx: u16,
  27 //     len: u16,
  28 // }
  29 // ```
  30 //
  31 // Since Rust doesn't acutally have dependent types and polymorphic recursion,
  32 // we make do with lots of unsafety.
  33
  34 use alloc::heap;
  35 use core::marker::PhantomData;
  36 use core::mem;
  37 use core::nonzero::NonZero;
  38 use core::ptr::{self, Unique};
  39 use core::slice;
  40
  41 use boxed::Box;
  42
  43 const B: usize = 6;
  44 pub const CAPACITY: usize = 2 * B - 1;
  45
  46 struct LeafNode<K, V> {
  47     keys: [K; CAPACITY],
  48     vals: [V; CAPACITY],
  49     parent: *const InternalNode<K, V>,
  50     parent_idx: u16,
  51     len: u16,
  52 }
  53
  54 impl<K, V> LeafNode<K, V> {
  55     unsafe fn new() -> Self {
  56         LeafNode {
  57             keys: mem::uninitialized(),
  58             vals: mem::uninitialized(),
  59             parent: ptr::null(),
  60             parent_idx: mem::uninitialized(),
  61             len: 0
  62         }
  63     }
  64 }
  65
  66 // We use repr(C) so that a pointer to an internal node can be
  67 // directly used as a pointer to a leaf node
  68 #[repr(C)]
  69 struct InternalNode<K, V> {
  70     data: LeafNode<K, V>,
  71     edges: [BoxedNode<K, V>; 2 * B],
  72 }
  73
  74 impl<K, V> InternalNode<K, V> {
  75     unsafe fn new() -> Self {
  76         InternalNode {
  77             data: LeafNode::new(),
  78             edges: mem::uninitialized()
  79         }
  80     }
  81 }
  82
  83 struct BoxedNode<K, V> {
  84     ptr: Unique<LeafNode<K, V>> // we don't know if this points to a leaf node or an internal node
  85 }
  86
  87 impl<K, V> BoxedNode<K, V> {
  88     fn from_leaf(node: Box<LeafNode<K, V>>) -> Self {
  89         unsafe {
  90             BoxedNode { ptr: Unique::new(Box::into_raw(node)) }
  91         }
  92     }
  93
  94     fn from_internal(node: Box<InternalNode<K, V>>) -> Self {
  95         unsafe {
  96             BoxedNode { ptr: Unique::new(Box::into_raw(node) as *mut LeafNode<K, V>) }
  97         }
  98     }
  99
 100     unsafe fn from_ptr(ptr: NonZero<*mut LeafNode<K, V>>) -> Self {
 101         BoxedNode { ptr: Unique::new(*ptr) }
 102     }
 103
 104     fn as_ptr(&self) -> NonZero<*mut LeafNode<K, V>> {
 105         unsafe {
 106             NonZero::new(*self.ptr)
 107         }
 108     }
 109 }
 110
 111 /// An owned tree. Note that despite being owned, this does not have a destructor,
 112 /// and must be cleaned up manually.
 113 pub struct Root<K, V> {
 114     node: BoxedNode<K, V>,
 115     height: usize
 116 }
 117
 118 unsafe impl<K: Sync, V: Sync> Sync for Root<K, V> { }
 119 unsafe impl<K: Send, V: Send> Send for Root<K, V> { }
 120
 121 impl<K, V> Root<K, V> {
 122     pub fn new_leaf() -> Self {
 123         Root {
 124             node: BoxedNode::from_leaf(Box::new(unsafe { LeafNode::new() })),
 125             height: 0
 126         }
 127     }
 128
 129     pub fn as_ref(&self)
 130             -> NodeRef<marker::Immut, K, V, marker::LeafOrInternal> {
 131         NodeRef {
 132             height: self.height,
 133             node: self.node.as_ptr(),
 134             root: self as *const _ as *mut _,
 135             _marker: PhantomData,
 136         }
 137     }
 138
 139     pub fn as_mut(&mut self)
 140             -> NodeRef<marker::Mut, K, V, marker::LeafOrInternal> {
 141         NodeRef {
 142             height: self.height,
 143             node: self.node.as_ptr(),
 144             root: self as *mut _,
 145             _marker: PhantomData,
 146         }
 147     }
 148
 149     pub fn into_ref(self)
 150             -> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
 151         NodeRef {
 152             height: self.height,
 153             node: self.node.as_ptr(),
 154             root: ptr::null_mut(), // FIXME: Is there anything better to do here?
 155             _marker: PhantomData,
 156         }
 157     }
 158
 159     /// Add a new internal node with a single edge, pointing to the previous root, and make that
 160     /// new node the root. This increases the height by 1 and is the opposite of `pop_level`.
 161     pub fn push_level(&mut self)
 162             -> NodeRef<marker::Mut, K, V, marker::Internal> {
 163         let mut new_node = Box::new(unsafe { InternalNode::new() });
 164         new_node.edges[0] = unsafe { BoxedNode::from_ptr(self.node.as_ptr()) };
 165
 166         self.node = BoxedNode::from_internal(new_node);
 167         self.height += 1;
 168
 169         let mut ret = NodeRef {
 170             height: self.height,
 171             node: self.node.as_ptr(),
 172             root: self as *mut _,
 173             _marker: PhantomData
 174         };
 175
 176         unsafe {
 177             ret.reborrow_mut().first_edge().correct_parent_link();
 178         }
 179
 180         ret
 181     }
 182
 183     /// Remove the root node, using its first child as the new root. This cannot be called when
 184     /// the tree consists only of a leaf node. As it is intended only to be called when the root
 185     /// has only one edge, no cleanup is done on any of the other children are elements of the root.
 186     /// This decreases the height by 1 and is the opposite of `push_level`.
 187     pub fn pop_level(&mut self) {
 188         debug_assert!(self.height > 0);
 189
 190         let top = *self.node.ptr as *mut u8;
 191
 192         self.node = unsafe {
 193             BoxedNode::from_ptr(self.as_mut()
 194                                     .cast_unchecked::<marker::Internal>()
 195                                     .first_edge()
 196                                     .descend()
 197                                     .node)
 198         };
 199         self.height -= 1;
 200         self.as_mut().as_leaf_mut().parent = ptr::null();
 201
 202         unsafe {
 203             heap::deallocate(
 204                 top,
 205                 mem::size_of::<InternalNode<K, V>>(),
 206                 mem::align_of::<InternalNode<K, V>>()
 207             );
 208         }
 209     }
 210 }
 211
 212 /// A reference to a node.
 213 ///
 214 /// This type has a number of paramaters that controls how it acts:
 215 /// - `BorrowType`: This can be `Immut<'a>` or `Mut<'a>` for some `'a` or `Owned`.
 216 ///    When this is `Immut<'a>`, the `NodeRef` acts roughly like `&'a Node`,
 217 ///    when this is `Mut<'a>`, the `NodeRef` acts roughly like `&'a mut Node`,
 218 ///    and when this is `Owned`, the `NodeRef` acts roughly like `Box<Node>`.
 219 /// - `K` and `V`: These control what types of things are stored in the nodes.
 220 /// - `Type`: This can be `Leaf`, `Internal`, or `LeafOrInternal`. When this is
 221 ///   `Leaf`, the `NodeRef` points to a leaf node, when this is `Internal` the
 222 ///   `NodeRef` points to an internal node, and when this is `LeafOrInternal` the
 223 ///   `NodeRef` could be pointing to either type of node.
 224 pub struct NodeRef<BorrowType, K, V, Type> {
 225     height: usize,
 226     node: NonZero<*mut LeafNode<K, V>>,
 227     root: *mut Root<K, V>,
 228     _marker: PhantomData<(BorrowType, Type)>
 229 }
 230
 231 impl<'a, K: 'a, V: 'a, Type> Copy for NodeRef<marker::Immut<'a>, K, V, Type> { }
 232 impl<'a, K: 'a, V: 'a, Type> Clone for NodeRef<marker::Immut<'a>, K, V, Type> {
 233     fn clone(&self) -> Self {
 234         *self
 235     }
 236 }
 237
 238 unsafe impl<BorrowType, K: Sync, V: Sync, Type> Sync
 239     for NodeRef<BorrowType, K, V, Type> { }
 240
 241 unsafe impl<'a, K: Sync + 'a, V: Sync + 'a, Type> Send
 242    for NodeRef<marker::Immut<'a>, K, V, Type> { }
 243 unsafe impl<'a, K: Send + 'a, V: Send + 'a, Type> Send
 244    for NodeRef<marker::Mut<'a>, K, V, Type> { }
 245 unsafe impl<K: Send, V: Send, Type> Send
 246    for NodeRef<marker::Owned, K, V, Type> { }
 247
 248 impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Internal> {
 249     fn as_internal(&self) -> &InternalNode<K, V> {
 250         unsafe {
 251             &*(*self.node as *const InternalNode<K, V>)
 252         }
 253     }
 254 }
 255
 256 impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
 257     fn as_internal_mut(&mut self) -> &mut InternalNode<K, V> {
 258         unsafe {
 259             &mut *(*self.node as *mut InternalNode<K, V>)
 260         }
 261     }
 262 }
 263
 264
 265 impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
 266     pub fn len(&self) -> usize {
 267         self.as_leaf().len as usize
 268     }
 269
 270     pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
 271         NodeRef {
 272             height: self.height,
 273             node: self.node,
 274             root: self.root,
 275             _marker: PhantomData
 276         }
 277     }
 278
 279     fn reborrow<'a>(&'a self) -> NodeRef<marker::Immut<'a>, K, V, Type> {
 280         NodeRef {
 281             height: self.height,
 282             node: self.node,
 283             root: self.root,
 284             _marker: PhantomData
 285         }
 286     }
 287
 288     fn as_leaf(&self) -> &LeafNode<K, V> {
 289         unsafe {
 290             &**self.node
 291         }
 292     }
 293
 294     pub fn keys(&self) -> &[K] {
 295         self.reborrow().into_slices().0
 296     }
 297
 298     pub fn vals(&self) -> &[V] {
 299         self.reborrow().into_slices().1
 300     }
 301
 302     pub fn ascend(self) -> Result<
 303         Handle<
 304             NodeRef<
 305                 BorrowType,
 306                 K, V,
 307                 marker::Internal
 308             >,
 309             marker::Edge
 310         >,
 311         Self
 312     > {
 313         if self.as_leaf().parent.is_null() {
 314             Err(self)
 315         } else {
 316             Ok(Handle {
 317                 node: NodeRef {
 318                     height: self.height + 1,
 319                     node: unsafe {
 320                         NonZero::new(self.as_leaf().parent as *mut LeafNode<K, V>)
 321                     },
 322                     root: self.root,
 323                     _marker: PhantomData
 324                 },
 325                 idx: self.as_leaf().parent_idx as usize,
 326                 _marker: PhantomData
 327             })
 328         }
 329     }
 330
 331     pub fn first_edge(self) -> Handle<Self, marker::Edge> {
 332         Handle::new_edge(self, 0)
 333     }
 334
 335     pub fn last_edge(self) -> Handle<Self, marker::Edge> {
 336         let len = self.len();
 337         Handle::new_edge(self, len)
 338     }
 339 }
 340
 341 impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
 342     pub unsafe fn deallocate_and_ascend(self) -> Option<
 343         Handle<
 344             NodeRef<
 345                 marker::Owned,
 346                 K, V,
 347                 marker::Internal
 348             >,
 349             marker::Edge
 350         >
 351     > {
 352         let ptr = self.as_leaf() as *const LeafNode<K, V> as *const u8 as *mut u8;
 353         let ret = self.ascend().ok();
 354         heap::deallocate(ptr, mem::size_of::<LeafNode<K, V>>(), mem::align_of::<LeafNode<K, V>>());
 355         ret
 356     }
 357 }
 358
 359 impl<K, V> NodeRef<marker::Owned, K, V, marker::Internal> {
 360     pub unsafe fn deallocate_and_ascend(self) -> Option<
 361         Handle<
 362             NodeRef<
 363                 marker::Owned,
 364                 K, V,
 365                 marker::Internal
 366             >,
 367             marker::Edge
 368         >
 369     > {
 370         let ptr = self.as_internal() as *const InternalNode<K, V> as *const u8 as *mut u8;
 371         let ret = self.ascend().ok();
 372         heap::deallocate(
 373             ptr,
 374             mem::size_of::<InternalNode<K, V>>(),
 375             mem::align_of::<InternalNode<K, V>>()
 376         );
 377         ret
 378     }
 379 }
 380
 381 impl<'a, K, V, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
 382     unsafe fn cast_unchecked<NewType>(&mut self)
 383             -> NodeRef<marker::Mut, K, V, NewType> {
 384
 385         NodeRef {
 386             height: self.height,
 387             node: self.node,
 388             root: self.root,
 389             _marker: PhantomData
 390         }
 391     }
 392
 393     unsafe fn reborrow_mut(&mut self) -> NodeRef<marker::Mut, K, V, Type> {
 394         NodeRef {
 395             height: self.height,
 396             node: self.node,
 397             root: self.root,
 398             _marker: PhantomData
 399         }
 400     }
 401
 402     fn as_leaf_mut(&mut self) -> &mut LeafNode<K, V> {
 403         unsafe {
 404             &mut **self.node
 405         }
 406     }
 407
 408     pub fn keys_mut(&mut self) -> &mut [K] {
 409         unsafe { self.reborrow_mut().into_slices_mut().0 }
 410     }
 411
 412     pub fn vals_mut(&mut self) -> &mut [V] {
 413         unsafe { self.reborrow_mut().into_slices_mut().1 }
 414     }
 415 }
 416
 417 impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Immut<'a>, K, V, Type> {
 418     pub fn into_slices(self) -> (&'a [K], &'a [V]) {
 419         unsafe {
 420             (
 421                 slice::from_raw_parts(
 422                     self.as_leaf().keys.as_ptr(),
 423                     self.len()
 424                 ),
 425                 slice::from_raw_parts(
 426                     self.as_leaf().vals.as_ptr(),
 427                     self.len()
 428                 )
 429             )
 430         }
 431     }
 432 }
 433
 434 impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
 435     pub fn into_root_mut(self) -> &'a mut Root<K, V> {
 436         unsafe {
 437             &mut *self.root
 438         }
 439     }
 440
 441     pub fn into_slices_mut(mut self) -> (&'a mut [K], &'a mut [V]) {
 442         unsafe {
 443             (
 444                 slice::from_raw_parts_mut(
 445                     &mut self.as_leaf_mut().keys as *mut [K] as *mut K,
 446                     self.len()
 447                 ),
 448                 slice::from_raw_parts_mut(
 449                     &mut self.as_leaf_mut().vals as *mut [V] as *mut V,
 450                     self.len()
 451                 )
 452             )
 453         }
 454     }
 455 }
 456
 457 impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::Leaf> {
 458     pub fn push(&mut self, key: K, val: V) {
 459         // Necessary for correctness, but this is an internal module
 460         debug_assert!(self.len() < CAPACITY);
 461
 462         let idx = self.len();
 463
 464         unsafe {
 465             ptr::write(self.keys_mut().get_unchecked_mut(idx), key);
 466             ptr::write(self.vals_mut().get_unchecked_mut(idx), val);
 467         }
 468
 469         self.as_leaf_mut().len += 1;
 470     }
 471
 472     pub fn push_front(&mut self, key: K, val: V) {
 473         // Necessary for correctness, but this is an internal module
 474         debug_assert!(self.len() < CAPACITY);
 475
 476         unsafe {
 477             slice_insert(self.keys_mut(), 0, key);
 478             slice_insert(self.vals_mut(), 0, val);
 479         }
 480
 481         self.as_leaf_mut().len += 1;
 482     }
 483 }
 484
 485 impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
 486     pub fn push(&mut self, key: K, val: V, edge: Root<K, V>) {
 487         // Necessary for correctness, but this is an internal module
 488         debug_assert!(edge.height == self.height - 1);
 489         debug_assert!(self.len() < CAPACITY);
 490
 491         let idx = self.len();
 492
 493         unsafe {
 494             ptr::write(self.keys_mut().get_unchecked_mut(idx), key);
 495             ptr::write(self.vals_mut().get_unchecked_mut(idx), val);
 496             ptr::write(self.as_internal_mut().edges.get_unchecked_mut(idx + 1), edge.node);
 497
 498             self.as_leaf_mut().len += 1;
 499
 500             Handle::new_edge(self.reborrow_mut(), idx + 1).correct_parent_link();
 501         }
 502     }
 503
 504     pub fn push_front(&mut self, key: K, val: V, edge: Root<K, V>) {
 505         // Necessary for correctness, but this is an internal module
 506         debug_assert!(edge.height == self.height - 1);
 507         debug_assert!(self.len() < CAPACITY);
 508
 509         unsafe {
 510             slice_insert(self.keys_mut(), 0, key);
 511             slice_insert(self.vals_mut(), 0, val);
 512             slice_insert(
 513                 slice::from_raw_parts_mut(
 514                     self.as_internal_mut().edges.as_mut_ptr(),
 515                     self.len()+1
 516                 ),
 517                 0,
 518                 edge.node
 519             );
 520
 521             self.as_leaf_mut().len += 1;
 522
 523             for i in 0..self.len()+1 {
 524                 Handle::new_edge(self.reborrow_mut(), i).correct_parent_link();
 525             }
 526         }
 527
 528     }
 529 }
 530
 531 impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
 532     pub fn pop(&mut self) -> (K, V, Option<Root<K, V>>) {
 533         // Necessary for correctness, but this is an internal module
 534         debug_assert!(self.len() > 0);
 535
 536         let idx = self.len() - 1;
 537
 538         unsafe {
 539             let key = ptr::read(self.keys().get_unchecked(idx));
 540             let val = ptr::read(self.vals().get_unchecked(idx));
 541             let edge = match self.reborrow_mut().force() {
 542                 ForceResult::Leaf(_) => None,
 543                 ForceResult::Internal(internal) => {
 544                     let edge = ptr::read(internal.as_internal().edges.get_unchecked(idx + 1));
 545                     let mut new_root = Root { node: edge, height: internal.height - 1 };
 546                     new_root.as_mut().as_leaf_mut().parent = ptr::null();
 547                     Some(new_root)
 548                 }
 549             };
 550
 551             self.as_leaf_mut().len -= 1;
 552             (key, val, edge)
 553         }
 554     }
 555
 556     pub fn pop_front(&mut self) -> (K, V, Option<Root<K, V>>) {
 557         // Necessary for correctness, but this is an internal module
 558         debug_assert!(self.len() > 0);
 559
 560         let old_len = self.len();
 561
 562         unsafe {
 563             let key = slice_remove(self.keys_mut(), 0);
 564             let val = slice_remove(self.vals_mut(), 0);
 565             let edge = match self.reborrow_mut().force() {
 566                 ForceResult::Leaf(_) => None,
 567                 ForceResult::Internal(mut internal) => {
 568                     let edge = slice_remove(
 569                         slice::from_raw_parts_mut(
 570                             internal.as_internal_mut().edges.as_mut_ptr(),
 571                             old_len+1
 572                         ),
 573                         0
 574                     );
 575
 576                     let mut new_root = Root { node: edge, height: internal.height - 1 };
 577                     new_root.as_mut().as_leaf_mut().parent = ptr::null();
 578
 579                     for i in 0..old_len {
 580                         Handle::new_edge(internal.reborrow_mut(), i).correct_parent_link();
 581                     }
 582
 583                     Some(new_root)
 584                 }
 585             };
 586
 587             self.as_leaf_mut().len -= 1;
 588
 589             (key, val, edge)
 590         }
 591     }
 592 }
 593
 594 impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
 595     pub fn force(self) -> ForceResult<
 596         NodeRef<BorrowType, K, V, marker::Leaf>,
 597         NodeRef<BorrowType, K, V, marker::Internal>
 598     > {
 599         if self.height == 0 {
 600             ForceResult::Leaf(NodeRef {
 601                 height: self.height,
 602                 node: self.node,
 603                 root: self.root,
 604                 _marker: PhantomData
 605             })
 606         } else {
 607             ForceResult::Internal(NodeRef {
 608                 height: self.height,
 609                 node: self.node,
 610                 root: self.root,
 611                 _marker: PhantomData
 612             })
 613         }
 614     }
 615 }
 616
 617 pub struct Handle<Node, Type> {
 618     node: Node,
 619     idx: usize,
 620     _marker: PhantomData<Type>
 621 }
 622
 623 impl<Node: Copy, Type> Copy for Handle<Node, Type> { }
 624 impl<Node: Copy, Type> Clone for Handle<Node, Type> {
 625     fn clone(&self) -> Self {
 626         *self
 627     }
 628 }
 629
 630 impl<Node, Type> Handle<Node, Type> {
 631     pub fn into_node(self) -> Node {
 632         self.node
 633     }
 634 }
 635
 636 impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::KV> {
 637     pub fn new_kv(node: NodeRef<BorrowType, K, V, NodeType>, idx: usize) -> Self {
 638         // Necessary for correctness, but in a private module
 639         debug_assert!(idx < node.len());
 640
 641         Handle {
 642             node: node,
 643             idx: idx,
 644             _marker: PhantomData
 645         }
 646     }
 647
 648     pub fn left_edge(self) -> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::Edge> {
 649         Handle::new_edge(self.node, self.idx)
 650     }
 651
 652     pub fn right_edge(self) -> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::Edge> {
 653         Handle::new_edge(self.node, self.idx + 1)
 654     }
 655 }
 656
 657 impl<BorrowType, K, V, NodeType, HandleType> PartialEq
 658         for Handle<NodeRef<BorrowType, K, V, NodeType>, HandleType> {
 659
 660     fn eq(&self, other: &Self) -> bool {
 661         self.node.node == other.node.node && self.idx == other.idx
 662     }
 663 }
 664
 665 impl<BorrowType, K, V, NodeType, HandleType>
 666         Handle<NodeRef<BorrowType, K, V, NodeType>, HandleType> {
 667
 668     pub fn reborrow(&self)
 669             -> Handle<NodeRef<marker::Immut, K, V, NodeType>, HandleType> {
 670
 671         // We can't use Handle::new_kv or Handle::new_edge because we don't know our type
 672         Handle {
 673             node: self.node.reborrow(),
 674             idx: self.idx,
 675             _marker: PhantomData
 676         }
 677     }
 678 }
 679
 680 impl<'a, K, V, NodeType, HandleType>
 681         Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, HandleType> {
 682
 683     pub unsafe fn reborrow_mut(&mut self)
 684             -> Handle<NodeRef<marker::Mut, K, V, NodeType>, HandleType> {
 685
 686         // We can't use Handle::new_kv or Handle::new_edge because we don't know our type
 687         Handle {
 688             node: self.node.reborrow_mut(),
 689             idx: self.idx,
 690             _marker: PhantomData
 691         }
 692     }
 693 }
 694
 695 impl<BorrowType, K, V, NodeType>
 696         Handle<NodeRef<BorrowType, K, V, NodeType>, marker::Edge> {
 697
 698     pub fn new_edge(node: NodeRef<BorrowType, K, V, NodeType>, idx: usize) -> Self {
 699         // Necessary for correctness, but in a private module
 700         debug_assert!(idx <= node.len());
 701
 702         Handle {
 703             node: node,
 704             idx: idx,
 705             _marker: PhantomData
 706         }
 707     }
 708
 709     pub fn left_kv(self)
 710             -> Result<Handle<NodeRef<BorrowType, K, V, NodeType>, marker::KV>, Self> {
 711
 712         if self.idx > 0 {
 713             Ok(Handle::new_kv(self.node, self.idx - 1))
 714         } else {
 715             Err(self)
 716         }
 717     }
 718
 719     pub fn right_kv(self)
 720             -> Result<Handle<NodeRef<BorrowType, K, V, NodeType>, marker::KV>, Self> {
 721
 722         if self.idx < self.node.len() {
 723             Ok(Handle::new_kv(self.node, self.idx))
 724         } else {
 725             Err(self)
 726         }
 727     }
 728 }
 729
 730 impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge> {
 731     fn insert_fit(&mut self, key: K, val: V) -> *mut V {
 732         // Necessary for correctness, but in a private module
 733         debug_assert!(self.node.len() < CAPACITY);
 734
 735         unsafe {
 736             slice_insert(self.node.keys_mut(), self.idx, key);
 737             slice_insert(self.node.vals_mut(), self.idx, val);
 738
 739             self.node.as_leaf_mut().len += 1;
 740
 741             self.node.vals_mut().get_unchecked_mut(self.idx)
 742         }
 743     }
 744
 745     pub fn insert(mut self, key: K, val: V)
 746             -> (InsertResult<'a, K, V, marker::Leaf>, *mut V) {
 747
 748         if self.node.len() < CAPACITY {
 749             let ptr = self.insert_fit(key, val);
 750             (InsertResult::Fit(Handle::new_kv(self.node, self.idx)), ptr)
 751         } else {
 752             let middle = Handle::new_kv(self.node, B);
 753             let (mut left, k, v, mut right) = middle.split();
 754             let ptr = if self.idx <= B {
 755                 unsafe {
 756                     Handle::new_edge(left.reborrow_mut(), self.idx).insert_fit(key, val)
 757                 }
 758             } else {
 759                 unsafe {
 760                     Handle::new_edge(
 761                         right.as_mut().cast_unchecked::<marker::Leaf>(),
 762                         self.idx - (B + 1)
 763                     ).insert_fit(key, val)
 764                 }
 765             };
 766             (InsertResult::Split(left, k, v, right), ptr)
 767         }
 768     }
 769 }
 770
 771 impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge> {
 772     fn correct_parent_link(mut self) {
 773         let idx = self.idx as u16;
 774         let ptr = self.node.as_internal_mut() as *mut _;
 775         let mut child = self.descend();
 776         child.as_leaf_mut().parent = ptr;
 777         child.as_leaf_mut().parent_idx = idx;
 778     }
 779
 780     unsafe fn cast_unchecked<NewType>(&mut self)
 781             -> Handle<NodeRef<marker::Mut, K, V, NewType>, marker::Edge> {
 782
 783         Handle::new_edge(self.node.cast_unchecked(), self.idx)
 784     }
 785
 786     fn insert_fit(&mut self, key: K, val: V, edge: Root<K, V>) {
 787         // Necessary for correctness, but in an internal module
 788         debug_assert!(self.node.len() < CAPACITY);
 789         debug_assert!(edge.height == self.node.height - 1);
 790
 791         unsafe {
 792             self.cast_unchecked::<marker::Leaf>().insert_fit(key, val);
 793
 794             slice_insert(
 795                 slice::from_raw_parts_mut(
 796                     self.node.as_internal_mut().edges.as_mut_ptr(),
 797                     self.node.len()
 798                 ),
 799                 self.idx + 1,
 800                 edge.node
 801             );
 802
 803             for i in (self.idx+1)..(self.node.len()+1) {
 804                 Handle::new_edge(self.node.reborrow_mut(), i).correct_parent_link();
 805             }
 806         }
 807     }
 808
 809     pub fn insert(mut self, key: K, val: V, edge: Root<K, V>)
 810             -> InsertResult<'a, K, V, marker::Internal> {
 811
 812         // Necessary for correctness, but this is an internal module
 813         debug_assert!(edge.height == self.node.height - 1);
 814
 815         if self.node.len() < CAPACITY {
 816             self.insert_fit(key, val, edge);
 817             InsertResult::Fit(Handle::new_kv(self.node, self.idx))
 818         } else {
 819             let middle = Handle::new_kv(self.node, B);
 820             let (mut left, k, v, mut right) = middle.split();
 821             if self.idx <= B {
 822                 unsafe {
 823                     Handle::new_edge(left.reborrow_mut(), self.idx).insert_fit(key, val, edge);
 824                 }
 825             } else {
 826                 unsafe {
 827                     Handle::new_edge(
 828                         right.as_mut().cast_unchecked::<marker::Internal>(),
 829                         self.idx - (B + 1)
 830                     ).insert_fit(key, val, edge);
 831                 }
 832             }
 833             InsertResult::Split(left, k, v, right)
 834         }
 835     }
 836 }
 837
 838 impl<BorrowType, K, V>
 839         Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge> {
 840
 841     pub fn descend(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
 842         NodeRef {
 843             height: self.node.height - 1,
 844             node: unsafe { self.node.as_internal().edges.get_unchecked(self.idx).as_ptr() },
 845             root: self.node.root,
 846             _marker: PhantomData
 847         }
 848     }
 849 }
 850
 851 impl<'a, K: 'a, V: 'a, NodeType>
 852         Handle<NodeRef<marker::Immut<'a>, K, V, NodeType>, marker::KV> {
 853
 854     pub fn into_kv(self) -> (&'a K, &'a V) {
 855         let (keys, vals) = self.node.into_slices();
 856         unsafe {
 857             (keys.get_unchecked(self.idx), vals.get_unchecked(self.idx))
 858         }
 859     }
 860 }
 861
 862 impl<'a, K: 'a, V: 'a, NodeType>
 863         Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, marker::KV> {
 864
 865     pub fn into_kv_mut(self) -> (&'a mut K, &'a mut V) {
 866         let (mut keys, mut vals) = self.node.into_slices_mut();
 867         unsafe {
 868             (keys.get_unchecked_mut(self.idx), vals.get_unchecked_mut(self.idx))
 869         }
 870     }
 871 }
 872
 873 impl<'a, K, V, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, marker::KV> {
 874     pub fn kv_mut(&mut self) -> (&mut K, &mut V) {
 875         unsafe {
 876             let (mut keys, mut vals) = self.node.reborrow_mut().into_slices_mut();
 877             (keys.get_unchecked_mut(self.idx), vals.get_unchecked_mut(self.idx))
 878         }
 879     }
 880 }
 881
 882 impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV> {
 883     pub fn split(mut self)
 884             -> (NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, K, V, Root<K, V>) {
 885         unsafe {
 886             let mut new_node = Box::new(LeafNode::new());
 887
 888             let k = ptr::read(self.node.keys().get_unchecked(self.idx));
 889             let v = ptr::read(self.node.vals().get_unchecked(self.idx));
 890
 891             let new_len = self.node.len() - self.idx - 1;
 892
 893             ptr::copy_nonoverlapping(
 894                 self.node.keys().as_ptr().offset(self.idx as isize + 1),
 895                 new_node.keys.as_mut_ptr(),
 896                 new_len
 897             );
 898             ptr::copy_nonoverlapping(
 899                 self.node.vals().as_ptr().offset(self.idx as isize + 1),
 900                 new_node.vals.as_mut_ptr(),
 901                 new_len
 902             );
 903
 904             self.node.as_leaf_mut().len = self.idx as u16;
 905             new_node.len = new_len as u16;
 906
 907             (
 908                 self.node,
 909                 k, v,
 910                 Root {
 911                     node: BoxedNode::from_leaf(new_node),
 912                     height: 0
 913                 }
 914             )
 915         }
 916     }
 917
 918     pub fn remove(mut self)
 919             -> (Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>, K, V) {
 920         unsafe {
 921             let k = slice_remove(self.node.keys_mut(), self.idx);
 922             let v = slice_remove(self.node.vals_mut(), self.idx);
 923             self.node.as_leaf_mut().len -= 1;
 924             (self.left_edge(), k, v)
 925         }
 926     }
 927 }
 928
 929 impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::KV> {
 930     pub fn split(mut self)
 931             -> (NodeRef<marker::Mut<'a>, K, V, marker::Internal>, K, V, Root<K, V>) {
 932         unsafe {
 933             let mut new_node = Box::new(InternalNode::new());
 934
 935             let k = ptr::read(self.node.keys().get_unchecked(self.idx));
 936             let v = ptr::read(self.node.vals().get_unchecked(self.idx));
 937
 938             let height = self.node.height;
 939             let new_len = self.node.len() - self.idx - 1;
 940
 941             ptr::copy_nonoverlapping(
 942                 self.node.keys().as_ptr().offset(self.idx as isize + 1),
 943                 new_node.data.keys.as_mut_ptr(),
 944                 new_len
 945             );
 946             ptr::copy_nonoverlapping(
 947                 self.node.vals().as_ptr().offset(self.idx as isize + 1),
 948                 new_node.data.vals.as_mut_ptr(),
 949                 new_len
 950             );
 951             ptr::copy_nonoverlapping(
 952                 self.node.as_internal().edges.as_ptr().offset(self.idx as isize + 1),
 953                 new_node.edges.as_mut_ptr(),
 954                 new_len + 1
 955             );
 956
 957             self.node.as_leaf_mut().len = self.idx as u16;
 958             new_node.data.len = new_len as u16;
 959
 960             let mut new_root = Root {
 961                 node: BoxedNode::from_internal(new_node),
 962                 height: height
 963             };
 964
 965             for i in 0..(new_len+1) {
 966                 Handle::new_edge(new_root.as_mut().cast_unchecked(), i).correct_parent_link();
 967             }
 968
 969             (
 970                 self.node,
 971                 k, v,
 972                 new_root
 973             )
 974         }
 975     }
 976
 977     pub fn can_merge(&self) -> bool {
 978         (
 979             self.reborrow()
 980                 .left_edge()
 981                 .descend()
 982                 .len()
 983           + self.reborrow()
 984                 .right_edge()
 985                 .descend()
 986                 .len()
 987           + 1
 988         ) <= CAPACITY
 989     }
 990
 991     pub fn merge(mut self)
 992             -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge> {
 993         let self1 = unsafe { ptr::read(&self) };
 994         let self2 = unsafe { ptr::read(&self) };
 995         let mut left_node = self1.left_edge().descend();
 996         let left_len = left_node.len();
 997         let mut right_node = self2.right_edge().descend();
 998         let right_len = right_node.len();
 999
1000         // necessary for correctness, but in a private module
1001         debug_assert!(left_len + right_len + 1 <= CAPACITY);
1002
1003         unsafe {
1004             ptr::write(left_node.keys_mut().get_unchecked_mut(left_len),
1005                        slice_remove(self.node.keys_mut(), self.idx));
1006             ptr::copy_nonoverlapping(
1007                 right_node.keys().as_ptr(),
1008                 left_node.keys_mut().as_mut_ptr().offset(left_len as isize + 1),
1009                 right_len
1010             );
1011             ptr::write(left_node.vals_mut().get_unchecked_mut(left_len),
1012                        slice_remove(self.node.vals_mut(), self.idx));
1013             ptr::copy_nonoverlapping(
1014                 right_node.vals().as_ptr(),
1015                 left_node.vals_mut().as_mut_ptr().offset(left_len as isize + 1),
1016                 right_len
1017             );
1018
1019             slice_remove(&mut self.node.as_internal_mut().edges, self.idx + 1);
1020             for i in self.idx+1..self.node.len() {
1021                 Handle::new_edge(self.node.reborrow_mut(), i).correct_parent_link();
1022             }
1023             self.node.as_leaf_mut().len -= 1;
1024
1025             if self.node.height > 1 {
1026                 ptr::copy_nonoverlapping(
1027                     right_node.cast_unchecked().as_internal().edges.as_ptr(),
1028                     left_node.cast_unchecked()
1029                              .as_internal_mut()
1030                              .edges
1031                              .as_mut_ptr()
1032                              .offset(left_len as isize + 1),
1033                     right_len + 1
1034                 );
1035
1036                 for i in left_len+1..left_len+right_len+2 {
1037                     Handle::new_edge(
1038                         left_node.cast_unchecked().reborrow_mut(),
1039                         i
1040                     ).correct_parent_link();
1041                 }
1042
1043                 heap::deallocate(
1044                     *right_node.node as *mut u8,
1045                     mem::size_of::<InternalNode<K, V>>(),
1046                     mem::align_of::<InternalNode<K, V>>()
1047                 );
1048             } else {
1049                 heap::deallocate(
1050                     *right_node.node as *mut u8,
1051                     mem::size_of::<LeafNode<K, V>>(),
1052                     mem::align_of::<LeafNode<K, V>>()
1053                 );
1054             }
1055
1056             left_node.as_leaf_mut().len += right_len as u16 + 1;
1057
1058             Handle::new_edge(self.node, self.idx)
1059         }
1060     }
1061 }
1062
1063 impl<BorrowType, K, V, HandleType>
1064         Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, HandleType> {
1065
1066     pub fn force(self) -> ForceResult<
1067         Handle<NodeRef<BorrowType, K, V, marker::Leaf>, HandleType>,
1068         Handle<NodeRef<BorrowType, K, V, marker::Internal>, HandleType>
1069     > {
1070         match self.node.force() {
1071             ForceResult::Leaf(node) => ForceResult::Leaf(Handle {
1072                 node: node,
1073                 idx: self.idx,
1074                 _marker: PhantomData
1075             }),
1076             ForceResult::Internal(node) => ForceResult::Internal(Handle {
1077                 node: node,
1078                 idx: self.idx,
1079                 _marker: PhantomData
1080             })
1081         }
1082     }
1083 }
1084
1085 pub enum ForceResult<Leaf, Internal> {
1086     Leaf(Leaf),
1087     Internal(Internal)
1088 }
1089
1090 pub enum InsertResult<'a, K, V, Type> {
1091     Fit(Handle<NodeRef<marker::Mut<'a>, K, V, Type>, marker::KV>),
1092     Split(NodeRef<marker::Mut<'a>, K, V, Type>, K, V, Root<K, V>)
1093 }
1094
1095 pub mod marker {
1096     use core::marker::PhantomData;
1097
1098     pub enum Leaf { }
1099     pub enum Internal { }
1100     pub enum LeafOrInternal { }
1101
1102     pub enum Owned { }
1103     pub struct Immut<'a>(PhantomData<&'a ()>);
1104     pub struct Mut<'a>(PhantomData<&'a mut ()>);
1105
1106     pub enum KV { }
1107     pub enum Edge { }
1108 }
1109
1110 unsafe fn slice_insert<T>(slice: &mut [T], idx: usize, val: T) {
1111     ptr::copy(
1112         slice.as_ptr().offset(idx as isize),
1113         slice.as_mut_ptr().offset(idx as isize + 1),
1114         slice.len() - idx
1115     );
1116     ptr::write(slice.get_unchecked_mut(idx), val);
1117 }
1118
1119 unsafe fn slice_remove<T>(slice: &mut [T], idx: usize) -> T {
1120     let ret = ptr::read(slice.get_unchecked(idx));
1121     ptr::copy(
1122         slice.as_ptr().offset(idx as isize + 1),
1123         slice.as_mut_ptr().offset(idx as isize),
1124         slice.len() - idx - 1
1125     );
1126     ret
1127 }