Implemented Merkle Airdrop Support for Initial Token Distribution

contracts/Scarb.toml

@@ -10,6 +10,7 @@ starknet = "2.3.1"
 
 # External dependencies
 openzeppelin = { git = "https://github.com/OpenZeppelin/cairo-contracts.git", tag = "v0.8.0" }
+alexandria_merkle_tree = { git = "https://github.com/keep-starknet-strange/alexandria" }
 
 [[target.starknet-contract]]
 sierra = true

contracts/src/tokens/interface.cairo

@@ -43,6 +43,11 @@ trait IUnruggableMemecoin<TState> {
     fn launched(self: @TState) -> bool;
     fn launch_memecoin(ref self: TState);
     fn get_team_allocation(self: @TState) -> u256;
+    fn set_merkle_root(ref self: TState, merkle_root: felt252);
+    fn get_merkle_root(self: @TState) -> felt252;
+    fn claim_airdrop(
+        ref self: TState, to: ContractAddress, amount: u256, leaf: felt252, proof: Span<felt252>,
+    );
 }
 
 #[starknet::interface]
@@ -75,4 +80,9 @@ trait IUnruggableAdditional<TState> {
     fn launched(self: @TState) -> bool;
     fn launch_memecoin(ref self: TState);
     fn get_team_allocation(self: @TState) -> u256;
+    fn set_merkle_root(ref self: TState, merkle_root: felt252);
+    fn get_merkle_root(self: @TState) -> felt252;
+    fn claim_airdrop(
+        ref self: TState, to: ContractAddress, amount: u256, leaf: felt252, proof: Span<felt252>,
+    );
 }

contracts/src/tokens/memecoin.cairo

@@ -13,12 +13,19 @@ mod UnruggableMemecoin {
         IUnruggableMemecoinSnake, IUnruggableMemecoinCamel, IUnruggableAdditional
     };
     use zeroable::Zeroable;
+    use alexandria_merkle_tree::merkle_tree::{
+        Hasher, MerkleTree, pedersen::PedersenHasherImpl, MerkleTreeTrait, MerkleTreeImpl
+    };
+    use openzeppelin::security::initializable::InitializableComponent::InternalTrait as InitializableTrait;
+    use openzeppelin::security::initializable::InitializableComponent;
 
     // Components.
     component!(path: OwnableComponent, storage: ownable, event: OwnableEvent);
     impl OwnableInternalImpl = OwnableComponent::InternalImpl<ContractState>;
 
     component!(path: ERC20Component, storage: erc20, event: ERC20Event);
+
+    component!(path: InitializableComponent, storage: initializable, event: InitializableEvent);
     // Internals
     impl ERC20InternalImpl = ERC20Component::InternalImpl<ContractState>;
 
@@ -46,7 +53,12 @@ mod UnruggableMemecoin {
         #[substorage(v0)]
         ownable: OwnableComponent::Storage,
         #[substorage(v0)]
-        erc20: ERC20Component::Storage
+        erc20: ERC20Component::Storage,
+        #[substorage(v0)]
+        initializable: InitializableComponent::Storage,
+        //Contract Storage
+        merkle_root: felt252,
+        has_claimed: LegacyMap::<ContractAddress, bool>,
     }
 
     #[event]
@@ -55,14 +67,24 @@ mod UnruggableMemecoin {
         #[flat]
         OwnableEvent: OwnableComponent::Event,
         #[flat]
-        ERC20Event: ERC20Component::Event
+        ERC20Event: ERC20Component::Event,
+        #[flat]
+        InitializableEvent: InitializableComponent::Event,
+        //Contract Events
+        ClaimedAirdrop: ClaimedAirdrop,
     }
 
     mod Errors {
         const MAX_HOLDERS_REACHED: felt252 = 'Unruggable: max holders reached';
         const ARRAYS_LEN_DIF: felt252 = 'Unruggable: arrays len dif';
     }
 
+    #[derive(Drop, starknet::Event)]
+    struct ClaimedAirdrop {
+        account: ContractAddress,
+        amount: u256
+    }
+
 
     /// Constructor called once when the contract is deployed.
     /// # Arguments
@@ -137,6 +159,45 @@ mod UnruggableMemecoin {
                 * MAX_SUPPLY_PERCENTAGE_TEAM_ALLOCATION.into()
                 / 100
         }
