zcash_client_backend: Correctly track heap usage of batch items

As of zcash/librustzcash#633, `SaplingDomain::IncomingViewingKey` now
allocates memory internally, and this memory persists as long as the
`BatchRunner` is alive. Now that we have decoupled the measurement of
heap usage for batch tasks from their internals, we can add bounds to
all of the generic parameters of `Batch` to enable correctly measuring
their actual heap usage.

We also add `DynamicUsage` impls for a bunch of `zcash_primitives` types
that will be used with `BatchRunner` (or its equivalent implementation
in `zcashd`) by callers.

Cargo.toml

@@ -22,4 +22,4 @@ codegen-units = 1
 zcash_encoding = { path = "components/zcash_encoding" }
 zcash_note_encryption = { path = "components/zcash_note_encryption" }
 orchard = { git = "https://github.com/zcash/orchard.git", rev = "33f1c1141e50adb68715f3359bd75378b4756cca" }
-group = { git = "https://github.com/zkcrypto/group.git", rev = "a7f3ceb2373e9fe536996f7b4d55c797f3e667f0" }
+group = { git = "https://github.com/zkcrypto/group.git", rev = "f61e3e420ed1220c8f1f80988f8c6c5e202d8715" }

zcash_client_backend/src/scan.rs

@@ -217,33 +217,33 @@ pub(crate) struct Batch<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOT
     repliers: Vec<OutputReplier<A, D>>,
 }
 
-fn base_vec_usage<T>(c: &Vec<T>) -> usize {
-    c.capacity() * mem::size_of::<T>()
-}
-
 impl<A, D, Output> DynamicUsage for Batch<A, D, Output>
 where
-    D: BatchDomain,
-    Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>,
+    A: DynamicUsage,
+    D: BatchDomain + DynamicUsage,
+    D::IncomingViewingKey: DynamicUsage,
+    Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + DynamicUsage,
 {
     fn dynamic_usage(&self) -> usize {
-        // We don't have a `DynamicUsage` bound on `A`, `D::IncomingViewingKey`, `D`, or
-        // `Output`, and we can't use newtypes because the batch decryption API takes
-        // slices. But we know that we don't allocate memory inside either of these, so we
-        // just compute the size directly.
-        base_vec_usage(&self.tags)
-            + base_vec_usage(&self.ivks)
-            + base_vec_usage(&self.outputs)
+        self.tags.dynamic_usage()
+            + self.ivks.dynamic_usage()
+            + self.outputs.dynamic_usage()
             + self.repliers.dynamic_usage()
     }
 
     fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
-        let base_usage =
-            base_vec_usage(&self.tags) + base_vec_usage(&self.ivks) + base_vec_usage(&self.outputs);
-        let bounds = self.repliers.dynamic_usage_bounds();
+        let (tags_lower, tags_upper) = self.tags.dynamic_usage_bounds();
+        let (ivks_lower, ivks_upper) = self.ivks.dynamic_usage_bounds();
+        let (outputs_lower, outputs_upper) = self.outputs.dynamic_usage_bounds();
+        let (repliers_lower, repliers_upper) = self.repliers.dynamic_usage_bounds();
+
         (
-            base_usage + bounds.0,
-            bounds.1.map(|upper| base_usage + upper),
+            tags_lower + ivks_lower + outputs_lower + repliers_lower,
+            tags_upper
+                .zip(ivks_upper)
+                .zip(outputs_upper)
+                .zip(repliers_upper)
+                .map(|(((a, b), c), d)| a + b + c + d),
         )
     }
 }
@@ -373,8 +373,10 @@ where
 
 impl<A, D, Output, T> DynamicUsage for BatchRunner<A, D, Output, T>
 where
