Merge pull request #2323 from dusk-network/wallet-core-inputs

wallet-core: ensure the pick notes code will works on FFI too

rusk-wallet/src/clients.rs

@@ -11,11 +11,14 @@ use execution_core::{
     transfer::{phoenix::Prove, Transaction},
     Error as ExecutionCoreError,
 };
+
+use execution_core::transfer::phoenix::{Note, NoteLeaf};
+
 use flume::Receiver;
 use tokio::time::{sleep, Duration};
 use wallet_core::{
-    input::try_input_notes,
     keys::{derive_phoenix_pk, derive_phoenix_sk, derive_phoenix_vk},
+    pick_notes,
 };
 use zeroize::Zeroize;
 
@@ -194,19 +197,19 @@ impl State {
             .into_iter()
             .map(|data| {
                 let note = data.note;
-                let nullifiers = note.gen_nullifier(&sk);
-                let value = note.value(Some(&vk)).unwrap();
-
-                Ok((note, value, nullifiers))
+                let block_height = data.height;
+                let nullifier = note.gen_nullifier(&sk);
+                let leaf = NoteLeaf { note, block_height };
+                Ok((nullifier, leaf))
             })
             .collect();
 
-        let inputs = try_input_notes(inputs?, target)
+        let inputs = pick_notes(&vk, inputs?.into(), target)
             .into_iter()
-            .map(|(note, scalar)| {
-                let opening = self.fetch_opening(&note)?;
+            .map(|(scalar, note)| {
+                let opening = self.fetch_opening(note.as_ref())?;
 
-                Ok((note, opening, scalar))
+                Ok((note.note.clone(), opening, *scalar))
             })
             .collect();
 

wallet-core/src/ffi.rs

@@ -126,7 +126,7 @@ pub unsafe fn map_owned(
     let notes: Vec<NoteLeaf> = from_bytes::<Vec<NoteLeaf>>(&notes)
         .or(Err(ErrorCode::UnarchivingError))?;
 
-    let owned = notes::map_owned(&keys, notes);
+    let owned = notes::owned::map(&keys, notes);
 
     keys.into_iter().for_each(|mut sk| sk.zeroize());
 

wallet-core/src/lib.rs

@@ -22,7 +22,6 @@ extern crate alloc;
 #[macro_use]
 mod ffi;
 
-pub mod input;
 pub mod keys;
 pub mod notes;
 pub mod transaction;
@@ -32,8 +31,12 @@ pub type Seed = [u8; 64];
 
 pub mod prelude {
     //! Re-export of the most commonly used types and traits.
-    pub use crate::input::MAX_INPUT_NOTES;
     pub use crate::keys;
+    pub use crate::notes::MAX_INPUT_NOTES;
 }
 
-pub use notes::{map_owned, phoenix_balance, BalanceInfo};
+pub use notes::balance::{
+    calculate as phoenix_balance, TotalAmount as BalanceInfo,
+};
+pub use notes::owned::map as map_owned;
+pub use notes::pick::notes as pick_notes;

wallet-core/src/notes.rs

@@ -6,150 +6,13 @@
 
 //! Provides functions and types for interacting with notes.
 
-use alloc::vec::Vec;
-use core::ops::Index;
-use dusk_bytes::{DeserializableSlice, Serializable, Write};
-use execution_core::transfer::phoenix::{Note, ViewKey as PhoenixViewKey};
-use execution_core::{
-    transfer::phoenix::{NoteLeaf, SecretKey as PhoenixSecretKey},
-    BlsScalar,
-};
-
-use rkyv::{Archive, Deserialize, Serialize};
-
-// The maximum amount of input notes that can be spend in one
-// phoenix-transaction
-const MAX_INPUT_NOTES: usize = 4;
-
-/// A collection of notes stored as key-value pairs.
-/// The key is a `BlsScalar` and the value is a `NoteLeaf`.
-/// Duplicates are allowed.
-#[derive(Default, Archive, Serialize, Deserialize, Debug)]
-pub struct OwnedList {
-    /// The underlying storage of key-value pairs where
-    /// `BlsScalar` is the key and `NoteLeaf` is the value.
-    entries: Vec<(BlsScalar, NoteLeaf)>,
-}
-
-impl OwnedList {
-    /// Inserts a new key-value pair into the collection.
-    pub fn insert(&mut self, key: BlsScalar, value: NoteLeaf) {
-        self.entries.push((key, value));
-    }
-
-    /// Returns the number of entries (key-value pairs) in the collection.
-    #[must_use]
-    pub fn len(&self) -> usize {
-        self.entries.len()
-    }
-
-    /// Checks if the collection is empty.
-    #[must_use]
-    pub fn is_empty(&self) -> bool {
-        self.entries.is_empty()
-    }
-
-    /// Retrieves the value (`NoteLeaf`) associated with a given key
-    #[must_use]
-    pub fn get(&self, key: &BlsScalar) -> Option<&NoteLeaf> {
-        self.entries.iter().find(|(k, _)| k == key).map(|(_, v)| v)
-    }
-
-    /// Retrieves all keys in the collection.
-    #[must_use]
-    pub fn keys(&self) -> Vec<BlsScalar> {
-        self.entries.iter().map(|(k, _)| *k).collect()
-    }
-}
-
-impl Index<&BlsScalar> for OwnedList {
-    type Output = NoteLeaf;
-
-    /// Retrieves the value (`NoteLeaf`) associated with a given key
-    /// (`BlsScalar`).
-    ///
-    /// Panics if the key is not found in the collection.
-    fn index(&self, index: &BlsScalar) -> &Self::Output {
-        self.get(index).expect("key not found")
-    }
-}
-
-/// 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<[NoteLeaf]>,
-) -> OwnedList {
-    notes.as_ref().iter().fold(
-        OwnedList::default(),
-        |mut notes_map, note_leaf| {
-            for sk in keys.as_ref() {
-                if sk.owns(note_leaf.note.stealth_address()) {
-                    let nullifier = note_leaf.note.gen_nullifier(sk);
-                    notes_map.insert(nullifier, note_leaf.clone());
-                    break;
-                }
-            }
-            notes_map
-        },
-    )
-}
-
-/// Calculate the sum for all the given [`Note`]s that belong to the given
-/// [`PhoenixViewKey`].
-pub fn phoenix_balance<T>(
-    phoenix_vk: &PhoenixViewKey,
-    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));
-
-    let spendable = values.iter().take(MAX_INPUT_NOTES).sum();
-    let value = spendable + values.iter().skip(MAX_INPUT_NOTES).sum::<u64>();
-
-    BalanceInfo { value, spendable }
-}
-
-/// Information about the balance of a particular key.
-#[derive(Debug, Default, Hash, Clone, Copy, PartialEq, Eq)]
-pub struct BalanceInfo {
-    /// The total value of the balance.
-    pub value: u64,
-    /// The maximum _spendable_ value in a single transaction. This is
-    /// different from `value` since there is a maximum number of notes one can
-    /// spend.
-    pub spendable: u64,
-}
-
-impl Serializable<{ 2 * u64::SIZE }> for BalanceInfo {
-    type Error = dusk_bytes::Error;
-
-    fn from_bytes(buf: &[u8; Self::SIZE]) -> Result<Self, Self::Error>
-    where
-        Self: Sized,
-    {
-        let mut reader = &buf[..];
-
-        let value = u64::from_reader(&mut reader)?;
-        let spendable = u64::from_reader(&mut reader)?;
-
-        Ok(Self { value, spendable })
-    }
-
-    #[allow(unused_must_use)]
-    fn to_bytes(&self) -> [u8; Self::SIZE] {
-        let mut buf = [0u8; Self::SIZE];
-        let mut writer = &mut buf[..];
-
-        writer.write(&self.value.to_bytes());
-        writer.write(&self.spendable.to_bytes());
-
-        buf
-    }
-}
+/// Module for balance information.
+pub mod balance;
+/// Module for owned notes.
+pub mod owned;
+/// Module for picking notes.
+pub mod pick;
+
+/// The maximum amount of input notes that can be spend in one
+/// phoenix-transaction
+pub const MAX_INPUT_NOTES: usize = 4;

wallet-core/src/notes/balance.rs

@@ -0,0 +1,73 @@
+// 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.
+
+//! Provides functions and types for calculate notes' balance.
+
+use alloc::vec::Vec;
+
+use dusk_bytes::{DeserializableSlice, Serializable, Write};
+use execution_core::transfer::phoenix::{Note, ViewKey as PhoenixViewKey};
+
+use crate::notes::MAX_INPUT_NOTES;
+
+/// Calculate the sum for all the given [`Note`]s that belong to the given
+/// [`PhoenixViewKey`].
+pub fn calculate<T>(
+    vk: &PhoenixViewKey,
+    notes: impl Iterator<Item = T>,
+) -> TotalAmount
+where
+    T: AsRef<Note>,
+{
+    let mut values: Vec<u64> = notes
+        .filter_map(|note| note.as_ref().value(Some(vk)).ok())
+        .collect();
+
+    values.sort_by(|a, b| b.cmp(a));
+
+    let spendable = values.iter().take(MAX_INPUT_NOTES).sum();
+    let value = spendable + values.iter().skip(MAX_INPUT_NOTES).sum::<u64>();
+
+    TotalAmount { value, spendable }
+}
+
+/// Information about the balance of a particular key.
+#[derive(Debug, Default, Hash, Clone, Copy, PartialEq, Eq)]
+pub struct TotalAmount {
+    /// The total value of the balance.
+    pub value: u64,
+    /// The maximum _spendable_ value in a single transaction. This is
+    /// different from `value` since there is a maximum number of notes one can
+    /// spend.
+    pub spendable: u64,
+}
+
+impl Serializable<{ 2 * u64::SIZE }> for TotalAmount {
+    type Error = dusk_bytes::Error;
+
+    fn from_bytes(buf: &[u8; Self::SIZE]) -> Result<Self, Self::Error>
+    where
+        Self: Sized,
+    {
+        let mut reader = &buf[..];
+
+        let value = u64::from_reader(&mut reader)?;
+        let spendable = u64::from_reader(&mut reader)?;
+
+        Ok(Self { value, spendable })
+    }
+
+    #[allow(unused_must_use)]
+    fn to_bytes(&self) -> [u8; Self::SIZE] {
+        let mut buf = [0u8; Self::SIZE];
+        let mut writer = &mut buf[..];
+
+        writer.write(&self.value.to_bytes());
+        writer.write(&self.spendable.to_bytes());
+
+        buf
+    }
+}