From ef50f8b4819fed43408c3c90ab5d3a18654dd258 Mon Sep 17 00:00:00 2001 From: teor Date: Wed, 13 Apr 2022 13:24:10 +1000 Subject: [PATCH] Refactor and fix documentation of the balance query code --- .../finalized_state/zebra_db/transparent.rs | 4 +-- .../src/service/non_finalized_state/chain.rs | 35 +++++++++++++------ zebra-state/src/service/read.rs | 15 ++++---- 3 files changed, 35 insertions(+), 19 deletions(-) diff --git a/zebra-state/src/service/finalized_state/zebra_db/transparent.rs b/zebra-state/src/service/finalized_state/zebra_db/transparent.rs index 85bd570744a..93b76618339 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/transparent.rs +++ b/zebra-state/src/service/finalized_state/zebra_db/transparent.rs @@ -255,8 +255,8 @@ impl ZebraDb { /// /// Callers should apply the non-finalized balance change for `addresses` to the returned balance. /// - /// The total balance will only be correct if this partial chain matches the finalized state. - /// Specifically, the root of this partial chain must be a child block of the finalized tip. + /// The total balance will only be correct if the non-finalized chain matches the finalized state. + /// Specifically, the root of the partial non-finalized chain must be a child block of the finalized tip. pub fn partial_finalized_transparent_balance( &self, addresses: &HashSet, diff --git a/zebra-state/src/service/non_finalized_state/chain.rs b/zebra-state/src/service/non_finalized_state/chain.rs index fc40a6cf291..aff43ec9085 100644 --- a/zebra-state/src/service/non_finalized_state/chain.rs +++ b/zebra-state/src/service/non_finalized_state/chain.rs @@ -12,7 +12,7 @@ use mset::MultiSet; use tracing::instrument; use zebra_chain::{ - amount::{self, Amount, NegativeAllowed, NonNegative}, + amount::{Amount, NegativeAllowed, NonNegative}, block, history_tree::HistoryTree, orchard, @@ -493,6 +493,26 @@ impl Chain { // Address index queries + /// Returns the transparent transfers for `addresses` in this non-finalized chain. + /// + /// If none of the addresses have an address index, returns an empty iterator. + /// + /// # Correctness + /// + /// Callers should apply the returned indexes to the corresponding finalized state indexes. + /// + /// The combined result will only be correct if the chains match. + /// The exact type of match varies by query. + pub fn partial_transparent_indexes<'a>( + &'a self, + addresses: &'a HashSet, + ) -> impl Iterator { + addresses + .iter() + .copied() + .flat_map(|address| self.partial_transparent_transfers.get(&address)) + } + /// Returns the transparent balance change for `addresses` in this non-finalized chain. /// /// If the balance doesn't change for any of the addresses, returns zero. @@ -507,16 +527,9 @@ impl Chain { &self, addresses: &HashSet, ) -> Amount { - let balance_change: amount::Result> = self - .partial_transparent_transfers - .iter() - .filter_map(|(address, transfers)| { - if addresses.contains(address) { - Some(transfers.balance()) - } else { - None - } - }) + let balance_change: Result, _> = self + .partial_transparent_indexes(addresses) + .map(|transfers| transfers.balance()) .sum(); balance_change.expect( diff --git a/zebra-state/src/service/read.rs b/zebra-state/src/service/read.rs index 3c9f063b657..b7fe623ea6d 100644 --- a/zebra-state/src/service/read.rs +++ b/zebra-state/src/service/read.rs @@ -21,12 +21,12 @@ use crate::{ #[cfg(test)] mod tests; -/// If the transparent address balance queries are interrupted by a new finalized block, +/// If the transparent address index queries are interrupted by a new finalized block, /// retry this many times. /// /// Once we're at the tip, we expect up to 2 blocks to arrive at the same time. /// If any more arrive, the client should wait until we're synchronised with our peers. -const FINALIZED_ADDRESS_BALANCE_RETRIES: usize = 3; +const FINALIZED_ADDRESS_INDEX_RETRIES: usize = 3; /// Returns the [`Block`] with [`block::Hash`](zebra_chain::block::Hash) or /// [`Height`](zebra_chain::block::Height), @@ -95,7 +95,7 @@ pub(crate) fn transparent_balance( let mut balance_result = finalized_transparent_balance(db, &addresses); // Retry the finalized balance query if it was interruped by a finalizing block - for _ in 0..FINALIZED_ADDRESS_BALANCE_RETRIES { + for _ in 0..FINALIZED_ADDRESS_INDEX_RETRIES { if balance_result.is_ok() { break; } @@ -121,7 +121,9 @@ pub(crate) fn transparent_balance( /// Returns the total transparent balance for `addresses` in the finalized chain, /// and the finalized tip height the balances were queried at. /// -/// If the addresses do not exist in the non-finalized `chain` or finalized `db`, returns zero. +/// If the addresses do not exist in the finalized `db`, returns zero. +// +// TODO: turn the return type into a struct? fn finalized_transparent_balance( db: &ZebraDb, addresses: &HashSet, @@ -147,9 +149,10 @@ fn finalized_transparent_balance( Ok((finalized_balance, finalized_tip)) } -/// Returns the total transparent balance for the supplied [`transparent::Address`]es. +/// Returns the total transparent balance change for `addresses` in the non-finalized chain, +/// matching the balance for the `finalized_tip`. /// -/// If the addresses do not exist in the non-finalized `chain` or finalized `db`, returns zero. +/// If the addresses do not exist in the non-finalized `chain`, returns zero. fn chain_transparent_balance_change( mut chain: Arc, addresses: &HashSet,