Skip to content

Commit

Permalink
improve gas efficiency and add more unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@moodysalem
moodysalem committed Mar 4, 2024
1 parent 941cc25 commit 47c7ba7
Show file tree
Hide file tree
Showing 2 changed files with 185 additions and 19 deletions.
16 changes: 9 additions & 7 deletions src/airdrop.cairo
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use core::array::{Array};
use core::array::{Array,Span};
use governance::interfaces::erc20::{IERC20Dispatcher};
use starknet::{ContractAddress};

Expand All @@ -21,7 +21,7 @@ pub trait IAirdrop<TStorage> {
fn get_token(self: @TStorage) -> IERC20Dispatcher;

// Claims the given allotment of tokens
fn claim(ref self: TStorage, claim: Claim, proof: Array<felt252>);
fn claim(ref self: TStorage, claim: Claim, proof: Span<felt252>);

// Return whether the claim with the given ID has been claimed
fn is_claimed(self: @TStorage, claim_id: u64) -> bool;
Expand Down Expand Up @@ -99,18 +99,20 @@ pub mod Airdrop {
self.token.read()
}

fn claim(ref self: ContractState, claim: Claim, proof: Array::<felt252>) {
fn claim(ref self: ContractState, claim: Claim, proof: Span<felt252>) {
let leaf = LegacyHash::hash(selector!("ekubo::governance::airdrop::Claim"), claim);
assert(self.root.read() == compute_pedersen_root(leaf, proof.span()), 'INVALID_PROOF');

assert(!self.is_claimed(claim.id), 'ALREADY_CLAIMED');
assert(self.root.read() == compute_pedersen_root(leaf, proof), 'INVALID_PROOF');

// this is copied in from is_claimed because we only want to read the bitmap once
let (word, index) = claim_id_to_bitmap_index(claim.id);
let bitmap = self.claimed_bitmap.read(word);
let already_claimed = (bitmap & exp2(index)).is_non_zero();

assert(!already_claimed, 'ALREADY_CLAIMED');

self.claimed_bitmap.write(word, bitmap | exp2(index.try_into().unwrap()));

self.token.read().transfer(claim.claimee, u256 { high: 0, low: claim.amount });
self.token.read().transfer(claim.claimee, claim.amount.into());

self.emit(Claimed { claim });
}
Expand Down
188 changes: 176 additions & 12 deletions src/airdrop_test.cairo
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use core::array::{ArrayTrait};
use core::array::{ArrayTrait, SpanTrait};
use core::hash::{LegacyHash};
use core::option::{OptionTrait};

Expand Down Expand Up @@ -97,9 +97,8 @@ fn test_claim_single_recipient() {
let airdrop = deploy(token.contract_address, leaf);

token.transfer(airdrop.contract_address, 6789);
let proof = ArrayTrait::new();

airdrop.claim(claim, proof);
airdrop.claim(claim, array![].span());

let log = pop_log::<Airdrop::Claimed>(airdrop.contract_address).unwrap();
assert_eq!(log.claim, claim);
Expand All @@ -125,10 +124,8 @@ fn test_double_claim() {
let airdrop = deploy(token.contract_address, leaf);

token.transfer(airdrop.contract_address, 6789);
let mut proof = ArrayTrait::new();
airdrop.claim(claim, proof);
proof = ArrayTrait::new();
airdrop.claim(claim, proof);
airdrop.claim(claim, array![].span());
airdrop.claim(claim, array![].span());
}

#[test]
Expand All @@ -143,7 +140,7 @@ fn test_invalid_proof_single_entry() {
let airdrop = deploy(token.contract_address, leaf);

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

#[test]
Expand All @@ -158,11 +155,11 @@ fn test_invalid_proof_fake_entry() {
let airdrop = deploy(token.contract_address, leaf);

token.transfer(airdrop.contract_address, 6789);
let proof = ArrayTrait::new();

airdrop
.claim(
Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789 + 1, }, proof
Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789 + 1, },
array![].span()
);
}

Expand All @@ -187,11 +184,11 @@ fn test_claim_two_claims() {
let airdrop = deploy(token.contract_address, root);
token.transfer(airdrop.contract_address, 6789 + 789 + 1);

airdrop.claim(claim_a, array![leaf_b]);
airdrop.claim(claim_a, array![leaf_b].span());
assert_eq!(token.balance_of(airdrop.contract_address), (789 + 1));
assert_eq!(token.balance_of(claim_a.claimee), 6789);

airdrop.claim(claim_b, array![leaf_a]);
airdrop.claim(claim_b, array![leaf_a].span());
assert_eq!(token.balance_of(airdrop.contract_address), 1);
assert_eq!(token.balance_of(claim_b.claimee), 789);
}
Expand Down Expand Up @@ -238,6 +235,7 @@ fn test_claims_from_generated_tree() {
22944591007266013337629529054088070826740344136663051917181912077498206093,
2846046884061389749777735515205600989814522753032574962636562486677935396074
]
.span()
);

let log = pop_log::<Airdrop::Claimed>(airdrop.contract_address).unwrap();
Expand All @@ -260,5 +258,171 @@ fn test_claims_from_generated_tree() {
22944591007266013337629529054088070826740344136663051917181912077498206093,
2846046884061389749777735515205600989814522753032574962636562486677935396074
]
.span()
);
let log = pop_log::<Airdrop::Claimed>(airdrop.contract_address).unwrap();
assert_eq!(log.claim, claim_0);
}

