wallet-core: Add filter_map function

Resolves #2129

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,9 +20,6 @@ use execution_core::{
         ViewKey as PhoenixViewKey,
     },
 };
-use zeroize::Zeroize;
-
-use core::ops::Range;
 
 use crate::RNG_SEED;
 
@@ -42,7 +43,7 @@ pub fn derive_phoenix_sk(seed: &[u8; RNG_SEED], index: u8) -> PhoenixSecretKey {
     PhoenixSecretKey::random(&mut rng_with_index(seed, index, b"PSK"))
 }
 
-/// Generates a [`PhoenixSecretKey`] from a seed and index.
+/// Generates multiple [`PhoenixSecretKey`] from a seed and a range of indices.
 ///
 /// The randomness is generated using [`rng_with_index`].
 #[must_use]
@@ -53,11 +54,7 @@ pub fn derive_multiple_phoenix_sk(
     let mut keys = Vec::new();
 
     for index in index_range {
-        // note that if we change the string used for the rng, all previously
-        // generated keys will become invalid
-        keys.push(PhoenixSecretKey::random(&mut rng_with_index(
-            seed, index, b"PSK",
-        )));
+        keys.push(derive_phoenix_sk(seed, index));
     }
 
     keys

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 filter_map(
+    keys: impl AsRef<[PhoenixSecretKey]>,
+    notes: impl AsRef<[Note]>,
+) -> BTreeMap<BlsScalar, Note> {
+    let mut notes_map = BTreeMap::new();
+
+    notes.as_ref().into_iter().for_each(|note| {
+        for sk in keys.as_ref().iter() {
+            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/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::{
+    filter_map, keys::derive_multiple_phoenix_sk, phoenix_balance, BalanceInfo,
+};
+
+#[test]
+fn test_filter_map() {
+    // 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 = filter_map(&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 = filter_map(&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)
+    }
+}