Skip to content

Commit

Permalink
wallet-core: Move some functionality from lib's root to notes module
Browse files Browse the repository at this point in the history
- Add `NOTES_BUFFER_SIZE` to indicate the scratch buffer size for notes
- Add `MAX_INPUT_NOTES`, `phoenix_balance` and `BalanceInfo` under `notes` module
- Remove `MAX_INPUT_NOTES`, `phoenix_balance` and `BalanceInfo` from lib.rs
  • Loading branch information
ZER0 committed Sep 9, 2024
1 parent d5365aa commit 44b8bfe
Show file tree
Hide file tree
Showing 3 changed files with 74 additions and 65 deletions.
8 changes: 6 additions & 2 deletions wallet-core/src/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,10 @@ use zeroize::Zeroize;

use rkyv::{from_bytes, to_bytes};

/// The size of the scratch buffer used for parsing the notes.
/// It can roughly contains less than 128 serialized notes.
const NOTES_BUFFER_SIZE: usize = 96 * 1024;

/// The alignment of the memory allocated by the FFI.
///
/// This is 1 because we're not allocating any complex data structures, and
Expand Down Expand Up @@ -126,8 +130,8 @@ pub unsafe fn map_owned(

keys.into_iter().for_each(|mut sk| sk.zeroize());

let bytes =
to_bytes::<_, 4096>(&owned).or(Err(ErrorCode::ArchivingError))?;
let bytes = to_bytes::<_, NOTES_BUFFER_SIZE>(&owned)
.or(Err(ErrorCode::ArchivingError))?;

let len = bytes.len().to_le_bytes();

Expand Down
62 changes: 2 additions & 60 deletions wallet-core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,8 @@ 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::{input::MAX_INPUT_NOTES, keys::RNG_SEED};
}

use alloc::vec::Vec;
Expand All @@ -44,62 +44,4 @@ use execution_core::transfer::phoenix::{Note, ViewKey as PhoenixViewKey};

pub use notes::map_owned;

/// 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(input::MAX_INPUT_NOTES).sum();
let value =
spendable + values.iter().skip(input::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
}
}
pub use notes::{map_owned, phoenix_balance, BalanceInfo};
69 changes: 66 additions & 3 deletions wallet-core/src/notes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,23 @@
//
// Copyright (c) DUSK NETWORK. All rights reserved.

//! Implementations of basic wallet functionalities to create transactions.
//! 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.
Expand Down Expand Up @@ -77,8 +83,6 @@ pub fn map_owned(
notes.as_ref().iter().fold(
OwnedList::default(),
|mut notes_map, note_leaf| {
eprintln!("Printing note...");
dbg!(note_leaf);
for sk in keys.as_ref() {
if sk.owns(note_leaf.note.stealth_address()) {
let nullifier = note_leaf.note.gen_nullifier(sk);
Expand All @@ -90,3 +94,62 @@ pub fn map_owned(
},
)
}

/// 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
}
}

0 comments on commit 44b8bfe

Please sign in to comment.