diff --git a/src/airdrop.cairo b/src/airdrop.cairo index f7c3a2c..b01f2e0 100644 --- a/src/airdrop.cairo +++ b/src/airdrop.cairo @@ -4,23 +4,27 @@ use starknet::{ContractAddress}; #[derive(Copy, Drop, Serde, Hash, PartialEq)] pub struct Claim { + // the unique ID of the claim + pub id: u64, + // the address that will receive the token pub claimee: ContractAddress, + // the amount of token the address is entitled to pub amount: u128, } #[starknet::interface] pub trait IAirdrop { - // Claims the given allotment of tokens - fn claim(ref self: TStorage, claim: Claim, proof: Array); - // Return the root of the airdrop fn get_root(self: @TStorage) -> felt252; // Return the token being dropped fn get_token(self: @TStorage) -> IERC20Dispatcher; - // Return whether the claim has been claimed (always false for invalid claims) - fn is_claimed(self: @TStorage, claim: Claim) -> bool; + // Claims the given allotment of tokens + fn claim(ref self: TStorage, claim: Claim, proof: Array); + + // Return whether the claim with the given ID has been claimed + fn is_claimed(self: @TStorage, claim_id: u64) -> bool; } #[starknet::contract] @@ -28,7 +32,9 @@ pub mod Airdrop { use core::array::{ArrayTrait, SpanTrait}; use core::hash::{LegacyHash}; use governance::interfaces::erc20::{IERC20DispatcherTrait}; + use governance::utils::exp2::{exp2}; use super::{IAirdrop, ContractAddress, Claim, IERC20Dispatcher}; + use core::num::traits::zero::{Zero}; pub(crate) fn lt, +Into>(lhs: @X, rhs: @X) -> bool { let a: u256 = (*lhs).into(); @@ -57,7 +63,7 @@ pub mod Airdrop { struct Storage { root: felt252, token: IERC20Dispatcher, - claimed: LegacyMap, + claimed_bitmap: LegacyMap, } #[derive(Drop, starknet::Event)] @@ -77,30 +83,43 @@ pub mod Airdrop { self.token.write(token); } + #[inline(always)] + fn claim_id_to_bitmap_index(claim_id: u64) -> (u64, u8) { + let (word, index) = DivRem::div_rem(claim_id, 128_u64.try_into().unwrap()); + (word, index.try_into().unwrap()) + } + #[abi(embed_v0)] impl AirdropImpl of IAirdrop { + fn get_root(self: @ContractState) -> felt252 { + self.root.read() + } + + fn get_token(self: @ContractState) -> IERC20Dispatcher { + self.token.read() + } + fn claim(ref self: ContractState, claim: Claim, proof: Array::) { - let leaf = LegacyHash::hash(0, claim); + assert(!self.is_claimed(claim.id), 'ALREADY_CLAIMED'); - assert(!self.claimed.read(leaf), 'ALREADY_CLAIMED'); + let leaf = LegacyHash::hash(0, claim); assert(self.root.read() == compute_pedersen_root(leaf, proof.span()), 'INVALID_PROOF'); - self.claimed.write(leaf, true); + + let (word, index) = claim_id_to_bitmap_index(claim.id); + let bitmap = self.claimed_bitmap.read(word); + + self.claimed_bitmap.write(word, bitmap | exp2(index.try_into().unwrap())); self.token.read().transfer(claim.claimee, u256 { high: 0, low: claim.amount }); self.emit(Claimed { claim }); } - fn get_root(self: @ContractState) -> felt252 { - self.root.read() - } - - fn get_token(self: @ContractState) -> IERC20Dispatcher { - self.token.read() - } - fn is_claimed(self: @ContractState, claim: Claim) -> bool { - self.claimed.read(LegacyHash::hash(0, claim)) + 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); + (bitmap & exp2(index)).is_non_zero() } } } diff --git a/src/airdrop_test.cairo b/src/airdrop_test.cairo index 5f48608..f814168 100644 --- a/src/airdrop_test.cairo +++ b/src/airdrop_test.cairo @@ -85,7 +85,7 @@ fn test_compute_pedersen_root_recursive() { fn test_claim_single_recipient() { let (_, token) = deploy_token('AIRDROP', 'AD', 1234567); - let claim = Claim { claimee: contract_address_const::<2345>(), amount: 6789, }; + let claim = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, }; let leaf = LegacyHash::hash(0, claim); @@ -114,7 +114,7 @@ fn test_claim_single_recipient() { fn test_double_claim() { let (_, token) = deploy_token('AIRDROP', 'AD', 1234567); - let claim = Claim { claimee: contract_address_const::<2345>(), amount: 6789, }; + let claim = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, }; let leaf = LegacyHash::hash(0, claim); @@ -134,7 +134,7 @@ fn test_double_claim() { fn test_invalid_proof_single_entry() { let (_, token) = deploy_token('AIRDROP', 'AD', 1234567); - let claim = Claim { claimee: contract_address_const::<2345>(), amount: 6789, }; + let claim = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, }; let leaf = LegacyHash::hash(0, claim); @@ -150,7 +150,7 @@ fn test_invalid_proof_single_entry() { fn test_invalid_proof_fake_entry() { let (_, token) = deploy_token('AIRDROP', 'AD', 1234567); - let claim = Claim { claimee: contract_address_const::<2345>(), amount: 6789, }; + let claim = Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789, }; let leaf = LegacyHash::hash(0, claim); @@ -159,7 +159,7 @@ fn test_invalid_proof_fake_entry() { token.transfer(airdrop.contract_address, 6789); let proof = ArrayTrait::new(); - airdrop.claim(Claim { claimee: contract_address_const::<2345>(), amount: 6789 + 1, }, proof); + airdrop.claim(Claim { id: 0, claimee: contract_address_const::<2345>(), amount: 6789 + 1, }, proof); } @@ -168,8 +168,8 @@ fn test_invalid_proof_fake_entry() { fn test_claim_two_claims() { let (_, token) = deploy_token('AIRDROP', 'AD', 1234567); - let claim_a = Claim { claimee: contract_address_const::<2345>(), amount: 6789, }; - let claim_b = Claim { claimee: contract_address_const::<3456>(), amount: 789, }; + 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(0, claim_a); let leaf_b = LegacyHash::hash(0, claim_b); diff --git a/src/lib.cairo b/src/lib.cairo index b318bca..c391e4d 100644 --- a/src/lib.cairo +++ b/src/lib.cairo @@ -24,4 +24,5 @@ pub(crate) mod interfaces { } pub(crate) mod utils { pub(crate) mod timestamps; + pub(crate) mod exp2; } diff --git a/src/utils/exp2.cairo b/src/utils/exp2.cairo new file mode 100644 index 0000000..a08db59 --- /dev/null +++ b/src/utils/exp2.cairo @@ -0,0 +1,136 @@ +// Returns 2^n +pub fn exp2(n: u8) -> u128 { + match n { + 0 => { 0x1 }, + 1 => { 0x2 }, + 2 => { 0x4 }, + 3 => { 0x8 }, + 4 => { 0x10 }, + 5 => { 0x20 }, + 6 => { 0x40 }, + 7 => { 0x80 }, + 8 => { 0x100 }, + 9 => { 0x200 }, + 10 => { 0x400 }, + 11 => { 0x800 }, + 12 => { 0x1000 }, + 13 => { 0x2000 }, + 14 => { 0x4000 }, + 15 => { 0x8000 }, + 16 => { 0x10000 }, + 17 => { 0x20000 }, + 18 => { 0x40000 }, + 19 => { 0x80000 }, + 20 => { 0x100000 }, + 21 => { 0x200000 }, + 22 => { 0x400000 }, + 23 => { 0x800000 }, + 24 => { 0x1000000 }, + 25 => { 0x2000000 }, + 26 => { 0x4000000 }, + 27 => { 0x8000000 }, + 28 => { 0x10000000 }, + 29 => { 0x20000000 }, + 30 => { 0x40000000 }, + 31 => { 0x80000000 }, + 32 => { 0x100000000 }, + 33 => { 0x200000000 }, + 34 => { 0x400000000 }, + 35 => { 0x800000000 }, + 36 => { 0x1000000000 }, + 37 => { 0x2000000000 }, + 38 => { 0x4000000000 }, + 39 => { 0x8000000000 }, + 40 => { 0x10000000000 }, + 41 => { 0x20000000000 }, + 42 => { 0x40000000000 }, + 43 => { 0x80000000000 }, + 44 => { 0x100000000000 }, + 45 => { 0x200000000000 }, + 46 => { 0x400000000000 }, + 47 => { 0x800000000000 }, + 48 => { 0x1000000000000 }, + 49 => { 0x2000000000000 }, + 50 => { 0x4000000000000 }, + 51 => { 0x8000000000000 }, + 52 => { 0x10000000000000 }, + 53 => { 0x20000000000000 }, + 54 => { 0x40000000000000 }, + 55 => { 0x80000000000000 }, + 56 => { 0x100000000000000 }, + 57 => { 0x200000000000000 }, + 58 => { 0x400000000000000 }, + 59 => { 0x800000000000000 }, + 60 => { 0x1000000000000000 }, + 61 => { 0x2000000000000000 }, + 62 => { 0x4000000000000000 }, + 63 => { 0x8000000000000000 }, + 64 => { 0x10000000000000000 }, + 65 => { 0x20000000000000000 }, + 66 => { 0x40000000000000000 }, + 67 => { 0x80000000000000000 }, + 68 => { 0x100000000000000000 }, + 69 => { 0x200000000000000000 }, + 70 => { 0x400000000000000000 }, + 71 => { 0x800000000000000000 }, + 72 => { 0x1000000000000000000 }, + 73 => { 0x2000000000000000000 }, + 74 => { 0x4000000000000000000 }, + 75 => { 0x8000000000000000000 }, + 76 => { 0x10000000000000000000 }, + 77 => { 0x20000000000000000000 }, + 78 => { 0x40000000000000000000 }, + 79 => { 0x80000000000000000000 }, + 80 => { 0x100000000000000000000 }, + 81 => { 0x200000000000000000000 }, + 82 => { 0x400000000000000000000 }, + 83 => { 0x800000000000000000000 }, + 84 => { 0x1000000000000000000000 }, + 85 => { 0x2000000000000000000000 }, + 86 => { 0x4000000000000000000000 }, + 87 => { 0x8000000000000000000000 }, + 88 => { 0x10000000000000000000000 }, + 89 => { 0x20000000000000000000000 }, + 90 => { 0x40000000000000000000000 }, + 91 => { 0x80000000000000000000000 }, + 92 => { 0x100000000000000000000000 }, + 93 => { 0x200000000000000000000000 }, + 94 => { 0x400000000000000000000000 }, + 95 => { 0x800000000000000000000000 }, + 96 => { 0x1000000000000000000000000 }, + 97 => { 0x2000000000000000000000000 }, + 98 => { 0x4000000000000000000000000 }, + 99 => { 0x8000000000000000000000000 }, + 100 => { 0x10000000000000000000000000 }, + 101 => { 0x20000000000000000000000000 }, + 102 => { 0x40000000000000000000000000 }, + 103 => { 0x80000000000000000000000000 }, + 104 => { 0x100000000000000000000000000 }, + 105 => { 0x200000000000000000000000000 }, + 106 => { 0x400000000000000000000000000 }, + 107 => { 0x800000000000000000000000000 }, + 108 => { 0x1000000000000000000000000000 }, + 109 => { 0x2000000000000000000000000000 }, + 110 => { 0x4000000000000000000000000000 }, + 111 => { 0x8000000000000000000000000000 }, + 112 => { 0x10000000000000000000000000000 }, + 113 => { 0x20000000000000000000000000000 }, + 114 => { 0x40000000000000000000000000000 }, + 115 => { 0x80000000000000000000000000000 }, + 116 => { 0x100000000000000000000000000000 }, + 117 => { 0x200000000000000000000000000000 }, + 118 => { 0x400000000000000000000000000000 }, + 119 => { 0x800000000000000000000000000000 }, + 120 => { 0x1000000000000000000000000000000 }, + 121 => { 0x2000000000000000000000000000000 }, + 122 => { 0x4000000000000000000000000000000 }, + 123 => { 0x8000000000000000000000000000000 }, + 124 => { 0x10000000000000000000000000000000 }, + 125 => { 0x20000000000000000000000000000000 }, + 126 => { 0x40000000000000000000000000000000 }, + _ => { + assert(n == 127, 'exp2'); + 0x80000000000000000000000000000000 + } + } +}