make the airdrop contract up to 128x more storage-efficient

EkuboProtocol · Mar 3, 2024 · cbfbbf5 · cbfbbf5
1 parent 465257b
commit cbfbbf5
Show file tree

Hide file tree

Showing 4 changed files with 181 additions and 25 deletions.
diff --git a/src/airdrop.cairo b/src/airdrop.cairo
@@ -4,31 +4,37 @@ 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<TStorage> {
-    // Claims the given allotment of tokens
-    fn claim(ref self: TStorage, claim: Claim, proof: Array<felt252>);
-
     // 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<felt252>);
+
+    // Return whether the claim with the given ID has been claimed
+    fn is_claimed(self: @TStorage, claim_id: u64) -> bool;
 }
 
 #[starknet::contract]
 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<X, +Copy<X>, +Into<X, u256>>(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<felt252, bool>,
+        claimed_bitmap: LegacyMap<u64, u128>,
     }
 
     #[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<ContractState> {
+        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::<felt252>) {
-            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
@@ -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
@@ -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
@@ -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
+        }
+    }
+}