Skip to content

Commit

Permalink
Prune list iterators (#3574)
Browse files Browse the repository at this point in the history 
* wip

* use range beneath subtree for efficient is_pruned check

* various iterators over the prune list

* improved prune list iter and subtree handling

* use take_while so unpruned iterators are not infinite
  • Loading branch information
@antiochp
antiochp committed Feb 24, 2021
1 parent adddff9 commit 3583028
Show file tree
Hide file tree
Showing 2 changed files with 176 additions and 25 deletions.
73 changes: 67 additions & 6 deletions store/src/prune_list.rs
View file
Expand Up @@ -21,8 +21,11 @@
//! must be shifted the appropriate amount when reading from the hash and data
//! files.

use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::{
io::{self, Write},
ops::Range,
};

use croaring::Bitmap;
use grin_core::core::pmmr;
Expand Down Expand Up @@ -256,11 +259,6 @@ impl PruneList {
self.bitmap.is_empty()
}

/// Convert the prune_list to a vec of pos.
pub fn to_vec(&self) -> Vec<u64> {
self.bitmap.iter().map(|x| x as u64).collect()
}

/// A pos is pruned if it is a pruned root directly or if it is
/// beneath the "next" pruned subtree.
/// We only need to consider the "next" subtree due to the append-only MMR structure.
Expand All @@ -283,4 +281,67 @@ impl PruneList {
assert!(pos > 0, "prune list 1-indexed, 0 not valid pos");
self.bitmap.contains(pos as u32)
}

/// Iterator over the entries in the prune list (pruned roots).
pub fn iter(&self) -> impl Iterator<Item = u64> + '_ {
self.bitmap.iter().map(|x| x as u64)
}

/// Iterator over the pruned "bintree range" for each pruned root.
pub fn pruned_bintree_range_iter(&self) -> impl Iterator<Item = Range<u64>> + '_ {
self.iter().map(|x| pmmr::bintree_range(x))
}

/// Iterator over all pos that are *not* pruned based on current prune_list.
pub fn unpruned_iter(&self, cutoff_pos: u64) -> impl Iterator<Item = u64> + '_ {
UnprunedIterator::new(self.pruned_bintree_range_iter())
.take_while(move |x| *x <= cutoff_pos)
}

/// Iterator over all leaf pos that are *not* pruned based on current prune_list.
/// Note this is not necessarily the same as the "leaf_set" as an output
/// can be spent but not yet pruned.
pub fn unpruned_leaf_iter(&self, cutoff_pos: u64) -> impl Iterator<Item = u64> + '_ {
self.unpruned_iter(cutoff_pos).filter(|x| pmmr::is_leaf(*x))
}
}

struct UnprunedIterator<I> {
inner: I,
current_excl_range: Option<Range<u64>>,
current_pos: u64,
}

impl<I: Iterator<Item = Range<u64>>> UnprunedIterator<I> {
fn new(mut inner: I) -> UnprunedIterator<I> {
let current_excl_range = inner.next();
UnprunedIterator {
inner,
current_excl_range,
current_pos: 1,
}
}
}

