Skip to content

Commit

Permalink
claim 128 at a time (#28)
Browse files Browse the repository at this point in the history 
* claim 128 at a time

* start adding some unit tests for the claim_128 function

* fix token re-entrance

* remove available gas annotations

* some refactoring

* address review comments

* add a unit test for claiming a large tree
  • Loading branch information
@moodysalem
moodysalem authored Mar 5, 2024
1 parent 20f3cc5 commit 797e290
Show file tree
Hide file tree
Showing 9 changed files with 201 additions and 66 deletions.
107 changes: 101 additions & 6 deletions src/airdrop.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,11 @@ pub trait IAirdrop<TStorage> {
// Panics if the proof is invalid.
fn claim(ref self: TStorage, claim: Claim, proof: Span<felt252>) -> bool;

// Claims the batch of up to 128 claims that must be aligned with a single bitmap, i.e. the id of the first must be a multiple of 128
// and the claims should be sequentially in order. The proof verification is optimized in this method.
// Returns the number of claims that were executed
fn claim_128(ref self: TStorage, claims: Span<Claim>, remaining_proof: Span<felt252>) -> u8;

// Return whether the claim with the given ID has been claimed
fn is_claimed(self: @TStorage, claim_id: u64) -> bool;
}
Expand All @@ -35,13 +40,13 @@ pub trait IAirdrop<TStorage> {
pub mod Airdrop {
use core::array::{ArrayTrait, SpanTrait};
use core::hash::{LegacyHash};
use core::num::traits::zero::{Zero};
use core::num::traits::one::{One};
use core::num::traits::zero::{Zero};
use governance::interfaces::erc20::{IERC20DispatcherTrait};
use governance::utils::exp2::{exp2};
use super::{IAirdrop, ContractAddress, Claim, IERC20Dispatcher};

#[inline(always)]

pub(crate) fn hash_function(a: felt252, b: felt252) -> felt252 {
let a_u256: u256 = a.into();
if a_u256 < b.into() {
Expand Down Expand Up @@ -85,17 +90,56 @@ pub mod Airdrop {
self.token.write(token);
}

#[inline(always)]
const BITMAP_SIZE: NonZero<u64> = 128;

fn claim_id_to_bitmap_index(claim_id: u64) -> (u64, u8) {
let (word, index) = DivRem::div_rem(claim_id, 128_u64.try_into().unwrap());
let (word, index) = DivRem::div_rem(claim_id, BITMAP_SIZE);
(word, index.try_into().unwrap())
}

#[inline(always)]
fn hash_claim(claim: Claim) -> felt252 {
pub fn hash_claim(claim: Claim) -> felt252 {
LegacyHash::hash(selector!("ekubo::governance::airdrop::Claim"), claim)
}

pub fn compute_root_of_group(mut claims: Span<Claim>) -> felt252 {
assert(!claims.is_empty(), 'NO_CLAIMS');

let mut claim_hashes: Array<felt252> = ArrayTrait::new();

let mut last_claim_id: Option<u64> = Option::None;

while let Option::Some(claim) = claims
.pop_front() {
if let Option::Some(last_id) = last_claim_id {
assert(last_id == (*claim.id - 1), 'SEQUENTIAL');
};

claim_hashes.append(hash_claim(*claim));
last_claim_id = Option::Some(*claim.id);
};

// will eventually contain an array of length 1
let mut current_layer: Span<felt252> = claim_hashes.span();

while current_layer
.len()
.is_non_one() {
let mut next_layer: Array<felt252> = ArrayTrait::new();

while let Option::Some(hash) = current_layer
.pop_front() {
next_layer
.append(
hash_function(*hash, *current_layer.pop_front().unwrap_or(hash))
);
};

current_layer = next_layer.span();
};

*current_layer.pop_front().unwrap()
}

#[abi(embed_v0)]
impl AirdropImpl of IAirdrop<ContractState> {
fn get_root(self: @ContractState) -> felt252 {
Expand Down Expand Up @@ -128,6 +172,57 @@ pub mod Airdrop {
}
}

fn claim_128(
ref self: ContractState, mut claims: Span<Claim>, remaining_proof: Span<felt252>
) -> u8 {
assert(claims.len() < 129, 'TOO_MANY_CLAIMS');

// groups that cross bitmap boundaries should just make multiple calls
// this code already reduces the number of pedersens in the verification by a factor of ~7
let (word, index_u64) = DivRem::div_rem(*claims.at(0).id, BITMAP_SIZE);
assert(index_u64 == 0, 'FIRST_CLAIM_MUST_BE_MULT_128');

let root_of_group = compute_root_of_group(claims);

assert(
self.root.read() == compute_pedersen_root(root_of_group, remaining_proof),
'INVALID_PROOF'
);

let mut bitmap = self.claimed_bitmap.read(word);

let mut index: u8 = 0;
let mut unclaimed: Array<Claim> = ArrayTrait::new();

while let Option::Some(claim) = claims
.pop_front() {
let already_claimed = (bitmap & exp2(index)).is_non_zero();

if !already_claimed {
bitmap = bitmap | exp2(index);
unclaimed.append(*claim);
}

index += 1;
};

self.claimed_bitmap.write(word, bitmap);

let num_claimed = unclaimed.len();

// the event emittance and transfers are separated from the above to prevent re-entrance
let token = self.token.read();

while let Option::Some(claim) = unclaimed
.pop_front() {
token.transfer(claim.claimee, claim.amount.into());
self.emit(Claimed { claim });
};

// never fails because we assert claims length at the beginning so we know it's less than 128
num_claimed.try_into().unwrap()
}

fn is_claimed(self: @ContractState, claim_id: u64) -> bool {
let (word, index) = claim_id_to_bitmap_index(claim_id);
let bitmap = self.claimed_bitmap.read(word);
Expand Down
109 changes: 96 additions & 13 deletions src/airdrop_test.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use core::result::{Result, ResultTrait};
use core::traits::{TryInto, Into};
use governance::airdrop::{
IAirdropDispatcher, IAirdropDispatcherTrait, Airdrop,
Airdrop::{compute_pedersen_root, hash_function}, Claim
Airdrop::{compute_pedersen_root, hash_function, hash_claim, compute_root_of_group}, Claim
};
use governance::governance_token::{
IGovernanceTokenDispatcherTrait, GovernanceToken, IGovernanceTokenDispatcher
Expand Down Expand Up @@ -41,10 +41,7 @@ fn test_selector() {
#[test]
fn test_hash() {
assert_eq!(
LegacyHash::hash(
selector!("ekubo::governance::airdrop::Claim"),
Claim { id: 123, claimee: contract_address_const::<456>(), amount: 789 }
),
hash_claim(Claim { id: 123, claimee: contract_address_const::<456>(), amount: 789 }),
0x0760b337026a91a6f2af99a0654f7fdff5d5c8d4e565277e787b99e17b1742a3
);
}
Expand Down Expand Up @@ -92,7 +89,7 @@ fn test_claim_single_recipient() {

let claim = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, };

let leaf = LegacyHash::hash(selector!("ekubo::governance::airdrop::Claim"), claim);
let leaf = hash_claim(claim);

let airdrop = deploy(token.contract_address, leaf);

Expand All @@ -112,13 +109,37 @@ fn test_claim_single_recipient() {
}

#[test]
#[available_gas(4000000)]
fn test_claim_128_single_recipient_tree() {
let (_, token) = deploy_token('AIRDROP', 'AD', 1234567);

let claim = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, };

let leaf = hash_claim(claim);

let airdrop = deploy(token.contract_address, leaf);

token.transfer(airdrop.contract_address, 6789);

assert_eq!(airdrop.claim_128(array![claim].span(), array![].span()), 1);

let log = pop_log::<Airdrop::Claimed>(airdrop.contract_address).unwrap();
assert_eq!(log.claim, claim);

pop_log::<GovernanceToken::Transfer>(token.contract_address).unwrap();
pop_log::<GovernanceToken::Transfer>(token.contract_address).unwrap();
let log = pop_log::<GovernanceToken::Transfer>(token.contract_address).unwrap();
assert_eq!(log.from, airdrop.contract_address);
assert_eq!(log.to, claim.claimee);
assert_eq!(log.value, claim.amount.into());
}

#[test]
fn test_double_claim() {
let (_, token) = deploy_token('AIRDROP', 'AD', 1234567);

let claim = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, };

let leaf = LegacyHash::hash(selector!("ekubo::governance::airdrop::Claim"), claim);
let leaf = hash_claim(claim);

let airdrop = deploy(token.contract_address, leaf);

Expand All @@ -127,14 +148,29 @@ fn test_double_claim() {
assert_eq!(airdrop.claim(claim, array![].span()), false);
}

#[test]
fn test_double_claim_128_single_recipient_tree() {
let (_, token) = deploy_token('AIRDROP', 'AD', 1234567);

let claim = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, };

let leaf = hash_claim(claim);

let airdrop = deploy(token.contract_address, leaf);

token.transfer(airdrop.contract_address, 6789);
assert_eq!(airdrop.claim_128(array![claim].span(), array![].span()), 1);
assert_eq!(airdrop.claim_128(array![claim].span(), array![].span()), 0);
}

#[test]
#[should_panic(expected: ('INVALID_PROOF', 'ENTRYPOINT_FAILED'))]
fn test_invalid_proof_single_entry() {
let (_, token) = deploy_token('AIRDROP', 'AD', 1234567);

let claim = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, };

let leaf = LegacyHash::hash(selector!("ekubo::governance::airdrop::Claim"), claim);
let leaf = hash_claim(claim);

let airdrop = deploy(token.contract_address, leaf);

Expand All @@ -149,7 +185,7 @@ fn test_invalid_proof_fake_entry() {

let claim = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, };

let leaf = LegacyHash::hash(selector!("ekubo::governance::airdrop::Claim"), claim);
let leaf = hash_claim(claim);

let airdrop = deploy(token.contract_address, leaf);

Expand All @@ -164,15 +200,14 @@ fn test_invalid_proof_fake_entry() {


#[test]
#[available_gas(30000000)]
fn test_claim_two_claims() {
let (_, token) = deploy_token('AIRDROP', 'AD', 1234567);

let claim_a = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, };
let claim_b = Claim { id: 1, claimee: contract_address_const::<3456>(), amount: 789, };

let leaf_a = LegacyHash::hash(selector!("ekubo::governance::airdrop::Claim"), claim_a);
let leaf_b = LegacyHash::hash(selector!("ekubo::governance::airdrop::Claim"), claim_b);
let leaf_a = hash_claim(claim_a);
let leaf_b = hash_claim(claim_b);

let root = hash_function(leaf_a, leaf_b);

Expand All @@ -188,6 +223,25 @@ fn test_claim_two_claims() {
assert_eq!(token.balance_of(claim_b.claimee), 789);
}

#[test]
fn test_claim_two_claims_via_claim_128() {
let (_, token) = deploy_token('AIRDROP', 'AD', 1234567);

let claim_a = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, };
let claim_b = Claim { id: 1, claimee: contract_address_const::<3456>(), amount: 789, };

let leaf_a = hash_claim(claim_a);
let leaf_b = hash_claim(claim_b);

let root = hash_function(leaf_a, leaf_b);

let airdrop = deploy(token.contract_address, root);
token.transfer(airdrop.contract_address, 6789 + 789 + 1);

assert_eq!(airdrop.claim_128(array![claim_a, claim_b].span(), array![].span()), 2);
assert_eq!(airdrop.claim_128(array![claim_a, claim_b].span(), array![].span()), 0);
}

fn test_claim_is_valid(root: felt252, claim: Claim, proof: Array<felt252>) {
let pspan = proof.span();
let (_, token) = deploy_token('AIRDROP', 'AD', claim.amount);
Expand Down Expand Up @@ -533,3 +587,32 @@ fn test_multiple_claims_from_generated_tree() {
let log = pop_log::<Airdrop::Claimed>(airdrop.contract_address).unwrap();
assert_eq!(log.claim, claim_0);
}


#[test]
fn test_claim_128_large_tree() {
let mut i: u64 = 0;

let mut claims: Array<Claim> = array![];

while (i < 320) {
claims.append(Claim { id: i, amount: 3, claimee: contract_address_const::<0xcdee>() });
i += 1;
};

let s1 = compute_root_of_group(claims.span().slice(0, 128));
let s2 = compute_root_of_group(claims.span().slice(128, 128));
let s3 = compute_root_of_group(claims.span().slice(256, 64));

let rl = hash_function(s1, s2);
let rr = hash_function(s3, s3);
let root = hash_function(rl, rr);

let (_, token) = deploy_token('AIRDROP', 'AD', 960);
let airdrop = deploy(token.contract_address, root);
token.transfer(airdrop.contract_address, 960);

assert_eq!(airdrop.claim_128(claims.span().slice(0, 128), array![s2, rr].span()), 128);
assert_eq!(airdrop.claim_128(claims.span().slice(128, 128), array![s1, rr].span()), 128);
assert_eq!(airdrop.claim_128(claims.span().slice(256, 64), array![s3, rl].span()), 64);
}
Loading

0 comments on commit 797e290

Please sign in to comment.