diff --git a/fuzz/fuzz_targets/fuzz_redb.rs b/fuzz/fuzz_targets/fuzz_redb.rs index a7809486..145df851 100644 --- a/fuzz/fuzz_targets/fuzz_redb.rs +++ b/fuzz/fuzz_targets/fuzz_redb.rs @@ -687,10 +687,12 @@ fn assert_multimap_value_eq( reference: Option<&BTreeSet>, ) -> Result<(), redb::Error> { if let Some(values) = reference { + assert_eq!(values.len() as u64, iter.len()); for value in values.iter() { assert_eq!(iter.next().unwrap()?.value().len(), *value); } } + assert!(iter.is_empty()); // This is basically assert!(iter.next().is_none()), but we also allow an Err such as a simulated IO error if let Some(Ok(_)) = iter.next() { panic!(); diff --git a/src/multimap_table.rs b/src/multimap_table.rs index 5f16e0d8..101b8ad9 100644 --- a/src/multimap_table.rs +++ b/src/multimap_table.rs @@ -505,6 +505,7 @@ impl DynamicCollection { let root = collection.value().as_subtree().root; MultimapValue::new_subtree( BtreeRangeIter::new::>(&(..), Some(root), mem)?, + collection.value().get_num_values(), guard, ) } @@ -518,6 +519,7 @@ impl DynamicCollection { guard: Arc, mem: Arc, ) -> Result> { + let num_values = collection.value().get_num_values(); Ok(match collection.value().collection_type() { Inline => { let leaf_iter = @@ -531,7 +533,14 @@ impl DynamicCollection { Some(root), mem.clone(), )?; - MultimapValue::new_subtree_free_on_drop(inner, freed_pages, pages, guard, mem) + MultimapValue::new_subtree_free_on_drop( + inner, + num_values, + freed_pages, + pages, + guard, + mem, + ) } }) } @@ -573,6 +582,7 @@ enum ValueIterState<'a, V: Key + 'static> { pub struct MultimapValue<'a, V: Key + 'static> { inner: Option>, + remaining: u64, freed_pages: Option>>>, free_on_drop: Vec, _transaction_guard: Arc, @@ -581,9 +591,14 @@ pub struct MultimapValue<'a, V: Key + 'static> { } impl<'a, V: Key + 'static> MultimapValue<'a, V> { - fn new_subtree(inner: BtreeRangeIter, guard: Arc) -> Self { + fn new_subtree( + inner: BtreeRangeIter, + num_values: u64, + guard: Arc, + ) -> Self { Self { inner: Some(ValueIterState::Subtree(inner)), + remaining: num_values, freed_pages: None, free_on_drop: vec![], _transaction_guard: guard, @@ -594,6 +609,7 @@ impl<'a, V: Key + 'static> MultimapValue<'a, V> { fn new_subtree_free_on_drop( inner: BtreeRangeIter, + num_values: u64, freed_pages: Arc>>, pages: Vec, guard: Arc, @@ -601,6 +617,7 @@ impl<'a, V: Key + 'static> MultimapValue<'a, V> { ) -> Self { Self { inner: Some(ValueIterState::Subtree(inner)), + remaining: num_values, freed_pages: Some(freed_pages), free_on_drop: pages, _transaction_guard: guard, @@ -610,8 +627,10 @@ impl<'a, V: Key + 'static> MultimapValue<'a, V> { } fn new_inline(inner: LeafKeyIter<'a, V>, guard: Arc) -> Self { + let remaining = inner.inline_collection.value().get_num_values(); Self { inner: Some(ValueIterState::InlineLeaf(inner)), + remaining, freed_pages: None, free_on_drop: vec![], _transaction_guard: guard, @@ -619,6 +638,17 @@ impl<'a, V: Key + 'static> MultimapValue<'a, V> { _value_type: Default::default(), } } + + /// Returns the number of times this iterator will return `Some(Ok(_))` + /// + /// Note that `Some` may be returned from `next()` more than `len()` times if `Some(Err(_))` is returned + pub fn len(&self) -> u64 { + self.remaining + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } } impl<'a, V: Key + 'static> Iterator for MultimapValue<'a, V> { @@ -635,6 +665,7 @@ impl<'a, V: Key + 'static> Iterator for MultimapValue<'a, V> { }, ValueIterState::InlineLeaf(ref mut iter) => iter.next_key()?.to_vec(), }; + self.remaining -= 1; Some(Ok(AccessGuard::with_owned_value(bytes))) } } @@ -1122,6 +1153,7 @@ impl<'txn, K: Key + 'static, V: Key + 'static> MultimapTable<'txn, K, V> { } else { MultimapValue::new_subtree( BtreeRangeIter::new::>(&(..), None, self.mem.clone())?, + 0, self.transaction.transaction_guard(), ) }; @@ -1169,6 +1201,7 @@ impl<'txn, K: Key + 'static, V: Key + 'static> ReadableMultimapTable } else { MultimapValue::new_subtree( BtreeRangeIter::new::>(&(..), None, self.mem.clone())?, + 0, guard, ) }; @@ -1311,6 +1344,7 @@ impl ReadOnlyMultimapTable { } else { MultimapValue::new_subtree( BtreeRangeIter::new::>(&(..), None, self.mem.clone())?, + 0, self.transaction_guard.clone(), ) }; @@ -1370,6 +1404,7 @@ impl ReadableMultimapTable } else { MultimapValue::new_subtree( BtreeRangeIter::new::>(&(..), None, self.mem.clone())?, + 0, self.transaction_guard.clone(), ) }; diff --git a/tests/multimap_tests.rs b/tests/multimap_tests.rs index 567e94d5..9a0d03c4 100644 --- a/tests/multimap_tests.rs +++ b/tests/multimap_tests.rs @@ -249,6 +249,7 @@ fn delete() { let read_txn = db.begin_read().unwrap(); let table = read_txn.open_multimap_table(STR_TABLE).unwrap(); + assert_eq!(3, table.get("hello").unwrap().len()); assert_eq!( vec![ "world".to_string(),