From 25f2e662f515a4e42e8b016bb4a9704d35f3d4f5 Mon Sep 17 00:00:00 2001
From: moana <moana@dusk.network>
Date: Sat, 31 Aug 2024 09:40:12 +0200
Subject: [PATCH 1/2] wallet-core: Store the `Note`s together with their
 block-height

---
 wallet-core/src/lib.rs     | 27 +++++++++++++++------------
 wallet-core/tests/notes.rs | 38 +++++++++++++++++++-------------------
 2 files changed, 34 insertions(+), 31 deletions(-)

diff --git a/wallet-core/src/lib.rs b/wallet-core/src/lib.rs
index 98e64c467b..5adcb7e87a 100644
--- a/wallet-core/src/lib.rs
+++ b/wallet-core/src/lib.rs
@@ -35,24 +35,27 @@ use execution_core::{
     BlsScalar,
 };
 
-/// Filter all notes that are owned by the given keys, mapped to their
-/// nullifiers.
+/// 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<[Note]>,
-) -> BTreeMap<BlsScalar, Note> {
-    notes
-        .as_ref()
-        .iter()
-        .fold(BTreeMap::new(), |mut notes_map, note| {
+    notes: impl AsRef<[EnrichedNote]>,
+) -> BTreeMap<BlsScalar, EnrichedNote> {
+    notes.as_ref().iter().fold(
+        BTreeMap::new(),
+        |mut notes_map, enriched_note| {
             for sk in keys.as_ref() {
-                if sk.owns(note.stealth_address()) {
-                    let nullifier = note.gen_nullifier(sk);
-                    notes_map.insert(nullifier, note.clone());
+                if sk.owns(enriched_note.0.stealth_address()) {
+                    let nullifier = enriched_note.0.gen_nullifier(sk);
+                    notes_map.insert(nullifier, enriched_note.clone());
                 }
             }
             notes_map
-        })
+        },
+    )
 }
 
 /// Calculate the sum for all the given [`Note`]s that belong to the given
diff --git a/wallet-core/tests/notes.rs b/wallet-core/tests/notes.rs
index c9ecfce44f..a6c425ad2b 100644
--- a/wallet-core/tests/notes.rs
+++ b/wallet-core/tests/notes.rs
@@ -44,37 +44,37 @@ fn test_map_owned() {
         PhoenixPublicKey::from(&PhoenixSecretKey::random(&mut rng));
 
     let value = 42;
-    let notes = 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
+    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
     ];
 
     // notes with idx 0, 1 and 4 are owned by owner_1
-    let notes_by_1 = map_owned(&owner_1_sks, &notes);
+    let notes_by_1 = map_owned(&owner_1_sks, &enriched_notes);
     assert_eq!(notes_by_1.len(), 3);
-    let note = &notes[0];
+    let note = &enriched_notes[0].0;
     let nullifier = note.gen_nullifier(&owner_1_sks[0]);
-    assert_eq!(&notes_by_1[&nullifier], note);
-    let note = &notes[1];
+    assert_eq!(&notes_by_1[&nullifier].0, note);
+    let note = &enriched_notes[1].0;
     let nullifier = note.gen_nullifier(&owner_1_sks[1]);
-    assert_eq!(&notes_by_1[&nullifier], note);
-    let note = &notes[4];
+    assert_eq!(&notes_by_1[&nullifier].0, note);
+    let note = &enriched_notes[4].0;
     let nullifier = note.gen_nullifier(&owner_1_sks[2]);
-    assert_eq!(&notes_by_1[&nullifier], note);
+    assert_eq!(&notes_by_1[&nullifier].0, note);
 
     // notes with idx 2 and 3 are owned by owner_2
-    let notes_by_2 = map_owned(&owner_2_sks, &notes);
+    let notes_by_2 = map_owned(&owner_2_sks, &enriched_notes);
     assert_eq!(notes_by_2.len(), 2);
-    let note = &notes[2];
+    let note = &enriched_notes[2].0;
     let nullifier = note.gen_nullifier(&owner_2_sks[0]);
-    assert_eq!(&notes_by_2[&nullifier], note);
-    let note = &notes[3];
+    assert_eq!(&notes_by_2[&nullifier].0, note);
+    let note = &enriched_notes[3].0;
     let nullifier = note.gen_nullifier(&owner_2_sks[1]);
-    assert_eq!(&notes_by_2[&nullifier], note);
+    assert_eq!(&notes_by_2[&nullifier].0, note);
 }
 
 #[test]

From 983de8f337bc10a12dd42fcfd41bb4490179b14c Mon Sep 17 00:00:00 2001
From: moana <moana@dusk.network>
Date: Sat, 31 Aug 2024 09:49:33 +0200
Subject: [PATCH 2/2] test-wallet: Move `EnrichedNote` to `wallet-core`

---
 test-wallet/src/lib.rs | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/test-wallet/src/lib.rs b/test-wallet/src/lib.rs
index 32d4fbe7bd..2df93dffec 100644
--- a/test-wallet/src/lib.rs
+++ b/test-wallet/src/lib.rs
@@ -31,8 +31,9 @@ use execution_core::{
     BlsScalar,
 };
 
-pub use wallet_core::keys::{
-    derive_bls_sk, derive_phoenix_pk, derive_phoenix_sk,
+pub use wallet_core::{
+    keys::{derive_bls_sk, derive_phoenix_pk, derive_phoenix_sk},
+    EnrichedNote,
 };
 
 pub use imp::*;
@@ -104,9 +105,6 @@ pub trait Store {
     }
 }
 
-/// Tuple containing Note and block height
-pub type EnrichedNote = (Note, u64);
-
 /// Types that are clients of the state API.
 pub trait StateClient {
     /// Error returned by the node client.