impl<I: Iterator<Item = Range<u64>>> Iterator for UnprunedIterator<I> {
type Item = u64;

fn next(&mut self) -> Option<Self::Item> {
if let Some(range) = &self.current_excl_range {
if self.current_pos < range.start {
let next = self.current_pos;
self.current_pos += 1;
Some(next)
} else {
// skip the entire excluded range, moving to next excluded range as necessary
self.current_pos = range.end;
self.current_excl_range = self.inner.next();
self.next()
}
} else {
let next = self.current_pos;
self.current_pos += 1;
Some(next)
}
}
}
128 changes: 109 additions & 19 deletions store/tests/prune_list.rs
View file
Expand Up @@ -42,38 +42,43 @@ fn test_is_pruned() {
assert_eq!(pl.is_pruned(3), false);

pl.add(2);
pl.flush().unwrap();

assert_eq!(pl.to_vec(), vec![2]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [2]);
assert_eq!(pl.is_pruned(1), false);
assert_eq!(pl.is_pruned(2), true);
assert_eq!(pl.is_pruned(3), false);
assert_eq!(pl.is_pruned(4), false);

pl.add(2);
pl.add(1);
pl.flush().unwrap();

assert_eq!(pl.len(), 1);
assert_eq!(pl.to_vec(), [3]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [3]);
assert_eq!(pl.is_pruned(1), true);
assert_eq!(pl.is_pruned(2), true);
assert_eq!(pl.is_pruned(3), true);
assert_eq!(pl.is_pruned(4), false);

pl.add(4);

// Flushing the prune_list removes any individual leaf positions.
// This assumes we will track these outside the prune_list via the leaf_set.
pl.flush().unwrap();

assert_eq!(pl.len(), 2);
assert_eq!(pl.to_vec(), [3, 4]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [3, 4]);
assert_eq!(pl.is_pruned(1), true);
assert_eq!(pl.is_pruned(2), true);
assert_eq!(pl.is_pruned(3), true);
assert_eq!(pl.is_pruned(4), true);
assert_eq!(pl.is_pruned(5), false);

// Test some poorly organized (out of order, overlapping) pruning.
let mut pl = PruneList::empty();
pl.add(2);
pl.add(4);
pl.add(3);
assert_eq!(pl.iter().collect::<Vec<_>>(), [3, 4]);

// now add a higher level pruned root clearing out the subtree.
pl.add(7);
assert_eq!(pl.iter().collect::<Vec<_>>(), [7]);
}

