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 5 use std::mem::replace; 6 use std::ops::Range; 7 8 use crate::nodes::chunk::{Chunk, CHUNK_SIZE}; 9 use crate::util::{ 10 Pool, PoolRef, 11 Side::{self, Left, Right}, 12 }; 13 use crate::vector::RRBPool; 14 15 use self::Entry::*; 16 17 pub(crate) const NODE_SIZE: usize = CHUNK_SIZE; 18 19 #[derive(Debug)] 20 enum Size { 21 Size(usize), 22 Table(PoolRef<Chunk<usize>>), 23 } 24 25 impl Clone for Size { clone(&self) -> Self26 fn clone(&self) -> Self { 27 match *self { 28 Size::Size(size) => Size::Size(size), 29 Size::Table(ref table) => Size::Table(table.clone()), 30 } 31 } 32 } 33 34 impl Size { size(&self) -> usize35 fn size(&self) -> usize { 36 match self { 37 Size::Size(s) => *s, 38 Size::Table(sizes) => *sizes.last().unwrap_or(&0), 39 } 40 } 41 is_size(&self) -> bool42 fn is_size(&self) -> bool { 43 match self { 44 Size::Size(_) => true, 45 Size::Table(_) => false, 46 } 47 } 48 table_from_size(pool: &Pool<Chunk<usize>>, level: usize, size: usize) -> Self49 fn table_from_size(pool: &Pool<Chunk<usize>>, level: usize, size: usize) -> Self { 50 let mut chunk = Chunk::new(); 51 let mut remaining = size; 52 if let Some(child_size) = NODE_SIZE.checked_pow(level as u32) { 53 while remaining > child_size { 54 let next_value = chunk.last().unwrap_or(&0) + child_size; 55 chunk.push_back(next_value); 56 remaining -= child_size; 57 } 58 } 59 if remaining > 0 { 60 let next_value = chunk.last().unwrap_or(&0) + remaining; 61 chunk.push_back(next_value); 62 } 63 Size::Table(PoolRef::new(pool, chunk)) 64 } 65 push(&mut self, pool: &Pool<Chunk<usize>>, side: Side, level: usize, value: usize)66 fn push(&mut self, pool: &Pool<Chunk<usize>>, side: Side, level: usize, value: usize) { 67 let size = match self { 68 Size::Size(ref mut size) => match side { 69 Left => *size, 70 Right => { 71 *size += value; 72 return; 73 } 74 }, 75 Size::Table(ref mut size_ref) => { 76 let size_table = PoolRef::make_mut(pool, size_ref); 77 debug_assert!(size_table.len() < NODE_SIZE); 78 match side { 79 Left => { 80 for entry in size_table.iter_mut() { 81 *entry += value; 82 } 83 size_table.push_front(value); 84 } 85 Right => { 86 let prev = *(size_table.last().unwrap_or(&0)); 87 size_table.push_back(value + prev); 88 } 89 } 90 return; 91 } 92 }; 93 *self = Size::table_from_size(pool, level, size); 94 self.push(pool, side, level, value); 95 } 96 pop(&mut self, pool: &Pool<Chunk<usize>>, side: Side, level: usize, value: usize)97 fn pop(&mut self, pool: &Pool<Chunk<usize>>, side: Side, level: usize, value: usize) { 98 let size = match self { 99 Size::Size(ref mut size) => match side { 100 Left => *size, 101 Right => { 102 *size -= value; 103 return; 104 } 105 }, 106 Size::Table(ref mut size_ref) => { 107 let size_table = PoolRef::make_mut(pool, size_ref); 108 match side { 109 Left => { 110 let first = size_table.pop_front(); 111 debug_assert_eq!(value, first); 112 for entry in size_table.iter_mut() { 113 *entry -= value; 114 } 115 } 116 Right => { 117 let pop = size_table.pop_back(); 118 let last = size_table.last().unwrap_or(&0); 119 debug_assert_eq!(value, pop - last); 120 } 121 } 122 return; 123 } 124 }; 125 *self = Size::table_from_size(pool, level, size); 126 self.pop(pool, side, level, value); 127 } 128 update(&mut self, pool: &Pool<Chunk<usize>>, index: usize, level: usize, value: isize)129 fn update(&mut self, pool: &Pool<Chunk<usize>>, index: usize, level: usize, value: isize) { 130 let size = match self { 131 Size::Size(ref size) => *size, 132 Size::Table(ref mut size_ref) => { 133 let size_table = PoolRef::make_mut(pool, size_ref); 134 for entry in size_table.iter_mut().skip(index) { 135 *entry = (*entry as isize + value) as usize; 136 } 137 return; 138 } 139 }; 140 *self = Size::table_from_size(pool, level, size); 141 self.update(pool, index, level, value); 142 } 143 } 144 145 pub(crate) enum PushResult<A> { 146 Full(A, usize), 147 Done, 148 } 149 150 pub(crate) enum PopResult<A> { 151 Done(A), 152 Drained(A), 153 Empty, 154 } 155 156 pub(crate) enum SplitResult { 157 Dropped(usize), 158 OutOfBounds, 159 } 160 161 // Invariants: Nodes only at level > 0, Values/Empty only at level = 0 162 enum Entry<A> { 163 Nodes(Size, PoolRef<Chunk<Node<A>>>), 164 Values(PoolRef<Chunk<A>>), 165 Empty, 166 } 167 168 impl<A: Clone> Clone for Entry<A> { clone(&self) -> Self169 fn clone(&self) -> Self { 170 match *self { 171 Nodes(ref size, ref nodes) => Nodes(size.clone(), nodes.clone()), 172 Values(ref values) => Values(values.clone()), 173 Empty => Empty, 174 } 175 } 176 } 177 178 impl<A: Clone> Entry<A> { len(&self) -> usize179 fn len(&self) -> usize { 180 match self { 181 Nodes(_, ref nodes) => nodes.len(), 182 Values(ref values) => values.len(), 183 Empty => 0, 184 } 185 } 186 is_full(&self) -> bool187 fn is_full(&self) -> bool { 188 match self { 189 Nodes(_, ref nodes) => nodes.is_full(), 190 Values(ref values) => values.is_full(), 191 Empty => false, 192 } 193 } 194 unwrap_values(&self) -> &Chunk<A>195 fn unwrap_values(&self) -> &Chunk<A> { 196 match self { 197 Values(ref values) => values, 198 _ => panic!("rrb::Entry::unwrap_values: expected values, found nodes"), 199 } 200 } 201 unwrap_nodes(&self) -> &Chunk<Node<A>>202 fn unwrap_nodes(&self) -> &Chunk<Node<A>> { 203 match self { 204 Nodes(_, ref nodes) => nodes, 205 _ => panic!("rrb::Entry::unwrap_nodes: expected nodes, found values"), 206 } 207 } 208 unwrap_values_mut(&mut self, pool: &RRBPool<A>) -> &mut Chunk<A>209 fn unwrap_values_mut(&mut self, pool: &RRBPool<A>) -> &mut Chunk<A> { 210 match self { 211 Values(ref mut values) => PoolRef::make_mut(&pool.value_pool, values), 212 _ => panic!("rrb::Entry::unwrap_values_mut: expected values, found nodes"), 213 } 214 } 215 unwrap_nodes_mut(&mut self, pool: &RRBPool<A>) -> &mut Chunk<Node<A>>216 fn unwrap_nodes_mut(&mut self, pool: &RRBPool<A>) -> &mut Chunk<Node<A>> { 217 match self { 218 Nodes(_, ref mut nodes) => PoolRef::make_mut(&pool.node_pool, nodes), 219 _ => panic!("rrb::Entry::unwrap_nodes_mut: expected nodes, found values"), 220 } 221 } 222 values(self) -> Chunk<A>223 fn values(self) -> Chunk<A> { 224 match self { 225 Values(values) => PoolRef::unwrap_or_clone(values), 226 _ => panic!("rrb::Entry::values: expected values, found nodes"), 227 } 228 } 229 nodes(self) -> Chunk<Node<A>>230 fn nodes(self) -> Chunk<Node<A>> { 231 match self { 232 Nodes(_, nodes) => PoolRef::unwrap_or_clone(nodes), 233 _ => panic!("rrb::Entry::nodes: expected nodes, found values"), 234 } 235 } 236 is_empty_node(&self) -> bool237 fn is_empty_node(&self) -> bool { 238 match self { 239 Empty => true, 240 _ => false, 241 } 242 } 243 } 244 245 // Node 246 247 pub(crate) struct Node<A> { 248 children: Entry<A>, 249 } 250 251 impl<A: Clone> Clone for Node<A> { clone(&self) -> Self252 fn clone(&self) -> Self { 253 Node { 254 children: self.children.clone(), 255 } 256 } 257 } 258 259 impl<A: Clone> Default for Node<A> { default() -> Self260 fn default() -> Self { 261 Self::new() 262 } 263 } 264 265 impl<A: Clone> Node<A> { new() -> Self266 pub(crate) fn new() -> Self { 267 Node { children: Empty } 268 } 269 parent(pool: &RRBPool<A>, level: usize, children: Chunk<Self>) -> Self270 pub(crate) fn parent(pool: &RRBPool<A>, level: usize, children: Chunk<Self>) -> Self { 271 let size = { 272 let mut size = Size::Size(0); 273 let mut it = children.iter().peekable(); 274 loop { 275 match it.next() { 276 None => break, 277 Some(child) => { 278 if size.is_size() 279 && !child.is_completely_dense(level - 1) 280 && it.peek().is_some() 281 { 282 size = Size::table_from_size(&pool.size_pool, level, size.size()); 283 } 284 size.push(&pool.size_pool, Right, level, child.len()) 285 } 286 } 287 } 288 size 289 }; 290 Node { 291 children: Nodes(size, PoolRef::new(&pool.node_pool, children)), 292 } 293 } 294 clear_node(&mut self)295 pub(crate) fn clear_node(&mut self) { 296 self.children = Empty; 297 } 298 from_chunk(pool: &RRBPool<A>, level: usize, chunk: PoolRef<Chunk<A>>) -> Self299 pub(crate) fn from_chunk(pool: &RRBPool<A>, level: usize, chunk: PoolRef<Chunk<A>>) -> Self { 300 let node = Node { 301 children: Values(chunk), 302 }; 303 node.elevate(pool, level) 304 } 305 single_parent(pool: &RRBPool<A>, node: Self) -> Self306 pub(crate) fn single_parent(pool: &RRBPool<A>, node: Self) -> Self { 307 let size = if node.is_dense() { 308 Size::Size(node.len()) 309 } else { 310 let size_table = Chunk::unit(node.len()); 311 Size::Table(PoolRef::new(&pool.size_pool, size_table)) 312 }; 313 let children = PoolRef::new(&pool.node_pool, Chunk::unit(node)); 314 Node { 315 children: Nodes(size, children), 316 } 317 } 318 join_dense(pool: &RRBPool<A>, left: Self, right: Self) -> Self319 pub(crate) fn join_dense(pool: &RRBPool<A>, left: Self, right: Self) -> Self { 320 let left_len = left.len(); 321 let right_len = right.len(); 322 Node { 323 children: { 324 let children = PoolRef::new(&pool.node_pool, Chunk::pair(left, right)); 325 Nodes(Size::Size(left_len + right_len), children) 326 }, 327 } 328 } 329 elevate(self, pool: &RRBPool<A>, level_increment: usize) -> Self330 pub(crate) fn elevate(self, pool: &RRBPool<A>, level_increment: usize) -> Self { 331 if level_increment > 0 { 332 Self::single_parent(pool, self.elevate(pool, level_increment - 1)) 333 } else { 334 self 335 } 336 } 337 join_branches(self, pool: &RRBPool<A>, right: Self, level: usize) -> Self338 pub(crate) fn join_branches(self, pool: &RRBPool<A>, right: Self, level: usize) -> Self { 339 let left_len = self.len(); 340 let right_len = right.len(); 341 let size = if self.is_completely_dense(level) && right.is_dense() { 342 Size::Size(left_len + right_len) 343 } else { 344 let size_table = Chunk::pair(left_len, left_len + right_len); 345 Size::Table(PoolRef::new(&pool.size_pool, size_table)) 346 }; 347 Node { 348 children: { 349 let children = Chunk::pair(self, right); 350 Nodes(size, PoolRef::new(&pool.node_pool, children)) 351 }, 352 } 353 } 354 len(&self) -> usize355 pub(crate) fn len(&self) -> usize { 356 match self.children { 357 Entry::Nodes(Size::Size(size), _) => size, 358 Entry::Nodes(Size::Table(ref size_table), _) => *(size_table.last().unwrap_or(&0)), 359 Entry::Values(ref values) => values.len(), 360 Entry::Empty => 0, 361 } 362 } 363 is_empty(&self) -> bool364 pub(crate) fn is_empty(&self) -> bool { 365 self.len() == 0 366 } 367 is_single(&self) -> bool368 pub(crate) fn is_single(&self) -> bool { 369 self.children.len() == 1 370 } 371 is_full(&self) -> bool372 pub(crate) fn is_full(&self) -> bool { 373 self.children.is_full() 374 } 375 376 #[allow(dead_code)] // this is only used by tests number_of_children(&self) -> usize377 pub(crate) fn number_of_children(&self) -> usize { 378 self.children.len() 379 } 380 first_child(&self) -> &Self381 pub(crate) fn first_child(&self) -> &Self { 382 self.children.unwrap_nodes().first().unwrap() 383 } 384 385 /// True if the node is dense and so doesn't have a size table is_dense(&self) -> bool386 fn is_dense(&self) -> bool { 387 match self.children { 388 Entry::Nodes(Size::Table(_), _) => false, 389 _ => true, 390 } 391 } 392 393 /// True if the node and its children are dense and at capacity 394 // TODO can use this technique to quickly test if a Size::Table 395 // should be converted back to a Size::Size is_completely_dense(&self, level: usize) -> bool396 fn is_completely_dense(&self, level: usize) -> bool { 397 // Size of a full node is NODE_SIZE at level 0, NODE_SIZE² at 398 // level 1, etc. 399 if let Some(expected_size) = NODE_SIZE.checked_pow(level as u32 + 1) { 400 self.size() == expected_size 401 } else { 402 // We overflowed a usize, there's no way we can be completely dense as we know the size 403 // fits in a usize. 404 false 405 } 406 } 407 408 #[inline] size(&self) -> usize409 fn size(&self) -> usize { 410 match self.children { 411 Entry::Nodes(ref size, _) => size.size(), 412 Entry::Values(ref values) => values.len(), 413 Entry::Empty => 0, 414 } 415 } 416 417 #[inline] push_size(&mut self, pool: &RRBPool<A>, side: Side, level: usize, value: usize)418 fn push_size(&mut self, pool: &RRBPool<A>, side: Side, level: usize, value: usize) { 419 if let Entry::Nodes(ref mut size, _) = self.children { 420 size.push(&pool.size_pool, side, level, value) 421 } 422 } 423 424 #[inline] pop_size(&mut self, pool: &RRBPool<A>, side: Side, level: usize, value: usize)425 fn pop_size(&mut self, pool: &RRBPool<A>, side: Side, level: usize, value: usize) { 426 if let Entry::Nodes(ref mut size, _) = self.children { 427 size.pop(&pool.size_pool, side, level, value) 428 } 429 } 430 431 #[inline] update_size(&mut self, pool: &RRBPool<A>, index: usize, level: usize, value: isize)432 fn update_size(&mut self, pool: &RRBPool<A>, index: usize, level: usize, value: isize) { 433 if let Entry::Nodes(ref mut size, _) = self.children { 434 size.update(&pool.size_pool, index, level, value) 435 } 436 } 437 size_up_to(&self, level: usize, index: usize) -> usize438 fn size_up_to(&self, level: usize, index: usize) -> usize { 439 if let Entry::Nodes(ref size, _) = self.children { 440 if index == 0 { 441 0 442 } else { 443 match size { 444 Size::Table(ref size_table) => size_table[index - 1], 445 Size::Size(_) => index * NODE_SIZE.pow(level as u32), 446 } 447 } 448 } else { 449 index 450 } 451 } 452 index_in(&self, level: usize, index: usize) -> Option<usize>453 fn index_in(&self, level: usize, index: usize) -> Option<usize> { 454 let mut target_idx = if let Some(child_size) = NODE_SIZE.checked_pow(level as u32) { 455 index / child_size 456 } else { 457 0 458 }; 459 if target_idx >= self.children.len() { 460 return None; 461 } 462 if let Entry::Nodes(Size::Table(ref size_table), _) = self.children { 463 while size_table[target_idx] <= index { 464 target_idx += 1; 465 if target_idx >= size_table.len() { 466 return None; 467 } 468 } 469 } 470 Some(target_idx) 471 } 472 index(&self, level: usize, index: usize) -> &A473 pub(crate) fn index(&self, level: usize, index: usize) -> &A { 474 if level == 0 { 475 &self.children.unwrap_values()[index] 476 } else { 477 let target_idx = self.index_in(level, index).unwrap(); 478 self.children.unwrap_nodes()[target_idx] 479 .index(level - 1, index - self.size_up_to(level, target_idx)) 480 } 481 } 482 index_mut(&mut self, pool: &RRBPool<A>, level: usize, index: usize) -> &mut A483 pub(crate) fn index_mut(&mut self, pool: &RRBPool<A>, level: usize, index: usize) -> &mut A { 484 if level == 0 { 485 &mut self.children.unwrap_values_mut(pool)[index] 486 } else { 487 let target_idx = self.index_in(level, index).unwrap(); 488 let offset = index - self.size_up_to(level, target_idx); 489 let child = &mut self.children.unwrap_nodes_mut(pool)[target_idx]; 490 child.index_mut(pool, level - 1, offset) 491 } 492 } 493 lookup_chunk( &self, level: usize, base: usize, index: usize, ) -> (Range<usize>, *const Chunk<A>)494 pub(crate) fn lookup_chunk( 495 &self, 496 level: usize, 497 base: usize, 498 index: usize, 499 ) -> (Range<usize>, *const Chunk<A>) { 500 if level == 0 { 501 ( 502 base..(base + self.children.len()), 503 self.children.unwrap_values() as *const Chunk<A>, 504 ) 505 } else { 506 let target_idx = self.index_in(level, index).unwrap(); 507 let offset = self.size_up_to(level, target_idx); 508 let child_base = base + offset; 509 let children = self.children.unwrap_nodes(); 510 let child = &children[target_idx]; 511 child.lookup_chunk(level - 1, child_base, index - offset) 512 } 513 } 514 lookup_chunk_mut( &mut self, pool: &RRBPool<A>, level: usize, base: usize, index: usize, ) -> (Range<usize>, *mut Chunk<A>)515 pub(crate) fn lookup_chunk_mut( 516 &mut self, 517 pool: &RRBPool<A>, 518 level: usize, 519 base: usize, 520 index: usize, 521 ) -> (Range<usize>, *mut Chunk<A>) { 522 if level == 0 { 523 ( 524 base..(base + self.children.len()), 525 self.children.unwrap_values_mut(pool) as *mut Chunk<A>, 526 ) 527 } else { 528 let target_idx = self.index_in(level, index).unwrap(); 529 let offset = self.size_up_to(level, target_idx); 530 let child_base = base + offset; 531 let children = self.children.unwrap_nodes_mut(pool); 532 let child = &mut children[target_idx]; 533 child.lookup_chunk_mut(pool, level - 1, child_base, index - offset) 534 } 535 } 536 push_child_node(&mut self, pool: &RRBPool<A>, side: Side, child: Node<A>)537 fn push_child_node(&mut self, pool: &RRBPool<A>, side: Side, child: Node<A>) { 538 let children = self.children.unwrap_nodes_mut(pool); 539 match side { 540 Left => children.push_front(child), 541 Right => children.push_back(child), 542 } 543 } 544 pop_child_node(&mut self, pool: &RRBPool<A>, side: Side) -> Node<A>545 fn pop_child_node(&mut self, pool: &RRBPool<A>, side: Side) -> Node<A> { 546 let children = self.children.unwrap_nodes_mut(pool); 547 match side { 548 Left => children.pop_front(), 549 Right => children.pop_back(), 550 } 551 } 552 push_chunk( &mut self, pool: &RRBPool<A>, level: usize, side: Side, mut chunk: PoolRef<Chunk<A>>, ) -> PushResult<PoolRef<Chunk<A>>>553 pub(crate) fn push_chunk( 554 &mut self, 555 pool: &RRBPool<A>, 556 level: usize, 557 side: Side, 558 mut chunk: PoolRef<Chunk<A>>, 559 ) -> PushResult<PoolRef<Chunk<A>>> { 560 if chunk.is_empty() { 561 return PushResult::Done; 562 } 563 let is_full = self.is_full(); 564 if level == 0 { 565 if self.children.is_empty_node() { 566 self.push_size(pool, side, level, chunk.len()); 567 self.children = Values(chunk); 568 PushResult::Done 569 } else { 570 let values = self.children.unwrap_values_mut(pool); 571 if values.len() + chunk.len() <= NODE_SIZE { 572 let chunk = PoolRef::make_mut(&pool.value_pool, &mut chunk); 573 match side { 574 Side::Left => { 575 chunk.append(values); 576 values.append(chunk); 577 } 578 Side::Right => values.append(chunk), 579 } 580 PushResult::Done 581 } else { 582 PushResult::Full(chunk, 0) 583 } 584 } 585 } else if level == 1 { 586 // If rightmost existing node has any room, merge as much as 587 // possible over from the new node. 588 let num_drained = match side { 589 Side::Right => { 590 if let Entry::Nodes(ref mut size, ref mut children) = self.children { 591 let rightmost = PoolRef::make_mut(&pool.node_pool, children) 592 .last_mut() 593 .unwrap(); 594 let old_size = rightmost.len(); 595 let chunk = PoolRef::make_mut(&pool.value_pool, &mut chunk); 596 let values = rightmost.children.unwrap_values_mut(pool); 597 let to_drain = chunk.len().min(NODE_SIZE - values.len()); 598 values.drain_from_front(chunk, to_drain); 599 size.pop(&pool.size_pool, Side::Right, level, old_size); 600 size.push(&pool.size_pool, Side::Right, level, values.len()); 601 to_drain 602 } else { 603 0 604 } 605 } 606 Side::Left => { 607 if let Entry::Nodes(ref mut size, ref mut children) = self.children { 608 let leftmost = PoolRef::make_mut(&pool.node_pool, children) 609 .first_mut() 610 .unwrap(); 611 let old_size = leftmost.len(); 612 let chunk = PoolRef::make_mut(&pool.value_pool, &mut chunk); 613 let values = leftmost.children.unwrap_values_mut(pool); 614 let to_drain = chunk.len().min(NODE_SIZE - values.len()); 615 values.drain_from_back(chunk, to_drain); 616 size.pop(&pool.size_pool, Side::Left, level, old_size); 617 size.push(&pool.size_pool, Side::Left, level, values.len()); 618 to_drain 619 } else { 620 0 621 } 622 } 623 }; 624 if is_full { 625 PushResult::Full(chunk, num_drained) 626 } else { 627 // If the chunk is empty after being drained, there might be 628 // more space in existing chunks. To keep the middle dense, we 629 // do not add it here. 630 if !chunk.is_empty() { 631 if side == Left && chunk.len() < NODE_SIZE { 632 if let Entry::Nodes(ref mut size, _) = self.children { 633 if let Size::Size(value) = *size { 634 *size = Size::table_from_size(&pool.size_pool, level, value); 635 } 636 } 637 } 638 self.push_size(pool, side, level, chunk.len()); 639 self.push_child_node(pool, side, Node::from_chunk(pool, 0, chunk)); 640 } 641 PushResult::Done 642 } 643 } else { 644 let chunk_size = chunk.len(); 645 let index = match side { 646 Right => self.children.len() - 1, 647 Left => 0, 648 }; 649 let new_child = { 650 let children = self.children.unwrap_nodes_mut(pool); 651 let child = &mut children[index]; 652 match child.push_chunk(pool, level - 1, side, chunk) { 653 PushResult::Done => None, 654 PushResult::Full(chunk, num_drained) => { 655 // Our chunk was too large for `child`, so it could not 656 // be pushed there. However, exactly `num_drained` 657 // elements were added to the child. We need to reflect 658 // that change in the size field of the node. 659 match side { 660 Right => match self.children { 661 Entry::Nodes(Size::Table(ref mut sizes), _) => { 662 let sizes = PoolRef::make_mut(&pool.size_pool, sizes); 663 sizes[index] += num_drained; 664 } 665 Entry::Nodes(Size::Size(ref mut size), _) => { 666 *size += num_drained; 667 } 668 Entry::Values(_) | Entry::Empty => (), 669 }, 670 Left => { 671 self.update_size(pool, 0, level, num_drained as isize); 672 } 673 } 674 if is_full { 675 return PushResult::Full(chunk, 0); 676 } else { 677 Some(Node::from_chunk(pool, level - 1, chunk)) 678 } 679 } 680 } 681 }; 682 match new_child { 683 None => { 684 self.update_size(pool, index, level, chunk_size as isize); 685 PushResult::Done 686 } 687 Some(child) => { 688 if side == Left && chunk_size < NODE_SIZE { 689 if let Entry::Nodes(ref mut size, _) = self.children { 690 if let Size::Size(value) = *size { 691 *size = Size::table_from_size(&pool.size_pool, level, value); 692 } 693 } 694 } 695 self.push_size(pool, side, level, child.len()); 696 self.push_child_node(pool, side, child); 697 PushResult::Done 698 } 699 } 700 } 701 } 702 pop_chunk( &mut self, pool: &RRBPool<A>, level: usize, side: Side, ) -> PopResult<PoolRef<Chunk<A>>>703 pub(crate) fn pop_chunk( 704 &mut self, 705 pool: &RRBPool<A>, 706 level: usize, 707 side: Side, 708 ) -> PopResult<PoolRef<Chunk<A>>> { 709 if self.is_empty() { 710 return PopResult::Empty; 711 } 712 if level == 0 { 713 // should only get here if the tree is just one leaf node 714 match replace(&mut self.children, Empty) { 715 Values(chunk) => PopResult::Drained(chunk), 716 Empty => panic!("rrb::Node::pop_chunk: non-empty tree with Empty leaf"), 717 Nodes(_, _) => panic!("rrb::Node::pop_chunk: branch node at leaf"), 718 } 719 } else if level == 1 { 720 let child_node = self.pop_child_node(pool, side); 721 self.pop_size(pool, side, level, child_node.len()); 722 let chunk = match child_node.children { 723 Values(ref chunk) => chunk.clone(), 724 Empty => panic!("rrb::Node::pop_chunk: non-empty tree with Empty leaf"), 725 Nodes(_, _) => panic!("rrb::Node::pop_chunk: branch node at leaf"), 726 }; 727 if self.is_empty() { 728 PopResult::Drained(chunk) 729 } else { 730 PopResult::Done(chunk) 731 } 732 } else { 733 let index = match side { 734 Right => self.children.len() - 1, 735 Left => 0, 736 }; 737 let mut drained = false; 738 let chunk = { 739 let children = self.children.unwrap_nodes_mut(pool); 740 let child = &mut children[index]; 741 match child.pop_chunk(pool, level - 1, side) { 742 PopResult::Empty => return PopResult::Empty, 743 PopResult::Done(chunk) => chunk, 744 PopResult::Drained(chunk) => { 745 drained = true; 746 chunk 747 } 748 } 749 }; 750 if drained { 751 self.pop_size(pool, side, level, chunk.len()); 752 self.pop_child_node(pool, side); 753 if self.is_empty() { 754 PopResult::Drained(chunk) 755 } else { 756 PopResult::Done(chunk) 757 } 758 } else { 759 self.update_size(pool, index, level, -(chunk.len() as isize)); 760 PopResult::Done(chunk) 761 } 762 } 763 } 764 split( &mut self, pool: &RRBPool<A>, level: usize, drop_side: Side, index: usize, ) -> SplitResult765 pub(crate) fn split( 766 &mut self, 767 pool: &RRBPool<A>, 768 level: usize, 769 drop_side: Side, 770 index: usize, 771 ) -> SplitResult { 772 if index == 0 && drop_side == Side::Left { 773 // Dropped nothing 774 return SplitResult::Dropped(0); 775 } 776 if level > 0 && index == 0 && drop_side == Side::Right { 777 // Dropped everything 778 let dropped = if let Entry::Nodes(ref size, _) = self.children { 779 size.size() 780 } else { 781 panic!("leaf node at non-leaf level!"); 782 }; 783 self.children = Entry::Empty; 784 return SplitResult::Dropped(dropped); 785 } 786 let mut dropped; 787 if level == 0 { 788 let len = self.children.len(); 789 if index >= len { 790 return SplitResult::OutOfBounds; 791 } 792 let children = self.children.unwrap_values_mut(pool); 793 match drop_side { 794 Side::Left => children.drop_left(index), 795 Side::Right => children.drop_right(index), 796 } 797 SplitResult::Dropped(match drop_side { 798 Left => index, 799 Right => len - index, 800 }) 801 } else if let Some(target_idx) = self.index_in(level, index) { 802 let size_up_to = self.size_up_to(level, target_idx); 803 let (size, children) = 804 if let Entry::Nodes(ref mut size, ref mut children) = self.children { 805 (size, PoolRef::make_mut(&pool.node_pool, children)) 806 } else { 807 unreachable!() 808 }; 809 let child_gone = 0 == { 810 let child_node = &mut children[target_idx]; 811 match child_node.split(pool, level - 1, drop_side, index - size_up_to) { 812 SplitResult::OutOfBounds => return SplitResult::OutOfBounds, 813 SplitResult::Dropped(amount) => dropped = amount, 814 } 815 child_node.len() 816 }; 817 match drop_side { 818 Left => { 819 let mut drop_from = target_idx; 820 if child_gone { 821 drop_from += 1; 822 } 823 children.drop_left(drop_from); 824 if let Size::Size(value) = *size { 825 *size = Size::table_from_size(&pool.size_pool, level, value); 826 } 827 let size_table = if let Size::Table(ref mut size_ref) = size { 828 PoolRef::make_mut(&pool.size_pool, size_ref) 829 } else { 830 unreachable!() 831 }; 832 let dropped_size = if target_idx > 0 { 833 size_table[target_idx - 1] 834 } else { 835 0 836 }; 837 dropped += dropped_size; 838 size_table.drop_left(drop_from); 839 for i in size_table.iter_mut() { 840 *i -= dropped; 841 } 842 } 843 Right => { 844 let at_last = target_idx == children.len() - 1; 845 let mut drop_from = target_idx + 1; 846 if child_gone { 847 drop_from -= 1; 848 } 849 if drop_from < children.len() { 850 children.drop_right(drop_from); 851 } 852 match size { 853 Size::Size(ref mut size) if at_last => { 854 *size -= dropped; 855 } 856 Size::Size(ref mut size) => { 857 let size_per_child = NODE_SIZE.pow(level as u32); 858 let remainder = (target_idx + 1) * size_per_child; 859 let new_size = remainder - dropped; 860 if new_size < *size { 861 dropped = *size - new_size; 862 *size = new_size; 863 } else { 864 unreachable!( 865 "this means node is empty, should be caught at start of method" 866 ); 867 } 868 } 869 Size::Table(ref mut size_ref) => { 870 let size_table = PoolRef::make_mut(&pool.size_pool, size_ref); 871 let dropped_size = 872 size_table[size_table.len() - 1] - size_table[target_idx]; 873 if drop_from < size_table.len() { 874 size_table.drop_right(drop_from); 875 } 876 if !child_gone { 877 size_table[target_idx] -= dropped; 878 } 879 dropped += dropped_size; 880 } 881 } 882 } 883 } 884 SplitResult::Dropped(dropped) 885 } else { 886 SplitResult::OutOfBounds 887 } 888 } 889 merge_leaves(pool: &RRBPool<A>, mut left: Self, mut right: Self) -> Self890 fn merge_leaves(pool: &RRBPool<A>, mut left: Self, mut right: Self) -> Self { 891 if left.children.is_empty_node() { 892 // Left is empty, just use right 893 Self::single_parent(pool, right) 894 } else if right.children.is_empty_node() { 895 // Right is empty, just use left 896 Self::single_parent(pool, left) 897 } else { 898 { 899 let left_vals = left.children.unwrap_values_mut(pool); 900 let left_len = left_vals.len(); 901 let right_vals = right.children.unwrap_values_mut(pool); 902 let right_len = right_vals.len(); 903 if left_len + right_len <= NODE_SIZE { 904 left_vals.append(right_vals); 905 } else { 906 let count = right_len.min(NODE_SIZE - left_len); 907 left_vals.drain_from_front(right_vals, count); 908 } 909 } 910 if right.is_empty() { 911 Self::single_parent(pool, left) 912 } else { 913 Self::join_dense(pool, left, right) 914 } 915 } 916 } 917 merge_rebalance( pool: &RRBPool<A>, level: usize, left: Self, middle: Self, right: Self, ) -> Self918 fn merge_rebalance( 919 pool: &RRBPool<A>, 920 level: usize, 921 left: Self, 922 middle: Self, 923 right: Self, 924 ) -> Self { 925 let left_nodes = left.children.nodes().into_iter(); 926 let middle_nodes = middle.children.nodes().into_iter(); 927 let right_nodes = right.children.nodes().into_iter(); 928 let mut subtree_still_balanced = true; 929 let mut next_leaf = Chunk::new(); 930 let mut next_node = Chunk::new(); 931 let mut next_subtree = Chunk::new(); 932 let mut root = Chunk::new(); 933 934 for subtree in left_nodes.chain(middle_nodes).chain(right_nodes) { 935 if subtree.is_empty() { 936 continue; 937 } 938 if subtree.is_completely_dense(level) && subtree_still_balanced { 939 root.push_back(subtree); 940 continue; 941 } 942 subtree_still_balanced = false; 943 944 if level == 1 { 945 for value in subtree.children.values() { 946 next_leaf.push_back(value); 947 if next_leaf.is_full() { 948 let new_node = 949 Node::from_chunk(pool, 0, PoolRef::new(&pool.value_pool, next_leaf)); 950 next_subtree.push_back(new_node); 951 next_leaf = Chunk::new(); 952 if next_subtree.is_full() { 953 let new_subtree = Node::parent(pool, level, next_subtree); 954 root.push_back(new_subtree); 955 next_subtree = Chunk::new(); 956 } 957 } 958 } 959 } else { 960 for node in subtree.children.nodes() { 961 next_node.push_back(node); 962 if next_node.is_full() { 963 let new_node = Node::parent(pool, level - 1, next_node); 964 next_subtree.push_back(new_node); 965 next_node = Chunk::new(); 966 if next_subtree.is_full() { 967 let new_subtree = Node::parent(pool, level, next_subtree); 968 root.push_back(new_subtree); 969 next_subtree = Chunk::new(); 970 } 971 } 972 } 973 } 974 } 975 if !next_leaf.is_empty() { 976 let new_node = Node::from_chunk(pool, 0, PoolRef::new(&pool.value_pool, next_leaf)); 977 next_subtree.push_back(new_node); 978 } 979 if !next_node.is_empty() { 980 let new_node = Node::parent(pool, level - 1, next_node); 981 next_subtree.push_back(new_node); 982 } 983 if !next_subtree.is_empty() { 984 let new_subtree = Node::parent(pool, level, next_subtree); 985 root.push_back(new_subtree); 986 } 987 Node::parent(pool, level + 1, root) 988 } 989 merge(pool: &RRBPool<A>, mut left: Self, mut right: Self, level: usize) -> Self990 pub(crate) fn merge(pool: &RRBPool<A>, mut left: Self, mut right: Self, level: usize) -> Self { 991 if level == 0 { 992 Self::merge_leaves(pool, left, right) 993 } else { 994 let merged = { 995 if level == 1 { 996 // We're going to rebalance all the leaves anyway, there's 997 // no need for a middle at level 1 998 Node::parent(pool, 0, Chunk::new()) 999 } else { 1000 let left_last = 1001 if let Entry::Nodes(ref mut size, ref mut children) = left.children { 1002 let node = PoolRef::make_mut(&pool.node_pool, children).pop_back(); 1003 if !node.is_empty() { 1004 size.pop(&pool.size_pool, Side::Right, level, node.len()); 1005 } 1006 node 1007 } else { 1008 panic!("expected nodes, found entries or empty"); 1009 }; 1010 let right_first = 1011 if let Entry::Nodes(ref mut size, ref mut children) = right.children { 1012 let node = PoolRef::make_mut(&pool.node_pool, children).pop_front(); 1013 if !node.is_empty() { 1014 size.pop(&pool.size_pool, Side::Left, level, node.len()); 1015 } 1016 node 1017 } else { 1018 panic!("expected nodes, found entries or empty"); 1019 }; 1020 Self::merge(pool, left_last, right_first, level - 1) 1021 } 1022 }; 1023 Self::merge_rebalance(pool, level, left, merged, right) 1024 } 1025 } 1026 1027 #[cfg(any(test, feature = "debug"))] assert_invariants(&self, level: usize) -> usize1028 pub(crate) fn assert_invariants(&self, level: usize) -> usize { 1029 // Verifies that the size table matches reality. 1030 match self.children { 1031 Entry::Empty => 0, 1032 Entry::Values(ref values) => { 1033 // An empty value node is pointless and should never occur. 1034 assert_ne!(0, values.len()); 1035 // Value nodes should only occur at level 0. 1036 assert_eq!(0, level); 1037 values.len() 1038 } 1039 Entry::Nodes(ref size, ref children) => { 1040 // A parent node with no children should never occur. 1041 assert_ne!(0, children.len()); 1042 // Parent nodes should never occur at level 0. 1043 assert_ne!(0, level); 1044 let mut lengths = Vec::new(); 1045 let should_be_dense = if let Size::Size(_) = size { 1046 true 1047 } else { 1048 false 1049 }; 1050 for (index, child) in children.iter().enumerate() { 1051 let len = child.assert_invariants(level - 1); 1052 if should_be_dense && index < children.len() - 1 { 1053 // Assert that non-end nodes without size tables are full. 1054 assert_eq!(len, NODE_SIZE.pow(level as u32)); 1055 } 1056 lengths.push(len); 1057 } 1058 match size { 1059 Size::Size(size) => { 1060 let total: usize = lengths.iter().sum(); 1061 assert_eq!(*size, total); 1062 } 1063 Size::Table(ref table) => { 1064 assert_eq!(table.iter().len(), children.len()); 1065 for (index, current) in table.iter().enumerate() { 1066 let expected: usize = lengths.iter().take(index + 1).sum(); 1067 assert_eq!(expected, *current); 1068 } 1069 } 1070 } 1071 lengths.iter().sum() 1072 } 1073 } 1074 } 1075 1076 // pub fn print<W>(&self, f: &mut W, indent: usize, level: usize) -> Result<(), fmt::Error> 1077 // where 1078 // W: fmt::Write, 1079 // A: fmt::Debug, 1080 // { 1081 // print_indent(f, indent)?; 1082 // if level == 0 { 1083 // if self.children.is_empty_node() { 1084 // writeln!(f, "Leaf: EMPTY") 1085 // } else { 1086 // writeln!(f, "Leaf: {:?}", self.children.unwrap_values()) 1087 // } 1088 // } else { 1089 // match &self.children { 1090 // Entry::Nodes(size, children) => { 1091 // writeln!(f, "Node level {} size_table {:?}", level, size)?; 1092 // for child in children.iter() { 1093 // child.print(f, indent + 4, level - 1)?; 1094 // } 1095 // Ok(()) 1096 // } 1097 // _ => unreachable!(), 1098 // } 1099 // } 1100 // } 1101 } 1102 1103 // fn print_indent<W>(f: &mut W, indent: usize) -> Result<(), fmt::Error> 1104 // where 1105 // W: fmt::Write, 1106 // { 1107 // for _i in 0..indent { 1108 // write!(f, " ")?; 1109 // } 1110 // Ok(()) 1111 // } 1112