+        fn set_merkle_root(ref self: ContractState, merkle_root: felt252) {
+            //Initializing the merkle root
+            self.ownable.assert_only_owner();
+            self.initializable.initialize();
+            self.merkle_root.write(merkle_root);
+        }
+
+        fn get_merkle_root(self: @ContractState) -> felt252 {
+            //Getting the merkle root
+            self.ownable.assert_only_owner();
+            self.merkle_root.read()
+        }
+
+        fn claim_airdrop(
+            ref self: ContractState,
+            to: ContractAddress,
+            amount: u256,
+            mut leaf: felt252,
+            mut proof: Span<felt252>,
+        ) {
+            //Initializing the Merkletree
+            let mut merkle_tree: MerkleTree<Hasher> = MerkleTreeTrait::new();
+            //Type casting it for pedersen hashing
+            let to_felt252: felt252 = starknet::contract_address_to_felt252(to);
+            let amount_felt252: felt252 = amount.try_into().unwrap();
+
+            //Verifying the proof
+            let valid_proof: bool = merkle_tree.verify(self.merkle_root.read(), leaf, proof);
+            assert(self.has_claimed.read(to) == false, 'Already Claimed');
+            assert(valid_proof == true, 'Invalid proof');
+
+            //Changing the has_claimed state to true
+            self.has_claimed.write(to, true);
+
+            //Minting the tokens
+            self.erc20._mint(to, amount);
+            //Emitting an event of ClaimedAirdrop
+            self.emit(ClaimedAirdrop { account: to, amount: amount });
+        }
     }
 
     #[abi(embed_v0)]

contracts/tests/test_unruggable_memecoin.cairo

@@ -1221,4 +1221,250 @@ mod memecoin_internals {
             index += 1;
         };
     }
