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