From 7b5010369a52ea8a3096ad63687b1f1602399c27 Mon Sep 17 00:00:00 2001
From: moana <moana@dusk.network>
Date: Tue, 13 Aug 2024 16:47:59 +0200
Subject: [PATCH] wallet-core: Add function to sum the balance for phoenix
 notes

Resolves #2117
---
 wallet-core/Cargo.toml                        |  2 +
 wallet-core/src/lib.rs                        |  2 +-
 wallet-core/src/{keys.rs => profile.rs}       | 15 ++++-
 .../tests/{keys_derive.rs => profile.rs}      | 62 ++++++++++++++++++-
 4 files changed, 78 insertions(+), 3 deletions(-)
 rename wallet-core/src/{keys.rs => profile.rs} (86%)
 rename wallet-core/tests/{keys_derive.rs => profile.rs} (60%)

diff --git a/wallet-core/Cargo.toml b/wallet-core/Cargo.toml
index f4e3630552..74b5a46a9c 100644
--- a/wallet-core/Cargo.toml
+++ b/wallet-core/Cargo.toml
@@ -12,5 +12,7 @@ rand_chacha = { version = "0.3", default-features = false }
 sha2 = { version = "0.10", default-features = false }
 
 [dev-dependencies]
+rand = "0.8"
+ff = { version = "0.13", default-features = false }
 
 [features]
diff --git a/wallet-core/src/lib.rs b/wallet-core/src/lib.rs
index 85f0d38093..8226338d56 100644
--- a/wallet-core/src/lib.rs
+++ b/wallet-core/src/lib.rs
@@ -11,7 +11,7 @@
 #![deny(rustdoc::broken_intra_doc_links)]
 #![deny(clippy::pedantic)]
 
-pub mod keys;
+pub mod profile;
 
 /// Length of the seed of the generated rng.
 pub const RNG_SEED: usize = 64;
diff --git a/wallet-core/src/keys.rs b/wallet-core/src/profile.rs
similarity index 86%
rename from wallet-core/src/keys.rs
rename to wallet-core/src/profile.rs
index 74eae1a81a..012d44ded1 100644
--- a/wallet-core/src/keys.rs
+++ b/wallet-core/src/profile.rs
@@ -12,7 +12,7 @@ use sha2::{Digest, Sha256};
 use execution_core::{
     signatures::bls::SecretKey as BlsSecretKey,
     transfer::phoenix::{
-        PublicKey as PhoenixPublicKey, SecretKey as PhoenixSecretKey,
+        Note, PublicKey as PhoenixPublicKey, SecretKey as PhoenixSecretKey,
         ViewKey as PhoenixViewKey,
     },
 };
@@ -20,6 +20,19 @@ use zeroize::Zeroize;
 
 use crate::RNG_SEED;
 
+/// Calculate the sum for all the given [`Note`]s that belong to the given
+/// [`PhoenixViewKey`].
+///
+/// This function assumes that all notes are owned by the given view-key.
+pub fn phoenix_balance(
+    pvk: &PhoenixViewKey,
+    notes: impl IntoIterator<Item = Note>,
+) -> u64 {
+    notes.into_iter().fold(0, |acc, note| {
+        acc + note.value(Some(pvk)).unwrap_or_default()
+    })
+}
+
 /// Generates a [`BlsSecretKey`] from a seed and index.
 ///
 /// The randomness is generated using [`rng_with_index`].
diff --git a/wallet-core/tests/keys_derive.rs b/wallet-core/tests/profile.rs
similarity index 60%
rename from wallet-core/tests/keys_derive.rs
rename to wallet-core/tests/profile.rs
index cf0e1fbcc5..c3fa4d4a9d 100644
--- a/wallet-core/tests/keys_derive.rs
+++ b/wallet-core/tests/profile.rs
@@ -5,13 +5,73 @@
 // Copyright (c) DUSK NETWORK. All rights reserved.
 
 use dusk_bytes::Serializable;
-use wallet_core::keys::{
+use ff::Field;
+use rand::rngs::StdRng;
+use rand::SeedableRng;
+
+use execution_core::{
+    transfer::phoenix::{
+        Note, PublicKey as PhoenixPublicKey, SecretKey as PhoenixSecretKey,
+    },
+    JubJubScalar,
+};
+
+use wallet_core::profile::{
     derive_bls_sk, derive_phoenix_pk, derive_phoenix_sk, derive_phoenix_vk,
+    phoenix_balance,
 };
 
 const SEED: [u8; 64] = [0; 64];
 const INDEX: u8 = 42;
 
+#[test]
+fn test_balance() {
+    let mut rng = StdRng::seed_from_u64(0xdab);
+
+    let owner_sk = PhoenixSecretKey::random(&mut rng);
+    let owner_pk = PhoenixPublicKey::from(&owner_sk);
+    let sender_pk = PhoenixPublicKey::from(&PhoenixSecretKey::random(&mut rng));
+
+    let mut notes = Vec::new();
+
+    // the sum of these notes should be 5 * 11 = 55
+    let expected_balance = 55;
+    for value in 0..=10 {
+        let value_blinder = JubJubScalar::random(&mut rng);
+        let sender_blinder = [
+            JubJubScalar::random(&mut rng),
+            JubJubScalar::random(&mut rng),
+        ];
+
+        // we want to test with a mix of transparent and obfuscated notes so we
+        // make every 10th note transparent
+        let note = if value % 10 == 0 {
+            Note::transparent(
+                &mut rng,
+                &sender_pk,
+                &owner_pk,
+                value,
+                sender_blinder,
+            )
+        } else {
+            Note::obfuscated(
+                &mut rng,
+                &sender_pk,
+                &owner_pk,
+                value,
+                value_blinder,
+                sender_blinder,
+            )
+        };
+        notes.push(note);
+    }
+
+    assert_eq!(
+        phoenix_balance(&(&owner_sk).into(), notes),
+        expected_balance
+    );
+}
+
 #[test]
 fn test_derive_phoenix_sk() {
     // it is important that we always derive the same key from a fixed seed