+    #[test]
+    fn test_set_merkle_root() {
+        let (
+            owner,
+            recipient,
+            name,
+            symbol,
+            initial_supply,
+            initial_holder_1,
+            initial_holder_2,
+            _,
+            initial_holders_amounts
+        ) =
+            instantiate_params();
+        let alice = contract_address_const::<53>();
+        let initial_holders = array![owner, initial_holder_1, initial_holder_2].span();
+
+        let contract_address =
+            match deploy_contract(
+                owner, owner, name, symbol, initial_supply, initial_holders, initial_holders_amounts
+            ) {
+            Result::Ok(address) => address,
+            Result::Err(msg) => panic(msg.panic_data),
+        };
+
+        let memecoin = IUnruggableMemecoinDispatcher { contract_address };
+
+        start_prank(CheatTarget::One(memecoin.contract_address), owner);
+        memecoin.set_merkle_root(0x3022e5c16bd1bf6e9c44b0d0de23ef6eb0bc84bd2b4eaca75306076eb99239a);
+
+        assert(
+            memecoin
+                .get_merkle_root() == 0x3022e5c16bd1bf6e9c44b0d0de23ef6eb0bc84bd2b4eaca75306076eb99239a,
+            'Inconsistency'
+        )
+    //TODO
+    }
+
+    #[test]
+    #[should_panic(expected: ('Caller is not the owner',))]
+    fn test_set_merkle_root_not_owner() {
+        let (
+            owner,
+            recipient,
+            name,
+            symbol,
+            initial_supply,
+            initial_holder_1,
+            initial_holder_2,
+            _,
+            initial_holders_amounts
+        ) =
+            instantiate_params();
+        let alice = contract_address_const::<53>();
+        let initial_holders = array![owner, initial_holder_1, initial_holder_2].span();
+
+        let contract_address =
+            match deploy_contract(
+                owner, owner, name, symbol, initial_supply, initial_holders, initial_holders_amounts
+            ) {
+            Result::Ok(address) => address,
+            Result::Err(msg) => panic(msg.panic_data),
+        };
+
+        let memecoin = IUnruggableMemecoinDispatcher { contract_address };
+
+        memecoin.set_merkle_root(0x4c5b879125d0fe0e0359dc87eea9c7370756635ca87c59148fb313c2cfb0579);
+    }
+
+    #[test]
+    #[should_panic(expected: ('Initializable: is initialized',))]
+    fn test_set_merkle_root_initialized() {
+        let (
+            owner,
+            recipient,
+            name,
+            symbol,
+            initial_supply,
+            initial_holder_1,
+            initial_holder_2,
+            _,
+            initial_holders_amounts
+        ) =
+            instantiate_params();
+        let alice = contract_address_const::<53>();
+        let initial_holders = array![owner, initial_holder_1, initial_holder_2].span();
+
+        let contract_address =
+            match deploy_contract(
+                owner, owner, name, symbol, initial_supply, initial_holders, initial_holders_amounts
+            ) {
+            Result::Ok(address) => address,
+            Result::Err(msg) => panic(msg.panic_data),
+        };
+
+        let memecoin = IUnruggableMemecoinDispatcher { contract_address };
+        start_prank(CheatTarget::One(memecoin.contract_address), owner);
+
+        memecoin.set_merkle_root(0x4c5b879125d0fe0e0359dc87eea9c7370756635ca87c59148fb313c2cfb0579);
+        memecoin.set_merkle_root(0x4c5b879125d0fe0e0359dc87eea9c7370756635ca87c59148fb313c2cfb0579);
+    }
+
+    #[test]
+    fn test_claimairdrop_memecoin() { //Testing claimairdrop 
+        let (
+            owner,
+            recipient,
+            name,
+            symbol,
+            initial_supply,
+            initial_holder_1,
+            initial_holder_2,
+            _,
+            initial_holders_amounts
+        ) =
+            instantiate_params();
+        let alice = contract_address_const::<53>();
+        let initial_holders = array![owner, initial_holder_1, initial_holder_2].span();
+
+        let contract_address =
+            match deploy_contract(
+                owner, owner, name, symbol, initial_supply, initial_holders, initial_holders_amounts
+            ) {
+            Result::Ok(address) => address,
+            Result::Err(msg) => panic(msg.panic_data),
+        };
+
+        let memecoin = IUnruggableMemecoinDispatcher { contract_address };
+
+        start_prank(CheatTarget::One(memecoin.contract_address), owner);
+
+        //These values have been generated through the merkle.ts script
+        memecoin.set_merkle_root(0x4c5b879125d0fe0e0359dc87eea9c7370756635ca87c59148fb313c2cfb0579);
+
+        let to = contract_address_const::<
+            0x02038e178565b977c99f3e6c8d4ba327356e1b279e84cdc2f1949022c91653bd
+        >();
+        let amount = 500_u256;
+
+        let leaf = 0x57d0e61d7b5849581495af551721710023a83e705710c58facfa3f4e36e8fac;
+        let valid_proof = array![0x3bf438e95d7428d14eb4270528ff8b1e2f9cb30113724626d5cf9943551ee4d]
+            .span();
+
+        memecoin.claim_airdrop(to, amount, leaf, valid_proof);
+
+        let balance = memecoin.balanceOf(to);
+        assert(balance == amount, 'No balance');
+    }
+
+    #[test]
+    #[should_panic(expected: ('Already Claimed',))]
+    fn test_claimairdrop_memecoin_should_fail() { //Already Claimed
+        let (
+            owner,
+            recipient,
+            name,
+            symbol,
+            initial_supply,
+            initial_holder_1,
+            initial_holder_2,
+            _,
+            initial_holders_amounts
+        ) =
+            instantiate_params();
+        let alice = contract_address_const::<53>();
+        let initial_holders = array![owner, initial_holder_1, initial_holder_2].span();
+
+        let contract_address =
+            match deploy_contract(
+                owner, owner, name, symbol, initial_supply, initial_holders, initial_holders_amounts
+            ) {
+            Result::Ok(address) => address,
+            Result::Err(msg) => panic(msg.panic_data),
+        };
+
+        let memecoin = IUnruggableMemecoinDispatcher { contract_address };
+
+        start_prank(CheatTarget::One(memecoin.contract_address), owner);
+
+        //These values have been generated through the merkle.ts script
+        memecoin.set_merkle_root(0x4c5b879125d0fe0e0359dc87eea9c7370756635ca87c59148fb313c2cfb0579);
+
+        let to = contract_address_const::<
+            0x02038e178565b977c99f3e6c8d4ba327356e1b279e84cdc2f1949022c91653bd
+        >();
+        let amount = 500_u256;
+
+        let leaf = 0x57d0e61d7b5849581495af551721710023a83e705710c58facfa3f4e36e8fac;
+        let valid_proof = array![0x3bf438e95d7428d14eb4270528ff8b1e2f9cb30113724626d5cf9943551ee4d]
+            .span();
+
+        memecoin.claim_airdrop(to, amount, leaf, valid_proof);
+
+        let balance = memecoin.balanceOf(to);
+        assert(balance == amount, 'No balance');
+
+        memecoin.claim_airdrop(to, amount, leaf, valid_proof);
+    }
+
+    #[test]
+    #[should_panic(expected: ('Invalid proof',))]
+    fn test_claimairdrop_memecoin_should_fail_2() { //Invalid Proof
+        let (
+            owner,
+            recipient,
+            name,
+            symbol,
+            initial_supply,
+            initial_holder_1,
+            initial_holder_2,
+            _,
+            initial_holders_amounts
+        ) =
+            instantiate_params();
+        let alice = contract_address_const::<53>();
+        let initial_holders = array![owner, initial_holder_1, initial_holder_2].span();
+
+        let contract_address =
+            match deploy_contract(
+                owner, owner, name, symbol, initial_supply, initial_holders, initial_holders_amounts
+            ) {
+            Result::Ok(address) => address,
+            Result::Err(msg) => panic(msg.panic_data),
+        };
+
+        let memecoin = IUnruggableMemecoinDispatcher { contract_address };
+
+        start_prank(CheatTarget::One(memecoin.contract_address), owner);
+
+        //These values have been generated through the merkle.ts script
+        memecoin.set_merkle_root(0x3022e5c16bd1bf6e9c44b0d0de23ef6eb0bc84bd2b4eaca75306076eb99239a);
+
+        let to = contract_address_const::<
+            0x02038e178565b977c99f3e6c8d4ba327356e1b279e84cdc2f1949022c91653bd
+        >();
+        let amount = 500_u256;
+
+        let leaf = 0x57d0e61d7b5849581495af551721710023a83e705710c58facfa3f4e36e8fac;
+        let valid_proof = array![0x3bf438e95d7428d14eb4270528ff8b1e2f9cb30113724626d5cf9943551ee4d]
+            .span();
+
+        memecoin.claim_airdrop(to, amount, leaf, valid_proof);
+
+        let balance = memecoin.balanceOf(to);
+        assert(balance == amount, 'No balance');
+    }
 }