Update RRB size table after draining in push_chunk

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

@@ -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)) {