Skip to content

Commit

Permalink
refactor: use shared immutable for slow tree (#4831)
Browse files Browse the repository at this point in the history
Fixes #4820
  • Loading branch information
LHerskind authored Feb 28, 2024
1 parent b7c2bc0 commit 821c25d
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 42 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,7 @@ contract TokenBlacklist {
use dep::aztec::protocol_types::{abis::function_selector::FunctionSelector, address::AztecAddress};
use dep::aztec::{
note::{note_getter_options::NoteGetterOptions, note_header::NoteHeader, utils as note_utils},
context::{PrivateContext, PublicContext, Context}, hash::{compute_secret_hash},
state_vars::{Map, PublicMutable, PrivateSet, PrivateImmutable}
hash::{compute_secret_hash}, state_vars::{Map, PublicMutable, PrivateSet, SharedImmutable}
};

use dep::field_note::field_note::FieldNote;
Expand All @@ -38,15 +37,12 @@ contract TokenBlacklist {
total_supply: PublicMutable<U128>,
pending_shields: PrivateSet<TransparentNote>,
public_balances: Map<AztecAddress, PublicMutable<U128>>,
slow_update: PrivateImmutable<FieldNote>,
public_slow_update: PublicMutable<AztecAddress>,
slow_update: SharedImmutable<AztecAddress>,
}

// docs:start:constructor
#[aztec(private)]
fn constructor(admin: AztecAddress, slow_updates_contract: AztecAddress) {
let mut slow_note = FieldNote::new(slow_updates_contract.to_field());
storage.slow_update.initialize(&mut slow_note, false);
// docs:end:constructor
let selector = FunctionSelector::from_signature("_initialize((Field),(Field))");
context.call_public_function(
Expand All @@ -55,18 +51,15 @@ contract TokenBlacklist {
[admin.to_field(), slow_updates_contract.to_field()]
);
// We cannot do the following atm
// let roles = UserFlags{is_admin: true, is_minter: false, is_blacklisted: false}.get_value() as Field;
// let roles = UserFlags { is_admin: true, is_minter: false, is_blacklisted: false }.get_value().to_field();
// SlowMap::at(slow_updates_contract).update_at_private(&mut context, admin.to_field(), roles);
}

////////
// Looking ugly because we cannot do constructor -> private calls

#[aztec(private)]
fn init_slow_tree(user: AztecAddress) {
let roles = UserFlags { is_admin: true, is_minter: false, is_blacklisted: false }.get_value().to_field();
// docs:start:get_and_update_private
let slow = SlowMap::at(AztecAddress::from_field(storage.slow_update.get_note().value));
let slow = SlowMap::at(storage.slow_update.read_private());
slow.update_at_private(&mut context, user.to_field(), roles);
// docs:end:get_and_update_private
context.call_public_function(
Expand All @@ -87,7 +80,7 @@ contract TokenBlacklist {
assert(!new_admin.is_zero(), "invalid admin");
storage.admin.write(new_admin);
// docs:start:write_slow_update_public
storage.public_slow_update.write(slow_updates_contract);
storage.slow_update.initialize(slow_updates_contract);
// docs:end:write_slow_update_public
// docs:start:slowmap_initialize
SlowMap::at(slow_updates_contract).initialize(context);
Expand All @@ -97,7 +90,7 @@ contract TokenBlacklist {
#[aztec(private)]
fn update_roles(user: AztecAddress, roles: Field) {
// docs:start:slowmap_at
let slow = SlowMap::at(AztecAddress::from_field(storage.slow_update.get_note().value));
let slow = SlowMap::at(storage.slow_update.read_private());
// docs:end:slowmap_at
let caller_roles = UserFlags::new(U128::from_integer(slow.read_at(&mut context, context.msg_sender().to_field())));
assert(caller_roles.is_admin, "caller is not admin");
Expand All @@ -108,7 +101,7 @@ contract TokenBlacklist {
#[aztec(public)]
fn mint_public(to: AztecAddress, amount: Field) {
// docs:start:get_public
let slow = SlowMap::at(storage.public_slow_update.read());
let slow = SlowMap::at(storage.slow_update.read_public());
// docs:end:get_public
// docs:start:read_at_pub
let to_roles = UserFlags::new(U128::from_integer(slow.read_at_pub(context, to.to_field())));
Expand All @@ -128,7 +121,7 @@ contract TokenBlacklist {

#[aztec(public)]
fn mint_private(amount: Field, secret_hash: Field) {
let slow = SlowMap::at(storage.public_slow_update.read());
let slow = SlowMap::at(storage.slow_update.read_public());
let caller_roles = UserFlags::new(U128::from_integer(slow.read_at_pub(context, context.msg_sender().to_field())));
assert(caller_roles.is_minter, "caller is not minter");

Expand All @@ -142,7 +135,7 @@ contract TokenBlacklist {

#[aztec(public)]
fn shield(from: AztecAddress, amount: Field, secret_hash: Field, nonce: Field) {
let slow = SlowMap::at(storage.public_slow_update.read());
let slow = SlowMap::at(storage.slow_update.read_public());
let from_roles = UserFlags::new(U128::from_integer(slow.read_at_pub(context, from.to_field())));
assert(!from_roles.is_blacklisted, "Blacklisted: Sender");

Expand All @@ -165,7 +158,7 @@ contract TokenBlacklist {

#[aztec(public)]
fn transfer_public(from: AztecAddress, to: AztecAddress, amount: Field, nonce: Field) {
let slow = SlowMap::at(storage.public_slow_update.read());
let slow = SlowMap::at(storage.slow_update.read_public());
let from_roles = UserFlags::new(U128::from_integer(slow.read_at_pub(context, from.to_field())));
assert(!from_roles.is_blacklisted, "Blacklisted: Sender");
let to_roles = UserFlags::new(U128::from_integer(slow.read_at_pub(context, to.to_field())));
Expand All @@ -187,7 +180,7 @@ contract TokenBlacklist {

#[aztec(public)]
fn burn_public(from: AztecAddress, amount: Field, nonce: Field) {
let slow = SlowMap::at(storage.public_slow_update.read());
let slow = SlowMap::at(storage.slow_update.read_public());
let from_roles = UserFlags::new(U128::from_integer(slow.read_at_pub(context, from.to_field())));
assert(!from_roles.is_blacklisted, "Blacklisted: Sender");

Expand All @@ -207,7 +200,7 @@ contract TokenBlacklist {

#[aztec(private)]
fn redeem_shield(to: AztecAddress, amount: Field, secret: Field) {
let slow = SlowMap::at(AztecAddress::from_field(storage.slow_update.get_note().value));
let slow = SlowMap::at(storage.slow_update.read_private());
// docs:start:slowmap_read_at
let to_roles = UserFlags::new(U128::from_integer(slow.read_at(&mut context, to.to_field())));
// docs:end:slowmap_read_at
Expand All @@ -229,7 +222,7 @@ contract TokenBlacklist {

#[aztec(private)]
fn unshield(from: AztecAddress, to: AztecAddress, amount: Field, nonce: Field) {
let slow = SlowMap::at(AztecAddress::from_field(storage.slow_update.get_note().value));
let slow = SlowMap::at(storage.slow_update.read_private());
let from_roles = UserFlags::new(U128::from_integer(slow.read_at(&mut context, from.to_field())));
assert(!from_roles.is_blacklisted, "Blacklisted: Sender");
let to_roles = UserFlags::new(U128::from_integer(slow.read_at(&mut context, to.to_field())));
Expand All @@ -250,7 +243,7 @@ contract TokenBlacklist {
// docs:start:transfer_private
#[aztec(private)]
fn transfer(from: AztecAddress, to: AztecAddress, amount: Field, nonce: Field) {
let slow = SlowMap::at(AztecAddress::from_field(storage.slow_update.get_note().value));
let slow = SlowMap::at(storage.slow_update.read_private());
let from_roles = UserFlags::new(U128::from_integer(slow.read_at(&mut context, from.to_field())));
assert(!from_roles.is_blacklisted, "Blacklisted: Sender");
let to_roles = UserFlags::new(U128::from_integer(slow.read_at(&mut context, to.to_field())));
Expand All @@ -270,7 +263,7 @@ contract TokenBlacklist {

#[aztec(private)]
fn burn(from: AztecAddress, amount: Field, nonce: Field) {
let slow = SlowMap::at(AztecAddress::from_field(storage.slow_update.get_note().value));
let slow = SlowMap::at(storage.slow_update.read_private());
let from_roles = UserFlags::new(U128::from_integer(slow.read_at(&mut context, from.to_field())));
assert(!from_roles.is_blacklisted, "Blacklisted: Sender");

Expand Down
24 changes: 4 additions & 20 deletions yarn-project/end-to-end/src/e2e_blacklist_token_contract.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import { BITSIZE_TOO_BIG_ERROR, U128_OVERFLOW_ERROR, U128_UNDERFLOW_ERROR } from
import { publicDeployAccounts, setup } from './fixtures/utils.js';
import { TokenSimulator } from './simulators/token_simulator.js';

const TIMEOUT = 90_000;
const TIMEOUT = 120_000;

describe('e2e_blacklist_token_contract', () => {
jest.setTimeout(TIMEOUT);
Expand Down Expand Up @@ -116,29 +116,13 @@ describe('e2e_blacklist_token_contract', () => {
const depth = 254;
slowUpdateTreeSimulator = await newTree(SparseTree, openTmpStore(), new Pedersen(), 'test', depth);

// Add account[0] as admin
await updateSlowTree(slowUpdateTreeSimulator, wallets[0], accounts[0].address, 4n);

const deployTx = TokenBlacklistContract.deploy(wallets[0], accounts[0], slowTree.address).send({});
const receipt = await deployTx.wait();
asset = receipt.contract;

// Add the note
const note = new Note([slowTree.address.toField()]);
const storageSlot = new Fr(6);
const noteTypeId = new Fr(7010510110810078111116101n); // FieldNote

for (const wallet of wallets) {
const extendedNote = new ExtendedNote(
note,
wallet.getCompleteAddress().address,
asset.address,
storageSlot,
noteTypeId,
receipt.txHash,
);
await wallet.addNote(extendedNote);
}

// Add account[0] as admin
await updateSlowTree(slowUpdateTreeSimulator, wallets[0], accounts[0].address, 4n);
await asset.methods.init_slow_tree(accounts[0].address).send().wait();

// Progress to next "epoch"
Expand Down

0 comments on commit 821c25d

Please sign in to comment.