Skip to content

Commit

Permalink
Merge pull request #2228 from dusk-network/mocello/2227_enriched_note
Browse files Browse the repository at this point in the history
wallet-core: Also store the block-height for the owned notes
  • Loading branch information
moCello authored Sep 1, 2024
2 parents 8879b48 + 983de8f commit 2de971c
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 36 deletions.
8 changes: 3 additions & 5 deletions test-wallet/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand Down Expand Up @@ -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.
Expand Down
27 changes: 15 additions & 12 deletions wallet-core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
38 changes: 19 additions & 19 deletions wallet-core/tests/notes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down

0 comments on commit 2de971c

Please sign in to comment.