wallet-core: Add filter_map function to filter owned notes

wallet-core/src/keys.rs

@@ -6,8 +6,12 @@
 
 //! Utilities to derive keys from the seed.
 
+use alloc::vec::Vec;
+use core::ops::Range;
+
 use rand_chacha::{rand_core::SeedableRng, ChaCha12Rng};
 use sha2::{Digest, Sha256};
+use zeroize::Zeroize;
 
 use execution_core::{
     signatures::bls::SecretKey as BlsSecretKey,
@@ -16,7 +20,6 @@ use execution_core::{
         ViewKey as PhoenixViewKey,
     },
 };
-use zeroize::Zeroize;
 
 use crate::RNG_SEED;
 
@@ -40,6 +43,19 @@ pub fn derive_phoenix_sk(seed: &[u8; RNG_SEED], index: u8) -> PhoenixSecretKey {
     PhoenixSecretKey::random(&mut rng_with_index(seed, index, b"PSK"))
 }
 
+/// Generates multiple [`PhoenixSecretKey`] from a seed and a range of indices.
+///
+/// The randomness is generated using [`rng_with_index`].
+#[must_use]
+pub fn derive_multiple_phoenix_sk(
+    seed: &[u8; RNG_SEED],
+    index_range: Range<u8>,
+) -> Vec<PhoenixSecretKey> {
+    index_range
+        .map(|index| derive_phoenix_sk(seed, index))
+        .collect()
+}
+
 /// Generates a [`PheonixPublicKey`] from its seed and index.
 ///
 /// First the [`PhoenixSecretKey`] is derived with [`derive_phoenix_sk`], then

wallet-core/src/lib.rs

@@ -22,9 +22,35 @@ pub const RNG_SEED: usize = 64;
 // phoenix-transaction
 const MAX_INPUT_NOTES: usize = 4;
 
+use alloc::collections::btree_map::BTreeMap;
 use alloc::vec::Vec;
 
-use execution_core::transfer::phoenix::{Note, ViewKey as PhoenixViewKey};
+use execution_core::{
+    transfer::phoenix::{
+        Note, SecretKey as PhoenixSecretKey, ViewKey as PhoenixViewKey,
+    },
+    BlsScalar,
+};
+
+/// Filter all notes 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| {
+            for sk in keys.as_ref() {
+                if sk.owns(note.stealth_address()) {
+                    let nullifier = note.gen_nullifier(sk);
+                    notes_map.insert(nullifier, note.clone());
+                }
+            }
+            notes_map
+        })
+}
 
 /// Calculate the sum for all the given [`Note`]s that belong to the given
 /// [`PhoenixViewKey`].

wallet-core/tests/keys.rs

@@ -7,7 +7,8 @@
 use dusk_bytes::Serializable;
 
 use wallet_core::keys::{
-    derive_bls_sk, derive_phoenix_pk, derive_phoenix_sk, derive_phoenix_vk,
+    derive_bls_sk, derive_multiple_phoenix_sk, derive_phoenix_pk,
+    derive_phoenix_sk, derive_phoenix_vk,
 };
 
 const SEED: [u8; 64] = [0; 64];
@@ -26,6 +27,29 @@ fn test_derive_phoenix_sk() {
     assert_eq!(derive_phoenix_sk(&SEED, INDEX).to_bytes(), sk_bytes);
 }
 
