zcash_client_sqlite: Improve estimate for tree size at chain tip

After splitting recover progress out from scan progress, the
overestimation of the chain tip tree size became much more noticeable.
We now extrapolate from the most recent known "notes per block" rate.

Author: str4d

zcash_client_sqlite/src/lib.rs

@@ -74,7 +74,7 @@ use zip32::fingerprint::SeedFingerprint;
 
 use crate::{error::SqliteClientError, wallet::commitment_tree::SqliteShardStore};
 
-#[cfg(any(feature = "test-dependencies", not(feature = "orchard")))]
+#[cfg(any(test, feature = "test-dependencies", not(feature = "orchard")))]
 use zcash_protocol::PoolType;
 
 #[cfg(feature = "orchard")]

zcash_client_sqlite/src/wallet.rs

@@ -815,19 +815,21 @@ pub(crate) struct Progress {
 }
 
 pub(crate) trait ScanProgress {
-    fn sapling_scan_progress(
+    fn sapling_scan_progress<P: consensus::Parameters>(
         &self,
         conn: &rusqlite::Connection,
+        params: &P,
         birthday_height: BlockHeight,
         recover_until_height: Option<BlockHeight>,
         fully_scanned_height: BlockHeight,
         chain_tip_height: BlockHeight,
     ) -> Result<Progress, SqliteClientError>;
 
     #[cfg(feature = "orchard")]
-    fn orchard_scan_progress(
+    fn orchard_scan_progress<P: consensus::Parameters>(
         &self,
         conn: &rusqlite::Connection,
+        params: &P,
         birthday_height: BlockHeight,
         recover_until_height: Option<BlockHeight>,
         fully_scanned_height: BlockHeight,
@@ -844,6 +846,7 @@ fn subtree_scan_progress(
     table_prefix: &'static str,
     output_count_col: &'static str,
     shard_height: u8,
+    pool_activation_height: BlockHeight,
     birthday_height: BlockHeight,
     recover_until_height: Option<BlockHeight>,
     fully_scanned_height: BlockHeight,
@@ -864,6 +867,11 @@ fn subtree_scan_progress(
         FROM blocks
         WHERE height <= :start_height",
     ))?;
+    let mut stmt_end_tree_size_at = conn.prepare_cached(&format!(
+        "SELECT {table_prefix}_commitment_tree_size
+        FROM blocks
+        WHERE height = :height",
+    ))?;
 
     if fully_scanned_height == chain_tip_height {
         // Compute the total blocks scanned since the wallet birthday on either side of
@@ -956,16 +964,13 @@ fn subtree_scan_progress(
             })
             .transpose()?;
 
-        // We don't have complete information on how many outputs will exist in the shard at
-        // the chain tip without having scanned the chain tip block, so we overestimate by
-        // computing the maximum possible number of notes directly from the shard indices.
-        //
-        // TODO: it would be nice to be able to reliably have the size of the commitment tree
-        // at the chain tip without having to have scanned that block.
-        let (min_tree_size, max_tree_size) = conn
+        // In case we didn't have information about the tree size at the recover-until
+        // height, get the tree size from a nearby subtree. It's fine for this to be
+        // approximate; it just shifts the boundary between scan and recover progress.
+        let min_tree_size = conn
             .query_row(
                 &format!(
-                    "SELECT MIN(shard_index), MAX(shard_index)
+                    "SELECT MIN(shard_index)
                     FROM {table_prefix}_tree_shards
                     WHERE subtree_end_height > :start_height
                     OR subtree_end_height IS NULL",
@@ -975,13 +980,183 @@ fn subtree_scan_progress(
                     let min_tree_size = row
                         .get::<_, Option<u64>>(0)?
                         .map(|min_idx| min_idx << shard_height);
-                    let max_tree_size = row
-                        .get::<_, Option<u64>>(1)?
-                        .map(|max_idx| (max_idx + 1) << shard_height);
-                    Ok(min_tree_size.zip(max_tree_size))
+                    Ok(min_tree_size)
                 },
             )
-            .optional()?.flatten().unzip();
+            .optional()?.flatten();
+
+        // If we've scanned the block at the chain tip, we know how many notes are
+        // currently in the tree.
+        let tip_tree_size = match stmt_end_tree_size_at
+            .query_row(
+                named_params! {":height": u32::from(chain_tip_height)},
+                |row| row.get::<_, Option<u64>>(0),
+            )
+            .optional()?
+            .flatten()
+        {
+            Some(tree_size) => Some(tree_size),
+            None => {
+                // Estimate the size of the tree by linear extrapolation from available
+                // data closest to the chain tip.
+                //
+                // - If we have scanned blocks within the incomplete subtree, and we know
+                //   the tree size for the end of the most recent scanned range, then we
+                //   extrapolate from the start of the incomplete subtree:
+                //
+                //         subtree
+                //         /     \
+                //       /         \
+                //     /             \
+                //   /                 \
+                //   |<--------->|  |
+                //     | scanned |  tip
+                //           last_scanned
+                //
+                //
+                //             subtree
+                //             /     \
+                //           /         \
+                //         /             \
+                //       /                 \
+                //       |<------->|    |
+                //   |   scanned   |    tip
+                //             last_scanned
+                //
+                // - If we don't have scanned blocks within the incomplete subtree, or we
+                //   don't know the tree size, then we extrapolate from the block-width of
+                //   the last complete subtree.
+                //
+                // This avoids having a sharp discontinuity in the progress percentages
+                // shown to users, and gets more accurate the closer to the chain tip we
+                // have scanned.
+                //
+                // TODO: it would be nice to be able to reliably have the size of the
+                // commitment tree at the chain tip without having to have scanned that
+                // block.
+
+                // Get the last scanned height.
+                let last_scanned = conn
+                    .query_row(
+                        "SELECT block_range_end
+                        FROM scan_queue
+                        WHERE priority = :priority
+                        ORDER BY block_range_start DESC
+                        LIMIT 1",
+                        named_params![":priority": priority_code(&ScanPriority::Scanned)],
+                        |row| {
+                            // Scan ranges are end-exclusive.
+                            Ok(BlockHeight::from_u32(row.get(0)?) - 1)
+                        },
+                    )
+                    // `None` if we haven't scanned anything.
+                    .optional()?;
+
+                // Get the tree size at the last scanned height, if known.
+                let last_scanned_tree_size = last_scanned
+                    .map(|last_scanned| {
+                        stmt_end_tree_size_at
+                            .query_row(named_params! {":height": u32::from(last_scanned)}, |row| {
+                                row.get::<_, Option<u64>>(0)
+                            })
+                    })
+                    .transpose()?
+                    .flatten();
+
+                // Get the last completed subtree.
+                let last_completed_subtree = conn
+                    .query_row(
+                        &format!(
+                            "SELECT shard_index, subtree_end_height
+                            FROM {table_prefix}_tree_shards
+                            WHERE subtree_end_height IS NOT NULL
+                            ORDER BY shard_index DESC
+                            LIMIT 1"
+                        ),
+                        [],
+                        |row| Ok((row.get::<_, u64>(0)?, BlockHeight::from_u32(row.get(1)?))),
+                    )
+                    // `None` if we have no subtree roots yet.
+                    .optional()?;
+
+                if let Some((last_completed_subtree_index, last_completed_subtree_height)) =
+                    last_completed_subtree
+                {
+                    // If we know the tree size at the last scanned height, extrapolate.
+                    let tip_tree_size = last_scanned.zip(last_scanned_tree_size).and_then(
+                        |(last_scanned, last_scanned_tree_size)| {
+                            let scanned_notes = last_scanned_tree_size
+                                - (last_completed_subtree_index << shard_height);
+                            let scanned_range =
+                                u64::from(last_scanned - last_completed_subtree_height);
+                            let unscanned_range = u64::from(chain_tip_height - last_scanned);
+
+                            (scanned_notes * unscanned_range)
+                                .checked_div(scanned_range)
+                                .map(|extrapolated_unscanned_notes| {
+                                    last_scanned_tree_size + extrapolated_unscanned_notes
+                                })
+                        },
+                    );
+
+                    if let Some(tree_size) = tip_tree_size {
+                        Some(tree_size)
+                    } else if let Some(second_to_last_completed_subtree_height) =
+                        last_completed_subtree_index
+                            .checked_sub(1)
+                            .and_then(|subtree_index| {
+                                conn.query_row(
+                                    &format!(
+                                        "SELECT subtree_end_height
+                                        FROM {table_prefix}_tree_shards
+                                        WHERE shard_index = :shard_index"
+                                    ),
+                                    named_params! {":shard_index": subtree_index},
+                                    |row| {
+                                        Ok(row.get::<_, Option<_>>(0)?.map(BlockHeight::from_u32))
+                                    },
+                                )
+                                .transpose()
+                            })
+                            .transpose()?
+                    {
+                        let notes_in_complete_subtrees =
+                            (last_completed_subtree_index + 1) << shard_height;
+
+                        let subtree_notes = 1 << shard_height;
+                        let subtree_range = u64::from(
+                            last_completed_subtree_height - second_to_last_completed_subtree_height,
+                        );
+                        let unscanned_range =
+                            u64::from(chain_tip_height - last_completed_subtree_height);
+
+                        (subtree_notes * unscanned_range)
+                            .checked_div(subtree_range)
+                            .map(|extrapolated_incomplete_subtree_notes| {
+                                notes_in_complete_subtrees + extrapolated_incomplete_subtree_notes
+                            })
+                    } else {
+                        // There's only one completed subtree; its start height must
+                        // be the activation height for this shielded protocol.
+                        let scanned_notes = last_completed_subtree_index << shard_height;
+
+                        let subtree_range =
+                            u64::from(last_completed_subtree_height - pool_activation_height);
+                        let unscanned_range =
+                            u64::from(chain_tip_height - last_completed_subtree_height);
+
+                        (scanned_notes * unscanned_range)
+                            .checked_div(subtree_range)
+                            .map(|extrapolated_unscanned_notes| {
+                                scanned_notes + extrapolated_unscanned_notes
+                            })
+                    }
+                } else {
+                    // We don't have subtree information, so give up. We'll get it soon.
+                    None
+                }
+            }
+        };
 
         let recover = recovered_count
             .zip(recover_until_size)
@@ -1018,9 +1193,9 @@ fn subtree_scan_progress(
             recover_until_size
                 .unwrap_or(start_size)
                 .or(min_tree_size)
-                .zip(max_tree_size)
-                .map(|(min_tree_size, max_tree_size)| {
-                    Ratio::new(scanned_count.unwrap_or(0), max_tree_size - min_tree_size)
+                .zip(tip_tree_size)
+                .map(|(start_size, tip_tree_size)| {
+                    Ratio::new(scanned_count.unwrap_or(0), tip_tree_size - start_size)
                 })
         };
 
@@ -1029,10 +1204,11 @@ fn subtree_scan_progress(
 }
 
 impl ScanProgress for SubtreeScanProgress {
-    #[tracing::instrument(skip(conn))]
-    fn sapling_scan_progress(
+    #[tracing::instrument(skip(conn, params))]
+    fn sapling_scan_progress<P: consensus::Parameters>(
         &self,
         conn: &rusqlite::Connection,
+        params: &P,
         birthday_height: BlockHeight,
         recover_until_height: Option<BlockHeight>,
         fully_scanned_height: BlockHeight,
@@ -1043,6 +1219,9 @@ impl ScanProgress for SubtreeScanProgress {
             SAPLING_TABLES_PREFIX,
             "sapling_output_count",
             SAPLING_SHARD_HEIGHT,
+            params
+                .activation_height(NetworkUpgrade::Sapling)
+                .expect("Sapling activation height must be available."),
             birthday_height,
             recover_until_height,
             fully_scanned_height,
@@ -1051,10 +1230,11 @@ impl ScanProgress for SubtreeScanProgress {
     }
 
     #[cfg(feature = "orchard")]
-    #[tracing::instrument(skip(conn))]
-    fn orchard_scan_progress(
+    #[tracing::instrument(skip(conn, params))]
+    fn orchard_scan_progress<P: consensus::Parameters>(
         &self,
         conn: &rusqlite::Connection,
+        params: &P,
         birthday_height: BlockHeight,
         recover_until_height: Option<BlockHeight>,
         fully_scanned_height: BlockHeight,
@@ -1065,6 +1245,9 @@ impl ScanProgress for SubtreeScanProgress {
             ORCHARD_TABLES_PREFIX,
             "orchard_action_count",
             ORCHARD_SHARD_HEIGHT,
+            params
+                .activation_height(NetworkUpgrade::Nu5)
+                .expect("NU5 activation height must be available."),
             birthday_height,
             recover_until_height,
             fully_scanned_height,
@@ -1104,6 +1287,7 @@ pub(crate) fn get_wallet_summary<P: consensus::Parameters>(
 
     let sapling_progress = progress.sapling_scan_progress(
         tx,
+        params,
         birthday_height,
         recover_until_height,
         fully_scanned_height,
@@ -1113,6 +1297,7 @@ pub(crate) fn get_wallet_summary<P: consensus::Parameters>(
     #[cfg(feature = "orchard")]
     let orchard_progress = progress.orchard_scan_progress(
         tx,
+        params,
         birthday_height,
         recover_until_height,
         fully_scanned_height,

zcash_client_sqlite/src/wallet/scanning.rs

@@ -570,7 +570,7 @@ pub(crate) mod tests {
             pool::ShieldedPoolTester, sapling::SaplingPoolTester, AddressType, FakeCompactOutput,
             InitialChainState, TestBuilder, TestState,
         },
-        AccountBirthday, Ratio, WalletRead, WalletWrite, SAPLING_SHARD_HEIGHT,
+        AccountBirthday, Ratio, WalletRead, WalletWrite,
     };
     use zcash_primitives::{
         block::BlockHash,
@@ -1351,14 +1351,14 @@ pub(crate) mod tests {
         );
 
         // Progress denominator depends on which pools are enabled (which changes the
-        // initial tree states). Here we compute the denominator based upon the fact that
-        // the trees are the same size at present.
-        let expected_denom = (1 << SAPLING_SHARD_HEIGHT) * 2 - frontier_tree_size;
+        // initial tree states), and is extrapolated from the scanned range.
+        let expected_denom: u64 = 2628;
         #[cfg(feature = "orchard")]
         let expected_denom = expected_denom * 2;
+        let expected_denom = expected_denom + 1;
         assert_eq!(
             summary.and_then(|s| s.scan_progress()),
-            Some(Ratio::new(1, u64::from(expected_denom)))
+            Some(Ratio::new(1, expected_denom))
         );
 
         // Now simulate shutting down, and then restarting 70 blocks later, after a shard
@@ -1430,11 +1430,14 @@ pub(crate) mod tests {
         // We've crossed a subtree boundary, but only in one pool. We still only have one scanned
         // note but in the pool where we crossed the subtree boundary we have two shards worth of
         // notes to scan.
+        let expected_denom: u64 = 2823;
+        #[cfg(feature = "orchard")]
+        let expected_denom = expected_denom + 11_794;
         let expected_denom = expected_denom + (1 << 16);
         let summary = st.get_wallet_summary(1);
         assert_eq!(
             summary.and_then(|s| s.scan_progress()),
-            Some(Ratio::new(1, u64::from(expected_denom)))
+            Some(Ratio::new(1, expected_denom))
         );
     }