#[test]
#[should_panic(expected: ('ALREADY_CLAIMED', 'ENTRYPOINT_FAILED'))]
fn test_double_claim_from_generated_tree() {
let claim_0 = Claim {
id: 0,
claimee: contract_address_const::<
1257981684727298919953780547925609938727371268283996697135018561811391002099
>(),
amount: 845608158412629999616,
};

let (_, token) = deploy_token('AIRDROP', 'AD', claim_0.amount.into());

let root = 2413984000256568988735068618807996871735886303454043475744972321149068137869;
let airdrop = deploy(token.contract_address, root);

token.transfer(airdrop.contract_address, claim_0.amount.into());

let proof: Span<felt252> = array![
390013443931943946052075510188945600544108471539235465760564815348896073043,
2591818886036301641799899841447556295494184204908229358406473782788431853617,
3433559452610196359109559589502585411529094342760420711041457728474879804685,
119111708719532621104568211251857481136318454621898627733025381039107349350,
1550418626007763899979956501892881046988353701960212721885621375458028218469,
218302537176435686946721821062002958322614343556723420712784506426080342216,
1753580693918376168416443301945093568141375497403576624304615426611458701443,
284161108154264923299661757093898525322488115499630822539338320558723810310,
3378969471732886394431481313236934101872088301949153794471811360320074526103,
2691963575009292057768595613759919396863463394980592564921927341908988940473,
22944591007266013337629529054088070826740344136663051917181912077498206093,
2846046884061389749777735515205600989814522753032574962636562486677935396074
]
.span();

airdrop.claim(claim_0, proof);
airdrop.claim(claim_0, proof);
}