#[test]
Expand All @@ -93,7 +98,7 @@ fn test_get_leaf_shift() {
pl.add(1);
pl.flush().unwrap();

assert_eq!(pl.to_vec(), vec![1]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [1]);
assert_eq!(pl.get_leaf_shift(1), 0);
assert_eq!(pl.get_leaf_shift(2), 0);
assert_eq!(pl.get_leaf_shift(3), 0);
Expand All @@ -119,7 +124,7 @@ fn test_get_leaf_shift() {
pl.flush().unwrap();

assert_eq!(pl.len(), 2);
assert_eq!(pl.to_vec(), [3, 4]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [3, 4]);
assert_eq!(pl.get_leaf_shift(1), 0);
assert_eq!(pl.get_leaf_shift(2), 0);
assert_eq!(pl.get_leaf_shift(3), 2);
Expand All @@ -138,7 +143,7 @@ fn test_get_leaf_shift() {
pl.flush().unwrap();

assert_eq!(pl.len(), 1);
assert_eq!(pl.to_vec(), [7]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [7]);
assert_eq!(pl.get_leaf_shift(1), 0);
assert_eq!(pl.get_leaf_shift(2), 0);
assert_eq!(pl.get_leaf_shift(3), 0);
Expand All @@ -159,7 +164,7 @@ fn test_get_leaf_shift() {
pl.flush().unwrap();

assert_eq!(pl.len(), 2);
assert_eq!(pl.to_vec(), [6, 13]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [6, 13]);
assert_eq!(pl.get_leaf_shift(2), 0);
assert_eq!(pl.get_leaf_shift(4), 0);
assert_eq!(pl.get_leaf_shift(8), 2);
Expand All @@ -182,7 +187,7 @@ fn test_get_shift() {
pl.add(1);
pl.flush().unwrap();

assert_eq!(pl.to_vec(), [1]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [1]);
assert_eq!(pl.get_shift(1), 0);
assert_eq!(pl.get_shift(2), 0);
assert_eq!(pl.get_shift(3), 0);
Expand All @@ -191,7 +196,7 @@ fn test_get_shift() {
pl.add(2);
pl.flush().unwrap();

assert_eq!(pl.to_vec(), [3]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [3]);
assert_eq!(pl.get_shift(1), 0);
assert_eq!(pl.get_shift(2), 0);
assert_eq!(pl.get_shift(3), 2);
Expand All @@ -204,7 +209,7 @@ fn test_get_shift() {
pl.add(3);
pl.flush().unwrap();

assert_eq!(pl.to_vec(), [3]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [3]);
assert_eq!(pl.get_shift(1), 0);
assert_eq!(pl.get_shift(2), 0);
assert_eq!(pl.get_shift(3), 2);
Expand All @@ -215,7 +220,7 @@ fn test_get_shift() {
pl.add(4);
pl.flush().unwrap();

assert_eq!(pl.to_vec(), [3, 4]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [3, 4]);
assert_eq!(pl.get_shift(1), 0);
assert_eq!(pl.get_shift(2), 0);
assert_eq!(pl.get_shift(3), 2);
Expand All @@ -227,7 +232,7 @@ fn test_get_shift() {
pl.add(5);
pl.flush().unwrap();

assert_eq!(pl.to_vec(), [7]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [7]);
assert_eq!(pl.get_shift(1), 0);
assert_eq!(pl.get_shift(2), 0);
assert_eq!(pl.get_shift(3), 0);
Expand All @@ -254,7 +259,7 @@ fn test_get_shift() {
pl.add(4);
pl.flush().unwrap();

assert_eq!(pl.to_vec(), [6, 10]);
assert_eq!(pl.iter().collect::<Vec<_>>(), [6, 10]);
assert_eq!(pl.get_shift(1), 0);
assert_eq!(pl.get_shift(2), 0);
assert_eq!(pl.get_shift(3), 0);
Expand All @@ -268,3 +273,88 @@ fn test_get_shift() {
assert_eq!(pl.get_shift(11), 4);
assert_eq!(pl.get_shift(12), 4);
}

#[test]
pub fn test_iter() {
let mut pl = PruneList::empty();
pl.add(1);
pl.add(2);
pl.add(4);
assert_eq!(pl.iter().collect::<Vec<_>>(), [3, 4]);

let mut pl = PruneList::empty();
pl.add(1);
pl.add(2);
pl.add(5);
assert_eq!(pl.iter().collect::<Vec<_>>(), [3, 5]);
}

#[test]
pub fn test_pruned_bintree_range_iter() {
let mut pl = PruneList::empty();
pl.add(1);
pl.add(2);
pl.add(4);
assert_eq!(
pl.pruned_bintree_range_iter().collect::<Vec<_>>(),
[1..4, 4..5]
);

let mut pl = PruneList::empty();
pl.add(1);
pl.add(2);
pl.add(5);
assert_eq!(
pl.pruned_bintree_range_iter().collect::<Vec<_>>(),
[1..4, 5..6]
);
}

#[test]
pub fn test_unpruned_iter() {
let pl = PruneList::empty();
assert_eq!(pl.unpruned_iter(5).collect::<Vec<_>>(), [1, 2, 3, 4, 5]);

let mut pl = PruneList::empty();
pl.add(2);
assert_eq!(pl.iter().collect::<Vec<_>>(), [2]);
assert_eq!(pl.pruned_bintree_range_iter().collect::<Vec<_>>(), [2..3]);
assert_eq!(pl.unpruned_iter(4).collect::<Vec<_>>(), [1, 3, 4]);

let mut pl = PruneList::empty();
pl.add(2);
pl.add(4);
pl.add(5);
assert_eq!(pl.iter().collect::<Vec<_>>(), [2, 6]);
assert_eq!(
pl.pruned_bintree_range_iter().collect::<Vec<_>>(),
[2..3, 4..7]
);
assert_eq!(pl.unpruned_iter(9).collect::<Vec<_>>(), [1, 3, 7, 8, 9]);
}

#[test]
fn test_unpruned_leaf_iter() {
let pl = PruneList::empty();
assert_eq!(
pl.unpruned_leaf_iter(8).collect::<Vec<_>>(),
[1, 2, 4, 5, 8]
);

let mut pl = PruneList::empty();
pl.add(2);
assert_eq!(pl.iter().collect::<Vec<_>>(), [2]);
assert_eq!(pl.pruned_bintree_range_iter().collect::<Vec<_>>(), [2..3]);
assert_eq!(pl.unpruned_leaf_iter(5).collect::<Vec<_>>(), [1, 4, 5]);

let mut pl = PruneList::empty();
pl.add(2);
pl.add(4);
pl.add(5);
assert_eq!(pl.iter().collect::<Vec<_>>(), [2, 6]);
assert_eq!(
pl.pruned_bintree_range_iter().collect::<Vec<_>>(),
[2..3, 4..7]
);
assert_eq!(pl.unpruned_leaf_iter(9).collect::<Vec<_>>(), [1, 8, 9]);
}

0 comments on commit 3583028

Please sign in to comment.