Merge pull request #2241 from dusk-network/mocello/2239_enriched_balance

wallet-core: Refactor `phoenix_balance` to work with `NoteLeaf`

execution-core/src/transfer/phoenix.rs

@@ -55,6 +55,12 @@ pub struct NoteLeaf {
     pub note: Note,
 }
 
+impl AsRef<Note> for NoteLeaf {
+    fn as_ref(&self) -> &Note {
+        &self.note
+    }
+}
+
 /// Label used for the ZK transcript initialization. Must be the same for prover
 /// and verifier.
 pub const TRANSCRIPT_LABEL: &[u8] = b"dusk-network";

test-wallet/src/imp.rs

@@ -26,7 +26,7 @@ use execution_core::{
         contract_exec::ContractExec,
         moonlight::{AccountData, Transaction as MoonlightTransaction},
         phoenix::{
-            Note, NoteOpening, PublicKey as PhoenixPublicKey,
+            Note, NoteLeaf, NoteOpening, PublicKey as PhoenixPublicKey,
             SecretKey as PhoenixSecretKey, ViewKey as PhoenixViewKey,
         },
         Transaction,
@@ -224,31 +224,31 @@ where
     fn unspent_notes_and_nullifiers(
         &self,
         sk: &PhoenixSecretKey,
-    ) -> Result<Vec<(Note, BlsScalar)>, Error<S, SC>> {
+    ) -> Result<Vec<(NoteLeaf, BlsScalar)>, Error<S, SC>> {
         let vk = PhoenixViewKey::from(sk);
 
-        let notes: Vec<Note> = self
-            .state
-            .fetch_notes(&vk)
-            .map_err(Error::from_state_err)?
-            .into_iter()
-            .map(|(note, _bh)| note)
-            .collect();
+        let note_leaves =
+            self.state.fetch_notes(&vk).map_err(Error::from_state_err)?;
 
-        let nullifiers: Vec<_> =
-            notes.iter().map(|n| n.gen_nullifier(sk)).collect();
+        let nullifiers: Vec<_> = note_leaves
+            .iter()
+            .map(|(note, _bh)| note.gen_nullifier(sk))
+            .collect();
 
         let existing_nullifiers = self
             .state
             .fetch_existing_nullifiers(&nullifiers)
             .map_err(Error::from_state_err)?;
 
-        let unspent_notes_and_nullifiers = notes
+        let unspent_notes_and_nullifiers = note_leaves
             .into_iter()
             .zip(nullifiers.into_iter())
             .filter(|(_note, nullifier)| {
                 !existing_nullifiers.contains(nullifier)
             })
+            .map(|((note, block_height), nullifier)| {
+                (NoteLeaf { note, block_height }, nullifier)
+            })
             .collect();
 
         Ok(unspent_notes_and_nullifiers)
@@ -271,12 +271,13 @@ where
             Vec::with_capacity(unspent_notes_nullifiers.len());
 
         let mut accumulated_value = 0;
-        for (note, nullifier) in unspent_notes_nullifiers {
-            let val = note
+        for (note_leaf, nullifier) in unspent_notes_nullifiers {
+            let val = note_leaf
+                .note
                 .value(Some(&sender_vk))
                 .map_err(|_| ExecutionError::PhoenixOwnership)?;
             accumulated_value += val;
-            notes_values_nullifiers.push((note, val, nullifier));
+            notes_values_nullifiers.push((note_leaf.note, val, nullifier));
         }
 
         if accumulated_value < transaction_cost {
@@ -670,12 +671,12 @@ where
         let mut phoenix_sk = derive_phoenix_sk(&seed, sk_index);
         let phoenix_vk = PhoenixViewKey::from(&phoenix_sk);
 
-        let unspent_notes: Vec<Note> = self
+        let unspent_notes: Vec<NoteLeaf> = self
             .unspent_notes_and_nullifiers(&phoenix_sk)?
             .into_iter()
-            .map(|(note, _nullifier)| note)
+            .map(|(note_leaf, _nul)| note_leaf)
             .collect();
-        let balance = phoenix_balance(&phoenix_vk, unspent_notes);
+        let balance = phoenix_balance(&phoenix_vk, unspent_notes.iter());
 
         seed.zeroize();
         phoenix_sk.zeroize();

wallet-core/src/lib.rs

@@ -10,7 +10,6 @@
 #![deny(missing_docs)]
 #![deny(rustdoc::broken_intra_doc_links)]
 #![deny(clippy::pedantic)]
-
 #[cfg(target_family = "wasm")]
 #[global_allocator]
 static ALLOC: dlmalloc::GlobalDlmalloc = dlmalloc::GlobalDlmalloc;
@@ -37,45 +36,44 @@ use dusk_bytes::{DeserializableSlice, Serializable, Write};
 
 use execution_core::{
     transfer::phoenix::{
-        Note, SecretKey as PhoenixSecretKey, ViewKey as PhoenixViewKey,
+        Note, NoteLeaf, SecretKey as PhoenixSecretKey,
+        ViewKey as PhoenixViewKey,
     },
     BlsScalar,
 };
 
-/// Tuple containing Note and block height
-pub type EnrichedNote = (Note, u64);
-
 /// Filter all notes and their block height that are owned by the given keys,
 /// mapped to their nullifiers.
 pub fn map_owned(
     keys: impl AsRef<[PhoenixSecretKey]>,
-    notes: impl AsRef<[EnrichedNote]>,
-) -> BTreeMap<BlsScalar, EnrichedNote> {
-    notes.as_ref().iter().fold(
-        BTreeMap::new(),
-        |mut notes_map, enriched_note| {
+    notes: impl AsRef<[NoteLeaf]>,
+) -> BTreeMap<BlsScalar, NoteLeaf> {
+    notes
+        .as_ref()
+        .iter()
+        .fold(BTreeMap::new(), |mut notes_map, note_leaf| {
             for sk in keys.as_ref() {
-                if sk.owns(enriched_note.0.stealth_address()) {
-                    let nullifier = enriched_note.0.gen_nullifier(sk);
-                    notes_map.insert(nullifier, enriched_note.clone());
+                if sk.owns(note_leaf.note.stealth_address()) {
+                    let nullifier = note_leaf.note.gen_nullifier(sk);
+                    notes_map.insert(nullifier, note_leaf.clone());
                 }
             }
             notes_map
-        },
-    )
+        })
 }
 
 /// Calculate the sum for all the given [`Note`]s that belong to the given
 /// [`PhoenixViewKey`].
-pub fn phoenix_balance(
+pub fn phoenix_balance<T>(
     phoenix_vk: &PhoenixViewKey,
-    unspent_notes: impl AsRef<[Note]>,
-) -> BalanceInfo {
-    let mut values: Vec<u64> = Vec::new();
-    let unspent_notes = unspent_notes.as_ref();
-    for note in unspent_notes {
-        values.push(note.value(Some(phoenix_vk)).unwrap_or_default());
-    }
+    notes: impl Iterator<Item = T>,
+) -> BalanceInfo
+where
+    T: AsRef<Note>,
+{
+    let mut values: Vec<u64> = notes
+        .filter_map(|note| note.as_ref().value(Some(phoenix_vk)).ok())
+        .collect();
 
     values.sort_by(|a, b| b.cmp(a));
 

wallet-core/tests/notes.rs

@@ -6,11 +6,12 @@
 
 use ff::Field;
 use rand::rngs::StdRng;
-use rand::SeedableRng;
+use rand::{Rng, SeedableRng};
 
 use execution_core::{
     transfer::phoenix::{
-        Note, PublicKey as PhoenixPublicKey, SecretKey as PhoenixSecretKey,
+        Note, NoteLeaf, PublicKey as PhoenixPublicKey,
+        SecretKey as PhoenixSecretKey,
     },
     JubJubScalar,
 };
@@ -44,37 +45,37 @@ fn test_map_owned() {
         PhoenixPublicKey::from(&PhoenixSecretKey::random(&mut rng));
 
     let value = 42;
-    let enriched_notes = vec![
-        (gen_note(&mut rng, true, &owner_1_pks[0], value), 1), // owner 1
-        (gen_note(&mut rng, true, &owner_1_pks[1], value), 1), // owner 1
-        (gen_note(&mut rng, true, &owner_2_pks[0], value), 1), // owner 2
-        (gen_note(&mut rng, true, &owner_2_pks[1], value), 1), // owner 2
-        (gen_note(&mut rng, true, &owner_1_pks[2], value), 1), // owner 1
-        (gen_note(&mut rng, true, &owner_3_pk, value), 1),     // owner 3
+    let note_leaves = vec![
+        gen_note(&mut rng, true, &owner_1_pks[0], value), // owner 1
+        gen_note(&mut rng, true, &owner_1_pks[1], value), // owner 1
+        gen_note(&mut rng, true, &owner_2_pks[0], value), // owner 2
+        gen_note(&mut rng, true, &owner_2_pks[1], value), // owner 2
+        gen_note(&mut rng, true, &owner_1_pks[2], value), // owner 1
+        gen_note(&mut rng, true, &owner_3_pk, value),     // owner 3
     ];
 
     // notes with idx 0, 1 and 4 are owned by owner_1
-    let notes_by_1 = map_owned(&owner_1_sks, &enriched_notes);
+    let notes_by_1 = map_owned(&owner_1_sks, &note_leaves);
     assert_eq!(notes_by_1.len(), 3);
-    let note = &enriched_notes[0].0;
+    let note = &note_leaves[0].note;
     let nullifier = note.gen_nullifier(&owner_1_sks[0]);
-    assert_eq!(&notes_by_1[&nullifier].0, note);
-    let note = &enriched_notes[1].0;
+    assert_eq!(&notes_by_1[&nullifier].note, note);
+    let note = &note_leaves[1].note;
     let nullifier = note.gen_nullifier(&owner_1_sks[1]);
-    assert_eq!(&notes_by_1[&nullifier].0, note);
-    let note = &enriched_notes[4].0;
+    assert_eq!(&notes_by_1[&nullifier].note, note);
+    let note = &note_leaves[4].note;
     let nullifier = note.gen_nullifier(&owner_1_sks[2]);
-    assert_eq!(&notes_by_1[&nullifier].0, note);
+    assert_eq!(&notes_by_1[&nullifier].note, note);
 
     // notes with idx 2 and 3 are owned by owner_2
-    let notes_by_2 = map_owned(&owner_2_sks, &enriched_notes);
+    let notes_by_2 = map_owned(&owner_2_sks, &note_leaves);
     assert_eq!(notes_by_2.len(), 2);
-    let note = &enriched_notes[2].0;
+    let note = &note_leaves[2].note;
     let nullifier = note.gen_nullifier(&owner_2_sks[0]);
-    assert_eq!(&notes_by_2[&nullifier].0, note);
-    let note = &enriched_notes[3].0;
+    assert_eq!(&notes_by_2[&nullifier].note, note);
+    let note = &note_leaves[3].note;
     let nullifier = note.gen_nullifier(&owner_2_sks[1]);
-    assert_eq!(&notes_by_2[&nullifier].0, note);
+    assert_eq!(&notes_by_2[&nullifier].note, note);
 }
 
 #[test]
@@ -90,7 +91,7 @@ fn test_balance() {
     for value in 0..=10 {
         // we want to test with a mix of transparent and obfuscated notes so we
         // make every 10th note transparent
-        let obfuscated_note = if value % 10 == 0 { false } else { true };
+        let obfuscated_note = value % 10 != 0;
 
         notes.push(gen_note(&mut rng, obfuscated_note, &owner_pk, value));
 
@@ -115,7 +116,7 @@ fn test_balance() {
     };
 
     assert_eq!(
-        phoenix_balance(&(&owner_sk).into(), notes),
+        phoenix_balance(&(&owner_sk).into(), notes.iter()),
         expected_balance
     );
 }
@@ -125,7 +126,7 @@ fn gen_note(
     obfuscated_note: bool,
     owner_pk: &PhoenixPublicKey,
     value: u64,
-) -> Note {
+) -> NoteLeaf {
     let sender_pk = PhoenixPublicKey::from(&PhoenixSecretKey::random(rng));
 
     let value_blinder = JubJubScalar::random(&mut *rng);
@@ -134,15 +135,27 @@ fn gen_note(
         JubJubScalar::random(&mut *rng),
     ];
     if obfuscated_note {
-        Note::obfuscated(
-            rng,
-            &sender_pk,
-            &owner_pk,
-            value,
-            value_blinder,
-            sender_blinder,
-        )
+        NoteLeaf {
+            note: Note::obfuscated(
+                rng,
+                &sender_pk,
+                &owner_pk,
+                value,
+                value_blinder,
+                sender_blinder,
+            ),
+            block_height: rng.gen(),
+        }
     } else {
-        Note::transparent(rng, &sender_pk, &owner_pk, value, sender_blinder)
+        NoteLeaf {
+            note: Note::transparent(
+                rng,
+                &sender_pk,
+                &owner_pk,
+                value,
+                sender_blinder,
+            ),
+            block_height: rng.gen(),
+        }
     }
 }