#[test]
#[should_panic(expected: ('ALREADY_CLAIMED', 'ENTRYPOINT_FAILED'))]
fn test_double_claim_after_other_claim() {
let claim_0 = Claim {
id: 0,
claimee: contract_address_const::<
1257981684727298919953780547925609938727371268283996697135018561811391002099
>(),
amount: 845608158412629999616,
};

let claim_1 = Claim {
id: 1,
claimee: contract_address_const::<
2446484730111463702450186103350698828806903266085688038950964576824849476058
>(),
amount: 758639984742607224832,
};

let (_, token) = deploy_token('AIRDROP', 'AD', claim_0.amount.into() + claim_1.amount.into());

let root = 2413984000256568988735068618807996871735886303454043475744972321149068137869;
let airdrop = deploy(token.contract_address, root);

token.transfer(airdrop.contract_address, claim_0.amount.into() + claim_1.amount.into());

airdrop
.claim(
claim_1,
array![
2879705852068751339326970574743249357626496859246711485336045655175496222574,
2591818886036301641799899841447556295494184204908229358406473782788431853617,
3433559452610196359109559589502585411529094342760420711041457728474879804685,
119111708719532621104568211251857481136318454621898627733025381039107349350,
1550418626007763899979956501892881046988353701960212721885621375458028218469,
218302537176435686946721821062002958322614343556723420712784506426080342216,
1753580693918376168416443301945093568141375497403576624304615426611458701443,
284161108154264923299661757093898525322488115499630822539338320558723810310,
3378969471732886394431481313236934101872088301949153794471811360320074526103,
2691963575009292057768595613759919396863463394980592564921927341908988940473,
22944591007266013337629529054088070826740344136663051917181912077498206093,
2846046884061389749777735515205600989814522753032574962636562486677935396074
]
.span()
);

airdrop
.claim(
claim_0,
array![
390013443931943946052075510188945600544108471539235465760564815348896073043,
2591818886036301641799899841447556295494184204908229358406473782788431853617,
3433559452610196359109559589502585411529094342760420711041457728474879804685,
119111708719532621104568211251857481136318454621898627733025381039107349350,
1550418626007763899979956501892881046988353701960212721885621375458028218469,
218302537176435686946721821062002958322614343556723420712784506426080342216,
1753580693918376168416443301945093568141375497403576624304615426611458701443,
284161108154264923299661757093898525322488115499630822539338320558723810310,
3378969471732886394431481313236934101872088301949153794471811360320074526103,
2691963575009292057768595613759919396863463394980592564921927341908988940473,
22944591007266013337629529054088070826740344136663051917181912077498206093,
2846046884061389749777735515205600989814522753032574962636562486677935396074
]
.span()
);

// double claim of claim id 1
airdrop
.claim(
claim_1,
array![
2879705852068751339326970574743249357626496859246711485336045655175496222574,
2591818886036301641799899841447556295494184204908229358406473782788431853617,
3433559452610196359109559589502585411529094342760420711041457728474879804685,
119111708719532621104568211251857481136318454621898627733025381039107349350,
1550418626007763899979956501892881046988353701960212721885621375458028218469,
218302537176435686946721821062002958322614343556723420712784506426080342216,
1753580693918376168416443301945093568141375497403576624304615426611458701443,
284161108154264923299661757093898525322488115499630822539338320558723810310,
3378969471732886394431481313236934101872088301949153794471811360320074526103,
2691963575009292057768595613759919396863463394980592564921927341908988940473,
22944591007266013337629529054088070826740344136663051917181912077498206093,
2846046884061389749777735515205600989814522753032574962636562486677935396074
]
.span()
);
}

#[test]
#[should_panic(expected: ('TRANSFER_INSUFFICIENT_BALANCE', 'ENTRYPOINT_FAILED', 'ENTRYPOINT_FAILED'))]
fn test_claim_before_funded() {
let claim_0 = Claim {
id: 0,
claimee: contract_address_const::<
1257981684727298919953780547925609938727371268283996697135018561811391002099
>(),
amount: 845608158412629999616,
};

let (_, token) = deploy_token('AIRDROP', 'AD', 0);

let root = 2413984000256568988735068618807996871735886303454043475744972321149068137869;
let airdrop = deploy(token.contract_address, root);

airdrop
.claim(
claim_0,
array![
390013443931943946052075510188945600544108471539235465760564815348896073043,
2591818886036301641799899841447556295494184204908229358406473782788431853617,
3433559452610196359109559589502585411529094342760420711041457728474879804685,
119111708719532621104568211251857481136318454621898627733025381039107349350,
1550418626007763899979956501892881046988353701960212721885621375458028218469,
218302537176435686946721821062002958322614343556723420712784506426080342216,
1753580693918376168416443301945093568141375497403576624304615426611458701443,
284161108154264923299661757093898525322488115499630822539338320558723810310,
3378969471732886394431481313236934101872088301949153794471811360320074526103,
2691963575009292057768595613759919396863463394980592564921927341908988940473,
22944591007266013337629529054088070826740344136663051917181912077498206093,
2846046884061389749777735515205600989814522753032574962636562486677935396074
]
.span()
);
}

0 comments on commit 47c7ba7

Please sign in to comment.