Skip to content

Commit

Permalink
wallet-core: Fix pick_notes
Browse files Browse the repository at this point in the history
  • Loading branch information
@moCello
moCello committed Nov 5, 2024
1 parent a678a81 commit 4bb2982
Show file tree
Hide file tree
Showing 2 changed files with 143 additions and 56 deletions.
84 changes: 59 additions & 25 deletions wallet-core/src/notes/pick.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,18 +19,23 @@ use execution_core::BlsScalar;
/// while minimizing the value employed. To do this we sort the notes in
/// ascending value order, and go through each combination in a
/// lexicographic order until we find the first combination whose sum is
/// larger or equal to the given value. If such a slice is not found, an
/// larger or equal to the given target-value. If such a slice is not found, an
/// empty vector is returned.
///
/// If the target sum is greater than the sum of the notes then an
/// empty vector is returned.
#[must_use]
pub fn notes(vk: &PhoenixViewKey, notes: NoteList, value: u64) -> NoteList {
pub fn notes(
vk: &PhoenixViewKey,
notes: NoteList,
target_value: u64,
) -> NoteList {
if notes.is_empty() {
return NoteList::default();
}

let mut notes_and_values: Vec<(NoteLeaf, u64, BlsScalar)> = notes
// decrypt the note-values
let mut notes_values_nullifier: Vec<(NoteLeaf, u64, BlsScalar)> = notes
.iter()
.filter_map(|(nullifier, leaf)| {
leaf.as_ref()
Expand All @@ -40,45 +45,74 @@ pub fn notes(vk: &PhoenixViewKey, notes: NoteList, value: u64) -> NoteList {
})
.collect();

let sum: u64 = notes_and_values
// return an empty list if the sum of all notes cannot cover the
// target-value
let sum: u64 = notes_values_nullifier
.iter()
.fold(0, |sum, &(_, value, _)| sum.saturating_add(value));

if sum < value {
if sum < target_value {
return NoteList::default();
}

// if there are less that MAX_INPUT_NOTES notes, we can return the list as
// it is
if notes.len() <= MAX_INPUT_NOTES {
return notes;
}

notes_and_values.sort_by(|(_, aval, _), (_, bval, _)| aval.cmp(bval));
pick_lexicographic(notes_and_values.len(), |indices| {
indices
.iter()
.map(|index| notes_and_values[*index].1)
.sum::<u64>()
>= value
})
.map(|index| notes_and_values[index].clone())
.map(|(n, _, b)| (b, n))
.to_vec()
.into()
// sort the input-notes from smallest to largest value
notes_values_nullifier.sort_by(|(_, aval, _), (_, bval, _)| aval.cmp(bval));

// return an empty list if the MAX_INPUT_NOTES highest notes do not cover
// the target-value
if notes_values_nullifier
.iter()
.skip(notes_values_nullifier.len() - MAX_INPUT_NOTES)
.map(|notes_values_nullifier| notes_values_nullifier.1)
.sum::<u64>()
< target_value
{
return NoteList::default();
}

// pick the MAX_INPUT_NOTES smallest notes that cover the target cost, if no
// combination of four notes cover the target-value, return an empty list
pick_lexicographic(&notes_values_nullifier, target_value)
.expect("a combination of notes should be possible")
.map(|index| notes_values_nullifier[index].clone())
.map(|(n, _, b)| (b, n))
.to_vec()
.into()
}

fn pick_lexicographic<F: Fn(&[usize; MAX_INPUT_NOTES]) -> bool>(
max_len: usize,
is_valid: F,
) -> [usize; MAX_INPUT_NOTES] {
fn is_valid(
notes_values_nullifier: &Vec<(NoteLeaf, u64, BlsScalar)>,
target: u64,
indices: &[usize; MAX_INPUT_NOTES],
) -> bool {
indices
.iter()
.map(|index| notes_values_nullifier[*index].1)
.sum::<u64>()
>= target
}

fn pick_lexicographic(
notes_values_nullifier: &Vec<(NoteLeaf, u64, BlsScalar)>,
target: u64,
) -> Option<[usize; MAX_INPUT_NOTES]> {
let max_len = notes_values_nullifier.len();

// initialize the indices
let mut indices = [0; MAX_INPUT_NOTES];
indices
.iter_mut()
.enumerate()
.for_each(|(i, index)| *index = i);

loop {
if is_valid(&indices) {
return indices;
if is_valid(notes_values_nullifier, target, &indices) {
return Some(indices);
}

let mut i = MAX_INPUT_NOTES - 1;
Expand All @@ -101,5 +135,5 @@ fn pick_lexicographic<F: Fn(&[usize; MAX_INPUT_NOTES]) -> bool>(
}
}

[0; MAX_INPUT_NOTES]
None
}
115 changes: 84 additions & 31 deletions wallet-core/tests/notes.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,7 @@ fn test_balance() {
}

#[test]
fn knapsack_works() {
fn test_pick_notes() {
use rand::SeedableRng;

let mut rng = rand_chacha::ChaCha12Rng::seed_from_u64(0xbeef);
Expand All @@ -182,58 +182,111 @@ fn knapsack_works() {
let vk = PhoenixViewKey::from(&sk);
let pk = PhoenixPublicKey::from(&sk);

// Check single input-note

// sanity check
assert!(pick_notes(&vk, NoteList::default(), 70).is_empty());

// basic check
// basic check with one note
let leaf = gen_note_leaf(&mut rng, true, &pk, 100);
let n = leaf.note.gen_nullifier(&sk);
let available = NoteList::from(vec![(n, leaf)]);
let all_notes = generate_note_list([leaf], &sk);

let input = pick_notes(&vk, available.clone(), 70);
assert_eq!(input, available);
let input_notes = pick_notes(&vk, all_notes.clone(), 70);
assert_eq!(input_notes, all_notes);

// out of balance basic check
let leaf = gen_note_leaf(&mut rng, true, &pk, 100);
let available = NoteList::from(vec![(n, leaf)]);
assert!(pick_notes(&vk, available, 101).is_empty());
let all_notes = generate_note_list([leaf], &sk);
assert!(pick_notes(&vk, all_notes, 101).is_empty());

// multiple inputs check
// note: this test is checking a naive, simple order-based output
// Check multiple input-notes

let leaf = [
// checking a naive, simple order-based output works
let leaves = [
gen_note_leaf(&mut rng, true, &pk, 100),
gen_note_leaf(&mut rng, true, &pk, 500),
gen_note_leaf(&mut rng, true, &pk, 300),
];
let all_notes = generate_note_list(leaves, &sk);
assert_eq!(pick_notes(&vk, all_notes.clone(), 600), all_notes);

let available: Vec<(_, _)> = leaf
.iter()
.map(|l| {
let n = l.note.gen_nullifier(&sk);
(n, l.clone())
})
.collect();

let available = NoteList::from(available);

assert_eq!(pick_notes(&vk, available.clone(), 600), available);

let leaf = [
// checking that spending more than the total doesn't work
let leaves = [
gen_note_leaf(&mut rng, true, &pk, 100),
gen_note_leaf(&mut rng, true, &pk, 500),
gen_note_leaf(&mut rng, true, &pk, 300),
];
let all_notes = generate_note_list(leaves, &sk);
assert_eq!(pick_notes(&vk, all_notes, 901), NoteList::default());

// checking that pick_notes works if spendable is smaller that the total

// generate 5 notes with 10 dusk owned by the same key
let leaves = [
gen_note_leaf(&mut rng, true, &pk, 10),
gen_note_leaf(&mut rng, true, &pk, 10),
gen_note_leaf(&mut rng, true, &pk, 10),
gen_note_leaf(&mut rng, true, &pk, 10),
gen_note_leaf(&mut rng, true, &pk, 20),
];
let all_notes = generate_note_list(leaves, &sk);

// with a target value of 20 it should pick the first 4 notes
let target = 20;
let expected_input_notes =
generate_expected_input_notes(&all_notes, &[0, 1, 2, 3]);
assert_eq!(
pick_notes(&vk, all_notes.clone(), target),
expected_input_notes
);

let available: Vec<(_, _)> = leaf
// with a target value of 50 it should also pick the last note
let target = 50;
let expected_input_notes =
generate_expected_input_notes(&all_notes, &[0, 1, 2, 4]);
assert_eq!(
pick_notes(&vk, all_notes.clone(), target),
expected_input_notes
);

// a target value of 51 however is above the max spendable
let target = 51;
let expected_input_notes = NoteList::default();
assert_eq!(
pick_notes(&vk, all_notes.clone(), target),
expected_input_notes
);
}

fn generate_expected_input_notes(
ordered_notes: &NoteList,
expected_indices: &[usize],
) -> NoteList {
let expected_input: Vec<(_, _)> = ordered_notes
.iter()
.map(|l| {
let n = l.note.gen_nullifier(&sk);
(n, l.clone())
.enumerate()
.filter_map(|(i, note_leaf)| {
if expected_indices.contains(&i) {
Some(note_leaf.clone())
} else {
None
}
})
.collect();
expected_input.into()
}

let available = NoteList::from(available);

assert_eq!(pick_notes(&vk, available, 901), NoteList::default());
fn generate_note_list(
note_leaves: impl AsRef<[NoteLeaf]>,
sk: &PhoenixSecretKey,
) -> NoteList {
note_leaves
.as_ref()
.iter()
.map(|leaf| {
let nullifier = leaf.note.gen_nullifier(sk);
(nullifier, leaf.clone())
})
.collect::<Vec<(_, _)>>()
.into()
}

0 comments on commit 4bb2982

Please sign in to comment.