Skip to content

Commit

Permalink
refactor orchard_subtrees() to use ranges
Browse files Browse the repository at this point in the history
  • Loading branch information
@arya2
arya2 committed Oct 12, 2023
1 parent 837061f commit a7edaa9
Show file tree
Hide file tree
Showing 4 changed files with 42 additions and 122 deletions.
27 changes: 21 additions & 6 deletions zebra-state/src/service.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1548,14 +1548,29 @@ impl Service<ReadRequest> for ReadStateService {

tokio::task::spawn_blocking(move || {
span.in_scope(move || {
let end_index = limit
.and_then(|limit| start_index.0.checked_add(limit.0))
.map(NoteCommitmentSubtreeIndex);

let orchard_subtrees = state.non_finalized_state_receiver.with_watch_data(
|non_finalized_state| {
read::orchard_subtrees(
non_finalized_state.best_chain(),
&state.db,
start_index,
limit,
)
if let Some(end_index) = end_index {
read::orchard_subtrees(
non_finalized_state.best_chain(),
&state.db,
start_index..end_index,
)
} else {
// If there is no end bound, just return all the trees.
// If the end bound would overflow, just returns all the trees, because that's what
// `zcashd` does. (It never calculates an end bound, so it just keeps iterating until
// the trees run out.)
read::orchard_subtrees(
non_finalized_state.best_chain(),
&state.db,
start_index..,
)
}
},
);

Expand Down
45 changes: 3 additions & 42 deletions zebra-state/src/service/finalized_state/zebra_db/shielded.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -382,8 +382,7 @@ impl ZebraDb {
Some(subtree_data.with_index(index))
}

/// Returns a list of Orchard [`NoteCommitmentSubtree`]s starting at `start_index`.
/// If `limit` is provided, the list is limited to `limit` entries.
/// Returns a list of Orchard [`NoteCommitmentSubtree`]s in the provided range.
///
/// If there is no subtree at `start_index`, the returned list is empty.
/// Otherwise, subtrees are continuous up to the finalized tip.
Expand All @@ -395,52 +394,14 @@ impl ZebraDb {
#[allow(clippy::unwrap_in_result)]
pub fn orchard_subtree_list_by_index_for_rpc(
&self,
start_index: NoteCommitmentSubtreeIndex,
limit: Option<NoteCommitmentSubtreeIndex>,
range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>> {
let orchard_subtrees = self
.db
.cf_handle("orchard_note_commitment_subtree")
.unwrap();

// Calculate the end bound, checking for overflow.
let exclusive_end_bound: Option<NoteCommitmentSubtreeIndex> = limit
.and_then(|limit| start_index.0.checked_add(limit.0))
.map(NoteCommitmentSubtreeIndex);

let list: BTreeMap<
NoteCommitmentSubtreeIndex,
NoteCommitmentSubtreeData<orchard::tree::Node>,
>;

if let Some(exclusive_end_bound) = exclusive_end_bound {
list = self
.db
.zs_range_iter(&orchard_subtrees, start_index..exclusive_end_bound)
.collect();
} else {
// If there is no end bound, just return all the trees.
// If the end bound would overflow, just returns all the trees, because that's what
// `zcashd` does. (It never calculates an end bound, so it just keeps iterating until
// the trees run out.)
list = self
.db
.zs_range_iter(&orchard_subtrees, start_index..)
.collect();
}

// Make sure the amount of retrieved subtrees does not exceed the given limit.
#[cfg(debug_assertions)]
if let Some(limit) = limit {
assert!(list.len() <= limit.0.into());
}

// Check that we got the start subtree.
if list.get(&start_index).is_some() {
list
} else {
BTreeMap::new()
}
self.db.zs_range_iter(&orchard_subtrees, range).collect()
}

/// Get the orchard note commitment subtress for the finalized tip.
Expand Down
14 changes: 3 additions & 11 deletions zebra-state/src/service/non_finalized_state/chain.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -902,8 +902,7 @@ impl Chain {
.map(|(index, subtree)| subtree.with_index(*index))
}

/// Returns a list of Orchard [`NoteCommitmentSubtree`]s at or after `start_index`.
/// If `limit` is provided, the list is limited to `limit` entries.
/// Returns a list of Orchard [`NoteCommitmentSubtree`]s in the provided range.
///
/// Unlike the finalized state and `ReadRequest::OrchardSubtrees`, the returned subtrees
/// can start after `start_index`. These subtrees are continuous up to the tip.
Expand All @@ -913,17 +912,10 @@ impl Chain {
/// finalized updates.
pub fn orchard_subtrees_in_range(
&self,
start_index: NoteCommitmentSubtreeIndex,
limit: Option<NoteCommitmentSubtreeIndex>,
range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>> {
let limit = limit
.map(|limit| usize::from(limit.0))
.unwrap_or(usize::MAX);

// Since we're working in memory, it's ok to iterate through the whole range here.
self.orchard_subtrees
.range(start_index..)
.take(limit)
.range(range)
.map(|(index, subtree)| (*index, *subtree))
.collect()
}
Expand Down
78 changes: 15 additions & 63 deletions zebra-state/src/service/read/tree.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -134,78 +134,30 @@ where
pub fn orchard_subtrees<C>(
chain: Option<C>,
db: &ZebraDb,
start_index: NoteCommitmentSubtreeIndex,
limit: Option<NoteCommitmentSubtreeIndex>,
range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex> + Clone,
) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>>
where
C: AsRef<Chain>,
{
let mut db_list = db.orchard_subtree_list_by_index_for_rpc(start_index, limit);

if let Some(limit) = limit {
let subtrees_num = u16::try_from(db_list.len())
.expect("There can't be more than `u16::MAX` Orchard subtrees.");

// Return the subtrees if the amount of them reached the given limit.
if subtrees_num == limit.0 {
return db_list;
}

// If not, make sure the amount is below the limit.
debug_assert!(subtrees_num < limit.0);
}

// If there's no chain, then we have the complete list.
let Some(chain) = chain else {
return db_list;
};

// Unlike the other methods, this returns any trees in the range,
// even if there is no tree for start_index.
let fork_list = chain.as_ref().orchard_subtrees_in_range(start_index, limit);
use std::ops::Bound::*;

// If there's no subtrees in chain, then we have the complete list.
if fork_list.is_empty() {
return db_list;
let start_index = match range.start_bound().cloned() {
Included(start_index) | Excluded(start_index) => start_index,
Unbounded => panic!("range provided to orchard_subtrees() must have start bound"),
};

// Check for inconsistent trees in the fork.
for (fork_index, fork_subtree) in fork_list {
// If there's no matching index, just update the list of trees.
let Some(db_subtree) = db_list.get(&fork_index) else {
db_list.insert(fork_index, fork_subtree);

// Stop adding new subtrees once their amount reaches the given limit.
if let Some(limit) = limit {
let subtrees_num = u16::try_from(db_list.len())
.expect("There can't be more than `u16::MAX` Orchard subtrees.");

if subtrees_num == limit.0 {
break;
}
}

continue;
};

// We have an outdated chain fork, so skip this subtree and all remaining subtrees.
if &fork_subtree != db_subtree {
break;
let results = match chain.map(|chain| chain.as_ref().orchard_subtrees_in_range(range.clone())) {
Some(results) if results.contains_key(&start_index) => results,
Some(mut results) => {
results.append(&mut db.orchard_subtree_list_by_index_for_rpc(range));
results
}
None => db.orchard_subtree_list_by_index_for_rpc(range),
};

// Otherwise, the subtree is already in the list, so we don't need to add it.
}

// Make sure the amount of retrieved subtrees does not exceed the given limit.
#[cfg(debug_assertions)]
if let Some(limit) = limit {
assert!(db_list.len() <= limit.0.into());
}

// Check that we got the start subtree from the non-finalized or finalized state.
// (The non-finalized state doesn't do this check.)
if db_list.get(&start_index).is_some() {
db_list
// Check that we got the start subtree
if results.contains_key(&start_index) {
results
} else {
BTreeMap::new()
}
Expand Down

0 comments on commit a7edaa9

Please sign in to comment.