im/nodes/
rrb.rs

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

im/nodes/
rrb.rs
