Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: rework balances map #8127

Merged
merged 1 commit into from
Aug 27, 2024
+34 −50
Merged

refactor: rework balances map #8127

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 16 additions & 16 deletions noir-projects/noir-contracts/contracts/token_contract/src/main.nr
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ contract Token {

use crate::types::{
transparent_note::TransparentNote,
token_note::{TokenNote, TOKEN_NOTE_LEN, TokenNoteHidingPoint}, balances_map::BalancesMap
token_note::{TokenNote, TOKEN_NOTE_LEN, TokenNoteHidingPoint}, balance_set::BalanceSet
};
// docs:end::imports

Expand Down Expand Up @@ -64,7 +64,7 @@ contract Token {
minters: Map<AztecAddress, PublicMutable<bool>>,
// docs:end:storage_minters
// docs:start:storage_balances
balances: BalancesMap<TokenNote>,
balances: Map<AztecAddress, BalanceSet<TokenNote>>,
// docs:end:storage_balances
total_supply: PublicMutable<U128>,
// docs:start:storage_pending_shields
Expand Down Expand Up @@ -223,7 +223,7 @@ contract Token {
fn privately_mint_private_note(amount: Field) {
let caller = context.msg_sender();
let caller_keys = get_current_public_keys(&mut context, caller);
storage.balances.add(caller, caller_keys.npk_m, U128::from_integer(amount)).emit(
storage.balances.at(caller).add(caller_keys.npk_m, U128::from_integer(amount)).emit(
encode_and_encrypt_note_with_keys(&mut context, caller_keys.ovpk_m, caller_keys.ivpk_m, caller)
);

Expand Down Expand Up @@ -320,7 +320,7 @@ contract Token {
let from = context.msg_sender();
let from_keys = get_current_public_keys(&mut context, from);
let to_keys = get_current_public_keys(&mut context, to);
storage.balances.add(to, to_keys.npk_m, U128::from_integer(amount)).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, to_keys.ivpk_m, to));
storage.balances.at(to).add(to_keys.npk_m, U128::from_integer(amount)).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, to_keys.ivpk_m, to));
}
// docs:end:redeem_shield

Expand All @@ -334,7 +334,7 @@ contract Token {
}

let from_keys = get_current_public_keys(&mut context, from);
storage.balances.sub(from, from_keys.npk_m, U128::from_integer(amount)).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, from_keys.ivpk_m, from));
storage.balances.at(from).sub(from_keys.npk_m, U128::from_integer(amount)).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, from_keys.ivpk_m, from));

Token::at(context.this_address())._increase_public_balance(to, amount).enqueue(&mut context);
}
Expand Down Expand Up @@ -364,11 +364,11 @@ contract Token {
INITIAL_TRANSFER_CALL_MAX_NOTES
);

storage.balances.add(from, from_keys.npk_m, change).emit(
storage.balances.at(from).add(from_keys.npk_m, change).emit(
encode_and_encrypt_note_with_keys_unconstrained(&mut context, from_keys.ovpk_m, from_keys.ivpk_m, from)
);

storage.balances.add(to, to_keys.npk_m, amount).emit(
storage.balances.at(to).add(to_keys.npk_m, amount).emit(
encode_and_encrypt_note_with_keys_unconstrained(&mut context, from_keys.ovpk_m, to_keys.ivpk_m, to)
);

Expand All @@ -390,7 +390,7 @@ contract Token {
amount: U128,
max_notes: u32
) -> U128 {
let subtracted = storage.balances.try_sub(account, amount, max_notes);
let subtracted = storage.balances.at(account).try_sub(amount, max_notes);

// Failing to subtract any amount means that the owner was unable to produce more notes that could be nullified.
// We could in some cases fail early inside try_sub if we detected that fewer notes than the maximum were
Expand Down Expand Up @@ -465,10 +465,10 @@ contract Token {
let amount = U128::from_integer(amount);
// docs:start:increase_private_balance
// docs:start:encrypted
storage.balances.sub(from, from_keys.npk_m, amount).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, from_keys.ivpk_m, from));
storage.balances.at(from).sub(from_keys.npk_m, amount).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, from_keys.ivpk_m, from));
// docs:end:encrypted
// docs:end:increase_private_balance
storage.balances.add(to, to_keys.npk_m, amount).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, to_keys.ivpk_m, to));
storage.balances.at(to).add(to_keys.npk_m, amount).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, to_keys.ivpk_m, to));
}
// docs:end:transfer_from

