Update RRB size table after draining in push_chunk

When a chunk does not fully fit in a node, we drain as many elements from the
chunk into the node and later push the leftover chunk to a sibling.

This commit makes sure that the size field of the parent node is updated
correctly. For dense nodes, we can simply add the number of drained elements to
the size. If there is a size table, we add that number to the size entry
corresponding to the current element and all entries to the right.

Closes #74

src/lib.rs

@@ -335,6 +335,7 @@
 
 #![deny(unsafe_code)]
 #![cfg_attr(has_specialisation, feature(specialization))]
+#![feature(stmt_expr_attributes)]
 
 #[cfg(test)]
 #[macro_use]

src/nodes/rrb.rs

@@ -120,7 +120,7 @@ impl Size {
 }
 
 pub enum PushResult<A> {
-    Full(A),
+    Full(A, usize),
     Done,
 }
 
@@ -551,13 +551,13 @@ impl<A: Clone> Node<A> {
                     }
                     PushResult::Done
                 } else {
-                    PushResult::Full(chunk)
+                    PushResult::Full(chunk, 0)
                 }
             }
         } else if level == 1 {
             // If rightmost existing node has any room, merge as much as
             // possible over from the new node.
-            match side {
+            let num_drained = match side {
                 Side::Right => {
                     if let Entry::Nodes(ref mut size, ref mut children) = self.children {
                         let rightmost = Ref::make_mut(Ref::make_mut(children).last_mut().unwrap());
@@ -568,6 +568,9 @@ impl<A: Clone> Node<A> {
                         values.drain_from_front(chunk, to_drain);
                         size.pop(Side::Right, old_size);
                         size.push(Side::Right, values.len());
+                        to_drain
+                    } else {
+                        0
                     }
                 }
                 Side::Left => {
@@ -580,11 +583,14 @@ impl<A: Clone> Node<A> {
                         values.drain_from_back(chunk, to_drain);
                         size.pop(Side::Left, old_size);
                         size.push(Side::Left, values.len());
+                        to_drain
+                    } else {
+                        0
                     }
                 }
-            }
+            };
             if is_full {
-                PushResult::Full(chunk)
+                PushResult::Full(chunk, num_drained)
             } else {
                 // If the chunk is empty after being drained, there might be
                 // more space in existing chunks. To keep the middle dense, we
@@ -606,9 +612,28 @@ impl<A: Clone> Node<A> {
                 let child = Ref::make_mut(&mut children[index]);
                 match child.push_chunk(level - 1, side, chunk) {
                     PushResult::Done => None,
-                    PushResult::Full(chunk) => {
+                    PushResult::Full(chunk, num_drained) => {
+                        // Our chunk was too large for `child`, so it could not
+                        // be pushed there. However, exactly `num_drained`
+                        // elements were added to the child. We need to reflect
+                        // that change in the size field of the node.
+                        match side {
+                            Right => match self.children {
+                                Entry::Nodes(Size::Table(ref mut sizes), _) => {
+                                    let sizes = Ref::make_mut(sizes);
+                                    sizes[index] += num_drained;
+                                }
+                                Entry::Nodes(Size::Size(ref mut size), _) => {
+                                    *size += num_drained;
+                                }
+                                Entry::Values(_) | Entry::Empty => (),
+                            },
+                            Left => {
+                                self.update_size(0, num_drained as isize);
+                            }
+                        }
                         if is_full {
-                            return PushResult::Full(chunk);
+                            return PushResult::Full(chunk, 0);
                         } else {
                             Some(Node::from_chunk(level - 1, chunk))
                         }

src/vector/mod.rs

@@ -54,7 +54,7 @@ use std::ops::{Add, Index, IndexMut, RangeBounds};
 
 use crate::nodes::chunk::{Chunk, Iter as ChunkIter, CHUNK_SIZE};
 use crate::nodes::rrb::{
-    ConsumingIter as ConsumingNodeIter, Node, PopResult, PushResult, SplitResult,
+    ConsumingIter as ConsumingNodeIter, Node, PopResult, PushResult, SplitResult, 
 };
 use crate::sort;
 use crate::util::{clone_ref, swap_indices, to_range, Ref, Side};
@@ -1555,7 +1555,7 @@ impl<A: Clone> RRB<A> {
             let middle = Ref::make_mut(&mut self.middle);
             match middle.push_chunk(self.middle_level, side, chunk) {
                 PushResult::Done => return,
-                PushResult::Full(chunk) => Ref::from({
+                PushResult::Full(chunk, _num_drained) => Ref::from({
                     match side {
                         Side::Left => Node::from_chunk(self.middle_level, chunk)
                             .join_branches(middle.clone(), self.middle_level),
@@ -2593,6 +2593,40 @@ mod test {
         assert_eq!(Some(&1), tail.get(0));
     }
 
+    #[test]
+    fn issue_74_simple_size() {
+        use crate::nodes::rrb::NODE_SIZE;
+        let mut x = Vector::new();
+        #[rustfmt::skip]
+        for _ in 0..(CHUNK_SIZE *
+            ( 1 // inner_f
+            + (2 * NODE_SIZE) // middle: two full Entry::Nodes (4096 elements each)
+            + 1 // inner_b
+            + 1 // outer_b
+            ))
+        {
+            x.push_back(0u32);
+        }
+        let middle_first_node_start = CHUNK_SIZE;
+        let middle_second_node_start = middle_first_node_start + NODE_SIZE * CHUNK_SIZE;
+        // This reduces the size of the second node to 4095.
+        x.remove(middle_second_node_start);
+        // As outer_b is full, this will cause inner_b (length 64) to be pushed
+        // to middle. The first element will be merged into the second node, the
+        // remaining 63 elements will end up in a new node.
+        x.push_back(0u32);
+        match x {
+            Vector::Full(tree) => {
+                assert_eq!(3, tree.middle.number_of_children());
+                assert_eq!(
+                    2 * NODE_SIZE * CHUNK_SIZE + CHUNK_SIZE - 1,
+                    tree.middle.len()
+                );
+            }
+            _ => unreachable!(),
+        }
+    }
+
     proptest! {
         #[test]
         fn iter(ref vec in vec(i32::ANY, 0..1000)) {