From 286949c36f576eee02ae178cb011c89b3e28147c Mon Sep 17 00:00:00 2001 From: Johannes Altmanninger Date: Fri, 8 Mar 2019 13:14:35 +0100 Subject: [PATCH] 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 | 1 + src/nodes/rrb.rs | 39 ++++++++++++++++++++++++++++++++------- src/vector/mod.rs | 36 +++++++++++++++++++++++++++++++++++- 3 files changed, 68 insertions(+), 8 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 02def48..d576702 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -335,6 +335,7 @@ #![deny(unsafe_code)] #![cfg_attr(has_specialisation, feature(specialization))] +#![feature(stmt_expr_attributes)] #[cfg(test)] #[macro_use] diff --git a/src/nodes/rrb.rs b/src/nodes/rrb.rs index 8cc941c..30f973c 100644 --- a/src/nodes/rrb.rs +++ b/src/nodes/rrb.rs @@ -120,7 +120,7 @@ impl Size { } pub enum PushResult { - Full(A), + Full(A, usize), Done, } @@ -551,13 +551,13 @@ impl Node { } 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 Node { 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 Node { 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 Node { 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)) } diff --git a/src/vector/mod.rs b/src/vector/mod.rs index f189c3a..a028e48 100644 --- a/src/vector/mod.rs +++ b/src/vector/mod.rs @@ -1555,7 +1555,7 @@ impl RRB { 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)) {