Expand All @@ -482,7 +482,7 @@ contract Token {
}

let from_keys = get_current_public_keys(&mut context, from);
storage.balances.sub(from, from_keys.npk_m, U128::from_integer(amount)).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, from_keys.ivpk_m, from));
storage.balances.at(from).sub(from_keys.npk_m, U128::from_integer(amount)).emit(encode_and_encrypt_note_with_keys(&mut context, from_keys.ovpk_m, from_keys.ivpk_m, from));

Token::at(context.this_address())._reduce_total_supply(amount).enqueue(&mut context);
}
Expand Down Expand Up @@ -521,7 +521,7 @@ contract Token {
let user_npk_m_hash = user_keys.npk_m.hash();

// 3. Deduct the funded amount from the user's balance - this is a maximum fee a user is willing to pay
// (called fee limit in aztec spec). The difference between fee limit and the actual tx fee will be refunded
// (called fee limit in aztec spec). The difference between fee limit and the actual tx fee will be refunded
// to the user in the `complete_refund(...)` function.
let change = subtract_balance(
&mut context,
Expand All @@ -530,7 +530,7 @@ contract Token {
U128::from_integer(funded_amount),
INITIAL_TRANSFER_CALL_MAX_NOTES
);
storage.balances.add(user, user_keys.npk_m, change).emit(
storage.balances.at(user).add(user_keys.npk_m, change).emit(
encode_and_encrypt_note_with_keys_unconstrained(&mut context, user_keys.ovpk_m, user_keys.ivpk_m, user)
);

Expand All @@ -540,7 +540,7 @@ contract Token {
header: NoteHeader {
contract_address: AztecAddress::zero(),
nonce: 0,
storage_slot: storage.balances.map.at(fee_payer).storage_slot,
storage_slot: storage.balances.at(fee_payer).set.storage_slot,
note_hash_counter: 0
},
amount: U128::zero(),
Expand All @@ -551,7 +551,7 @@ contract Token {
header: NoteHeader {
contract_address: AztecAddress::zero(),
nonce: 0,
storage_slot: storage.balances.map.at(user).storage_slot,
storage_slot: storage.balances.at(user).set.storage_slot,
note_hash_counter: 0
},
amount: U128::zero(),
Expand Down Expand Up @@ -640,7 +640,7 @@ contract Token {
// docs:start:balance_of_private
unconstrained fn balance_of_private(owner: AztecAddress) -> pub Field {
storage.balances.balance_of(owner).to_field()
storage.balances.at(owner).balance_of().to_field()
}
// docs:end:balance_of_private
}
Expand Down
2 changes: 1 addition & 1 deletion noir-projects/noir-contracts/contracts/token_contract/src/types.nr
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
mod transparent_note;
mod balances_map;
mod balance_set;
mod token_note;
50 changes: 17 additions & 33 deletions .../token_contract/src/types/balances_map.nr → ...s/token_contract/src/types/balance_set.nr
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,67 +1,53 @@
use dep::aztec::prelude::{AztecAddress, NoteGetterOptions, NoteViewerOptions, NoteHeader, NoteInterface, PrivateSet, Map, Point};
use dep::aztec::prelude::{NoteGetterOptions, NoteViewerOptions, NoteInterface, PrivateSet, Point};
use dep::aztec::{
context::{PrivateContext, UnconstrainedContext},
protocol_types::constants::MAX_NOTE_HASH_READ_REQUESTS_PER_CALL,
note::{
note_getter::view_notes, note_getter_options::SortOrder,
note_emission::{NoteEmission, OuterNoteEmission}
},
note::{note_getter::view_notes, note_emission::{NoteEmission, OuterNoteEmission}},
keys::{getters::get_current_public_keys, public_keys::NpkM}
};
use crate::types::{token_note::{TokenNote, OwnedNote}};
use crate::types::token_note::OwnedNote;

struct BalancesMap<T, Context> {
map: Map<AztecAddress, PrivateSet<T, Context>, Context>
struct BalanceSet<T, Context> {
set: PrivateSet<T, Context>,
}

impl<T, Context> BalancesMap<T, Context> {
impl<T, Context> BalanceSet<T, Context> {
pub fn new(context: Context, storage_slot: Field) -> Self {
assert(storage_slot != 0, "Storage slot 0 not allowed. Storage slots must start from 1.");
Self {
map: Map::new(
context,
storage_slot,
|context, slot| PrivateSet::new(context, slot)
)
}
Self { set: PrivateSet::new(context, storage_slot) }
}
}

impl<T> BalancesMap<T, UnconstrainedContext> {
unconstrained pub fn balance_of<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN>(
self: Self,
owner: AztecAddress
) -> U128 where T: NoteInterface<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN> + OwnedNote {
self.balance_of_with_offset(owner, 0)
impl<T> BalanceSet<T, UnconstrainedContext> {
unconstrained pub fn balance_of<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN>(self: Self) -> U128 where T: NoteInterface<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN> + OwnedNote {
self.balance_of_with_offset(0)
}

unconstrained pub fn balance_of_with_offset<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN>(
self: Self,
owner: AztecAddress,
offset: u32
) -> U128 where T: NoteInterface<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN> + OwnedNote {
let mut balance = U128::from_integer(0);
// docs:start:view_notes
let mut options = NoteViewerOptions::new();
let notes = self.map.at(owner).view_notes(options.set_offset(offset));
let notes = self.set.view_notes(options.set_offset(offset));
// docs:end:view_notes
for i in 0..options.limit {
if i < notes.len() {
balance = balance + notes.get_unchecked(i).get_amount();
}
}
if (notes.len() == options.limit) {
balance = balance + self.balance_of_with_offset(owner, offset + options.limit);
balance = balance + self.balance_of_with_offset(offset + options.limit);
}

balance
}
}

impl<T> BalancesMap<T, &mut PrivateContext> {
impl<T> BalanceSet<T, &mut PrivateContext> {
pub fn add<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN>(
self: Self,
owner: AztecAddress,
owner_npk_m: NpkM,
addend: U128
) -> OuterNoteEmission<T> where T: NoteInterface<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN> + OwnedNote + Eq {
Expand All @@ -72,23 +58,22 @@ impl<T> BalancesMap<T, &mut PrivateContext> {
let mut addend_note = T::new(addend, owner_npk_m.hash());

// docs:start:insert
OuterNoteEmission::new(Option::some(self.map.at(owner).insert(&mut addend_note)))
OuterNoteEmission::new(Option::some(self.set.insert(&mut addend_note)))
// docs:end:insert
}
}

pub fn sub<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN>(
self: Self,
owner: AztecAddress,
owner_npk_m: NpkM,
amount: U128
) -> OuterNoteEmission<T> where T: NoteInterface<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN> + OwnedNote + Eq {
let subtracted = self.try_sub(owner, amount, MAX_NOTE_HASH_READ_REQUESTS_PER_CALL);
let subtracted = self.try_sub(amount, MAX_NOTE_HASH_READ_REQUESTS_PER_CALL);

// try_sub may have substracted more or less than amount. We must ensure that we subtracted at least as much as
// we needed, and then create a new note for the owner for the change (if any).
assert(subtracted >= amount, "Balance too low");
self.add(owner, owner_npk_m, subtracted - amount)
self.add(owner_npk_m, subtracted - amount)
}

// Attempts to remove 'target_amount' from the owner's balance. try_sub returns how much was actually subtracted
Expand All @@ -102,7 +87,6 @@ impl<T> BalancesMap<T, &mut PrivateContext> {
// `try_sub` subtracting an amount smaller than `target_amount`.
pub fn try_sub<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN>(
self: Self,
owner: AztecAddress,
target_amount: U128,
max_notes: u32
) -> U128 where T: NoteInterface<T_SERIALIZED_LEN, T_SERIALIZED_BYTES_LEN> + OwnedNote + Eq {
Expand All @@ -112,7 +96,7 @@ impl<T> BalancesMap<T, &mut PrivateContext> {
// might result in is simply higher DA costs due to more nullifiers being emitted. Since we don't care
// about proving optimal note usage, we can save these constraints and make the circuit smaller.
let options = NoteGetterOptions::with_preprocessor(preprocess_notes_min_sum, target_amount).set_limit(max_notes);
let notes = self.map.at(owner).pop_notes(options);
let notes = self.set.pop_notes(options);

let mut subtracted = U128::from_integer(0);
for i in 0..options.limit {
Expand Down
Loading