Refactor and fix documentation of the balance query code

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<transparent::Address>,

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<transparent::Address>,
+    ) -> impl Iterator<Item = &TransparentTransfers> {
+        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<transparent::Address>,
     ) -> Amount<NegativeAllowed> {
-        let balance_change: amount::Result<Amount<NegativeAllowed>> = self
-            .partial_transparent_transfers
-            .iter()
-            .filter_map(|(address, transfers)| {
-                if addresses.contains(address) {
-                    Some(transfers.balance())
-                } else {
-                    None
-                }
-            })
+        let balance_change: Result<Amount<NegativeAllowed>, _> = self
+            .partial_transparent_indexes(addresses)
+            .map(|transfers| transfers.balance())
             .sum();
 
         balance_change.expect(

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<transparent::Address>,
@@ -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<Chain>,
     addresses: &HashSet<transparent::Address>,