From 050b2f231ef59064843ea7f940da423a55556fa2 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@electriccoin.co>
Date: Sat, 15 Oct 2022 23:29:06 +0000
Subject: [PATCH] Add `memuse::DynamicUsage` impls for types used for batch
 scanning

---
 CHANGELOG.md           |  3 +++
 Cargo.toml             |  4 ++--
 src/keys.rs            | 10 ++++++++++
 src/note/nullifier.rs  |  4 ++++
 src/note_encryption.rs | 10 ++++++++++
 src/spec.rs            | 11 +++++++++++
 6 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6deb97bc..2b164e00 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,9 @@ and this project adheres to Rust's notion of
 - `orchard::note`:
   - `RandomSeed`
   - `Note::{from_parts, rseed}`
+  - `impl memuse::DynamicUsage for Nullifier`
+- `orchard::note_encryption`:
+  - `impl memuse::DynamicUsage for OrchardDomain`
 - `orchard::builder::SpendInfo::new`
 - `orchard::circuit::Circuit::from_action_context`
 
diff --git a/Cargo.toml b/Cargo.toml
index 0c4ffeb5..87490630 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -28,12 +28,12 @@ bitvec = "1"
 blake2b_simd = "1"
 ff = "0.12"
 fpe = "0.5"
-group = "0.12.1"
+group = { version = "0.12.1", features = ["wnaf-memuse"] }
 halo2_gadgets = "0.2"
 halo2_proofs = "0.2"
 hex = "0.4"
 lazy_static = "1"
-memuse = { version = "0.2", features = ["nonempty"] }
+memuse = { version = "0.2.1", features = ["nonempty"] }
 pasta_curves = "0.4"
 proptest = { version = "1.0.0", optional = true }
 rand = "0.8"
diff --git a/src/keys.rs b/src/keys.rs
index 24145c23..8224a1ca 100644
--- a/src/keys.rs
+++ b/src/keys.rs
@@ -710,6 +710,16 @@ impl IncomingViewingKey {
 #[derive(Clone, Debug)]
 pub struct PreparedIncomingViewingKey(PreparedNonZeroScalar);
 
+impl memuse::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 an `IncomingViewingKey` for note
     /// decryption.
diff --git a/src/note/nullifier.rs b/src/note/nullifier.rs
index 74efe5cd..591ea655 100644
--- a/src/note/nullifier.rs
+++ b/src/note/nullifier.rs
@@ -1,5 +1,6 @@
 use group::{ff::PrimeField, Group};
 use halo2_proofs::arithmetic::CurveExt;
+use memuse::DynamicUsage;
 use pasta_curves::pallas;
 use rand::RngCore;
 use subtle::CtOption;
@@ -14,6 +15,9 @@ use crate::{
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
 pub struct Nullifier(pub(crate) pallas::Base);
 
+// We know that `pallas::Base` doesn't allocate internally.
+memuse::impl_no_dynamic_usage!(Nullifier);
+
 impl Nullifier {
     /// Generates a dummy nullifier for use as $\rho$ in dummy spent notes.
     ///
diff --git a/src/note_encryption.rs b/src/note_encryption.rs
index 762cc560..8ad06c8a 100644
--- a/src/note_encryption.rs
+++ b/src/note_encryption.rs
@@ -85,6 +85,16 @@ pub struct OrchardDomain {
     rho: Nullifier,
 }
 
+impl memuse::DynamicUsage for OrchardDomain {
+    fn dynamic_usage(&self) -> usize {
+        self.rho.dynamic_usage()
+    }
+
+    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
+        self.rho.dynamic_usage_bounds()
+    }
+}
+
 impl OrchardDomain {
     /// Constructs a domain that can be used to trial-decrypt this action's output note.
     pub fn for_action<T>(act: &Action<T>) -> Self {
diff --git a/src/spec.rs b/src/spec.rs
index ab5413ae..3af8b97d 100644
--- a/src/spec.rs
+++ b/src/spec.rs
@@ -7,6 +7,7 @@ use ff::{Field, PrimeField, PrimeFieldBits};
 use group::{Curve, Group, GroupEncoding, WnafBase, WnafScalar};
 use halo2_gadgets::{poseidon::primitives as poseidon, sinsemilla::primitives as sinsemilla};
 use halo2_proofs::arithmetic::{CurveAffine, CurveExt, FieldExt};
+use memuse::DynamicUsage;
 use pasta_curves::pallas;
 use subtle::{ConditionallySelectable, CtOption};
 
@@ -153,6 +154,16 @@ impl PreparedNonIdentityBase {
 #[derive(Clone, Debug)]
 pub(crate) struct PreparedNonZeroScalar(WnafScalar<pallas::Scalar, PREPARED_WINDOW_SIZE>);
 
+impl DynamicUsage for PreparedNonZeroScalar {
+    fn dynamic_usage(&self) -> usize {
+        self.0.dynamic_usage()
+    }
+
+    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
+        self.0.dynamic_usage_bounds()
+    }
+}
+
 impl PreparedNonZeroScalar {
     pub(crate) fn new(scalar: &NonZeroPallasScalar) -> Self {
         PreparedNonZeroScalar(WnafScalar::new(scalar))