+#[test]
+fn test_derive_multiple_phoenix_sk() {
+    // it is important that we always derive the same key from a fixed seed
+    let sk_bytes_0 = [
+        160, 210, 234, 8, 94, 23, 76, 60, 130, 143, 137, 225, 37, 83, 68, 218,
+        207, 192, 171, 235, 252, 130, 133, 62, 18, 232, 6, 49, 245, 123, 220,
+        12, 250, 111, 39, 88, 24, 41, 156, 174, 241, 14, 118, 173, 11, 53, 192,
+        126, 7, 119, 70, 69, 212, 230, 124, 79, 223, 140, 93, 153, 33, 147,
+        163, 0,
+    ];
+    let sk_bytes_1 = [
+        0, 229, 97, 222, 152, 25, 163, 173, 84, 216, 211, 69, 205, 122, 63,
+        227, 98, 234, 163, 66, 145, 71, 217, 221, 29, 78, 36, 77, 68, 29, 144,
+        2, 235, 80, 237, 21, 95, 54, 16, 89, 74, 200, 124, 248, 119, 216, 38,
+        167, 19, 129, 205, 138, 218, 57, 198, 4, 4, 202, 115, 62, 55, 213, 141,
+        0,
+    ];
+
+    let keys = derive_multiple_phoenix_sk(&SEED, INDEX..INDEX + 2);
+    assert_eq!(keys[0].to_bytes(), sk_bytes_0,);
+    assert_eq!(keys[1].to_bytes(), sk_bytes_1,);
+}
+
 #[test]
 fn test_derive_phoenix_pk() {
     // it is important that we always derive the same key from a fixed seed

wallet-core/tests/notes.rs

@@ -0,0 +1,148 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+//
+// Copyright (c) DUSK NETWORK. All rights reserved.
+
+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::{
+    keys::derive_multiple_phoenix_sk, map_owned, phoenix_balance, BalanceInfo,
+};
+
+#[test]
+fn test_map_owned() {
+    // Assuming this set of notes where the number used as suffix is the
+    // "owner":
+    // notes := [A1, B1, C2, D2, E1, F3]
+
+    let mut rng = StdRng::seed_from_u64(0xdab);
+    const SEED_1: [u8; 64] = [1; 64];
+    const SEED_2: [u8; 64] = [2; 64];
+
+    let owner_1_sks = derive_multiple_phoenix_sk(&SEED_1, 0..3);
+    let owner_1_pks = [
+        PhoenixPublicKey::from(&owner_1_sks[0]),
+        PhoenixPublicKey::from(&owner_1_sks[1]),
+        PhoenixPublicKey::from(&owner_1_sks[2]),
+    ];
+    let owner_2_sks = derive_multiple_phoenix_sk(&SEED_2, 0..2);
+    let owner_2_pks = [
+        PhoenixPublicKey::from(&owner_2_sks[0]),
+        PhoenixPublicKey::from(&owner_2_sks[1]),
+    ];
+    let owner_3_pk =
+        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
+    ];
+
+    // notes with idx 0, 1 and 4 are owned by owner_1
+    let notes_by_1 = map_owned(&owner_1_sks, &notes);
+    assert_eq!(notes_by_1.len(), 3);
+    let note = &notes[0];
+    let nullifier = note.gen_nullifier(&owner_1_sks[0]);
+    assert_eq!(&notes_by_1[&nullifier], note);
+    let note = &notes[1];
+    let nullifier = note.gen_nullifier(&owner_1_sks[1]);
+    assert_eq!(&notes_by_1[&nullifier], note);
+    let note = &notes[4];
+    let nullifier = note.gen_nullifier(&owner_1_sks[2]);
+    assert_eq!(&notes_by_1[&nullifier], note);
+
+    // notes with idx 2 and 3 are owned by owner_2
+    let notes_by_2 = map_owned(&owner_2_sks, &notes);
+    assert_eq!(notes_by_2.len(), 2);
+    let note = &notes[2];
+    let nullifier = note.gen_nullifier(&owner_2_sks[0]);
+    assert_eq!(&notes_by_2[&nullifier], note);
+    let note = &notes[3];
+    let nullifier = note.gen_nullifier(&owner_2_sks[1]);
+    assert_eq!(&notes_by_2[&nullifier], note);
+}
+
+#[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 mut notes = Vec::new();
+
+    // create the notes
+    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 };
+
+        notes.push(gen_note(&mut rng, obfuscated_note, &owner_pk, value));
+
+        // also push some notes that are not owned
+        if value % 4 == 0 {
+            let not_owner_pk =
+                PhoenixPublicKey::from(&PhoenixSecretKey::random(&mut rng));
+            notes.push(gen_note(
+                &mut rng,
+                obfuscated_note,
+                &not_owner_pk,
+                value,
+            ));
+        }
+    }
+
+    // the sum of these notes should be 5 * 11 = 55
+    // and the spendable notes are 7 + 8 + 9 + 10 = 34
+    let expected_balance = BalanceInfo {
+        value: 55,
+        spendable: 34,
+    };
+
+    assert_eq!(
+        phoenix_balance(&(&owner_sk).into(), notes),
+        expected_balance
+    );
+}
+
+fn gen_note(
+    rng: &mut StdRng,
+    obfuscated_note: bool,
+    owner_pk: &PhoenixPublicKey,
+    value: u64,
+) -> Note {
+    let sender_pk = PhoenixPublicKey::from(&PhoenixSecretKey::random(rng));
+
+    let value_blinder = JubJubScalar::random(&mut *rng);
+    let sender_blinder = [
+        JubJubScalar::random(&mut *rng),
+        JubJubScalar::random(&mut *rng),
+    ];
+    if obfuscated_note {
+        Note::obfuscated(
+            rng,
+            &sender_pk,
+            &owner_pk,
+            value,
+            value_blinder,
+            sender_blinder,
+        )
+    } else {
+        Note::transparent(rng, &sender_pk, &owner_pk, value, sender_blinder)
+    }
+}