-    D: BatchDomain,
-    Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>,
+    A: DynamicUsage,
+    D: BatchDomain + DynamicUsage,
+    D::IncomingViewingKey: DynamicUsage,
+    Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + DynamicUsage,
     T: Tasks<Batch<A, D, Output>> + DynamicUsage,
 {
     fn dynamic_usage(&self) -> usize {

zcash_primitives/CHANGELOG.md

@@ -36,6 +36,20 @@ and this library adheres to Rust's notion of
   - `JSDescription::net_value`
 - Added in `zcash_primitives::transaction::components::transparent`
   - `Bundle::value_balance`
+- Implementations of `memuse::DynamicUsage` for the following types:
+  - `zcash_primitives::block::BlockHash`
+  - `zcash_primitives::consensus`:
+    - `BlockHeight`
+    - `MainNetwork`, `TestNetwork`, `Network`
+    - `NetworkUpgrade`, `BranchId`
+  - `zcash_primitives::sapling`:
+    - `keys::Scope`
+    - `note_encryption::SaplingDomain`
+  - `zcash_primitives::transaction`:
+    - `TxId`
+    - `components::sapling::CompactOutputDescription`
+    - `components::sapling::{OutputDescription, OutputDescriptionV5}`
+  - `zcash_primitives::zip32::AccountId`
 
 ### Changed
 - Migrated to `group 0.13`.

zcash_primitives/Cargo.toml

@@ -30,13 +30,13 @@ chacha20poly1305 = "0.10"
 equihash = { version = "0.2", path = "../components/equihash" }
 ff = "0.12"
 fpe = "0.5"
-group = "0.12"
+group = { version = "0.12", features = ["wnaf-memuse"] }
 hdwallet = { version = "0.3.1", optional = true }
 hex = "0.4"
 incrementalmerkletree = "0.3"
 jubjub = "0.9"
 lazy_static = "1"
-memuse = "0.2"
+memuse = "0.2.1"
 nonempty = "0.7"
 orchard = "0.2"
 proptest = { version = "1.0.0", optional = true }

zcash_primitives/src/block.rs

@@ -1,6 +1,7 @@
 //! Structs and methods for handling Zcash block headers.
 
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+use memuse::DynamicUsage;
 use sha2::{Digest, Sha256};
 use std::fmt;
 use std::io::{self, Read, Write};
@@ -12,6 +13,8 @@ pub use equihash;
 #[derive(Clone, Copy, PartialEq, Eq, Hash)]
 pub struct BlockHash(pub [u8; 32]);
 
+memuse::impl_no_dynamic_usage!(BlockHash);
+
 impl fmt::Debug for BlockHash {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         // The (byte-flipped) hex string is more useful than the raw bytes, because we can

zcash_primitives/src/consensus.rs

@@ -1,5 +1,6 @@
 //! Consensus logic and parameters.
 
+use memuse::DynamicUsage;
 use std::cmp::{Ord, Ordering};
 use std::convert::TryFrom;
 use std::fmt;
@@ -14,6 +15,8 @@ use crate::constants;
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub struct BlockHeight(u32);
 
+memuse::impl_no_dynamic_usage!(BlockHeight);
+
 pub const H0: BlockHeight = BlockHeight(0);
 
 impl BlockHeight {
@@ -190,6 +193,8 @@ pub trait Parameters: Clone {
 #[derive(PartialEq, Eq, Copy, Clone, Debug)]
 pub struct MainNetwork;
 
+memuse::impl_no_dynamic_usage!(MainNetwork);
+
 pub const MAIN_NETWORK: MainNetwork = MainNetwork;
 
 impl Parameters for MainNetwork {
@@ -239,6 +244,8 @@ impl Parameters for MainNetwork {
 #[derive(PartialEq, Eq, Copy, Clone, Debug)]
 pub struct TestNetwork;
 
+memuse::impl_no_dynamic_usage!(TestNetwork);
+
 pub const TEST_NETWORK: TestNetwork = TestNetwork;
 
 impl Parameters for TestNetwork {
@@ -290,6 +297,8 @@ pub enum Network {
     TestNetwork,
 }
 
+memuse::impl_no_dynamic_usage!(Network);
+
 impl Parameters for Network {
     fn activation_height(&self, nu: NetworkUpgrade) -> Option<BlockHeight> {
         match self {
@@ -387,6 +396,8 @@ pub enum NetworkUpgrade {
     ZFuture,
 }
 
+memuse::impl_no_dynamic_usage!(NetworkUpgrade);
+
 impl fmt::Display for NetworkUpgrade {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
@@ -467,6 +478,8 @@ pub enum BranchId {
     ZFuture,
 }
 
+memuse::impl_no_dynamic_usage!(BranchId);
+
 impl TryFrom<u32> for BranchId {
     type Error = &'static str;
 

zcash_primitives/src/sapling/keys.rs

@@ -13,6 +13,7 @@ use crate::{
 };
 use ff::PrimeField;
 use group::{Group, GroupEncoding};
+use memuse::DynamicUsage;
 use subtle::CtOption;
 
 use super::{NullifierDerivingKey, PaymentAddress, ProofGenerationKey, SaplingIvk, ViewingKey};
@@ -201,6 +202,8 @@ pub enum Scope {
     Internal,
 }
 
+memuse::impl_no_dynamic_usage!(Scope);
+
 /// A Sapling key that provides the capability to view incoming and outgoing transactions.
 ///
 /// This key is useful anywhere you need to maintain accurate balance, but do not want the

zcash_primitives/src/sapling/note_encryption.rs

@@ -7,6 +7,7 @@ use byteorder::{LittleEndian, WriteBytesExt};
 use ff::PrimeField;
 use group::{cofactor::CofactorGroup, GroupEncoding, WnafBase, WnafScalar};
 use jubjub::{AffinePoint, ExtendedPoint};
+use memuse::DynamicUsage;
 use rand_core::RngCore;
 
 use zcash_note_encryption::{
@@ -39,6 +40,16 @@ type PreparedScalar = WnafScalar<jubjub::Scalar, PREPARED_WINDOW_SIZE>;
 #[derive(Clone, Debug)]
 pub struct PreparedIncomingViewingKey(PreparedScalar);
 
+impl DynamicUsage for PreparedIncomingViewingKey {
+    fn dynamic_usage(&self) -> usize {
+        self.0.dynamic_usage()
+    }
+
+    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
+        self.0.dynamic_usage_bounds()
+    }
+}
+
 impl PreparedIncomingViewingKey {
     /// Performs the necessary precomputations to use a `SaplingIvk` for note decryption.
     pub fn new(ivk: &SaplingIvk) -> Self {
@@ -146,6 +157,21 @@ pub struct SaplingDomain<P: consensus::Parameters> {
     height: BlockHeight,
 }
 
+impl<P: consensus::Parameters + DynamicUsage> DynamicUsage for SaplingDomain<P> {
+    fn dynamic_usage(&self) -> usize {
+        self.params.dynamic_usage() + self.height.dynamic_usage()
+    }
+
+    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
+        let (params_lower, params_upper) = self.params.dynamic_usage_bounds();
+        let (height_lower, height_upper) = self.height.dynamic_usage_bounds();
+        (
+            params_lower + height_lower,
+            params_upper.zip(height_upper).map(|(a, b)| a + b),
+        )
+    }
+}
+
 impl<P: consensus::Parameters> SaplingDomain<P> {
     pub fn for_height(params: P, height: BlockHeight) -> Self {
         Self { params, height }

zcash_primitives/src/transaction/components/amount.rs

@@ -2,6 +2,7 @@ use std::convert::TryFrom;
 use std::iter::Sum;
 use std::ops::{Add, AddAssign, Neg, Sub, SubAssign};
 
+use memuse::DynamicUsage;
 use orchard::value as orchard;
 
 pub const COIN: i64 = 1_0000_0000;
@@ -23,17 +24,7 @@ pub const DEFAULT_FEE: Amount = Amount(1000);
 #[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Eq, Ord)]
 pub struct Amount(i64);
 
-impl memuse::DynamicUsage for Amount {
-    #[inline(always)]
-    fn dynamic_usage(&self) -> usize {
-        0
-    }
-
-    #[inline(always)]
-    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
-        (0, Some(0))
-    }
-}
+memuse::impl_no_dynamic_usage!(Amount);
 
 impl Amount {
     /// Returns a zero-valued Amount.

zcash_primitives/src/transaction/components/sapling.rs

@@ -2,6 +2,8 @@ use core::fmt::Debug;
 
 use ff::PrimeField;
 use group::GroupEncoding;
+use memuse::DynamicUsage;
+
 use std::io::{self, Read, Write};
 
 use zcash_note_encryption::{
@@ -261,6 +263,16 @@ pub struct OutputDescription<Proof> {
     pub zkproof: Proof,
 }
 
+impl<Proof: DynamicUsage> DynamicUsage for OutputDescription<Proof> {
+    fn dynamic_usage(&self) -> usize {
+        self.zkproof.dynamic_usage()
+    }
+
+    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
+        self.zkproof.dynamic_usage_bounds()
+    }
+}
+
 impl<P: consensus::Parameters, A> ShieldedOutput<SaplingDomain<P>, ENC_CIPHERTEXT_SIZE>
     for OutputDescription<A>
 {
@@ -348,6 +360,8 @@ pub struct OutputDescriptionV5 {
     pub out_ciphertext: [u8; 80],
 }
 
+memuse::impl_no_dynamic_usage!(OutputDescriptionV5);
+
 impl OutputDescriptionV5 {
     pub fn read<R: Read>(mut reader: &mut R) -> io::Result<Self> {
         let cv = read_point(&mut reader, "cv")?;
@@ -395,6 +409,8 @@ pub struct CompactOutputDescription {
     pub enc_ciphertext: [u8; COMPACT_NOTE_SIZE],
 }
 
+memuse::impl_no_dynamic_usage!(CompactOutputDescription);
+
 impl<A> From<OutputDescription<A>> for CompactOutputDescription {
     fn from(out: OutputDescription<A>) -> CompactOutputDescription {
         CompactOutputDescription {

zcash_primitives/src/transaction/mod.rs

@@ -13,6 +13,7 @@ mod tests;
 use blake2b_simd::Hash as Blake2bHash;
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use ff::PrimeField;
+use memuse::DynamicUsage;
 use std::convert::TryFrom;
 use std::fmt;
 use std::fmt::Debug;
@@ -65,6 +66,8 @@ const ZFUTURE_TX_VERSION: u32 = 0x0000FFFF;
 #[derive(Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Hash)]
 pub struct TxId([u8; 32]);
 
+memuse::impl_no_dynamic_usage!(TxId);
+
 impl fmt::Debug for TxId {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         // The (byte-flipped) hex string is more useful than the raw bytes, because we can

zcash_primitives/src/zip32.rs

@@ -6,6 +6,7 @@ use aes::Aes256;
 use blake2b_simd::Params as Blake2bParams;
 use byteorder::{ByteOrder, LittleEndian, ReadBytesExt, WriteBytesExt};
 use fpe::ff1::{BinaryNumeralString, FF1};
+use memuse::DynamicUsage;
 use std::ops::AddAssign;
 use subtle::{Choice, ConditionallySelectable};
 
@@ -31,6 +32,8 @@ pub const ZIP32_SAPLING_INT_PERSONALIZATION: &[u8; 16] = b"Zcash_SaplingInt";
 #[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct AccountId(u32);
 
+memuse::impl_no_dynamic_usage!(AccountId);
+
 impl From<u32> for AccountId {
     fn from(id: u32) -> Self {
         Self(id)