add referral contracts

.gitignore

@@ -0,0 +1,4 @@
+.DS_Store
+
+.snfoundry_cache/
+target/

Scarb.lock

@@ -0,0 +1,20 @@
+# Code generated by scarb DO NOT EDIT.
+version = 1
+
+[[package]]
+name = "gmx_referral_cairo"
+version = "0.1.0"
+dependencies = [
+ "openzeppelin",
+ "snforge_std",
+]
+
+[[package]]
+name = "openzeppelin"
+version = "0.11.0"
+source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.11.0#a83f36b23f1af6e160288962be4a2701c3ecbcda"
+
+[[package]]
+name = "snforge_std"
+version = "0.20.1"
+source = "git+https://github.com/foundry-rs/starknet-foundry.git?tag=v0.20.1#fea2db8f2b20148cc15ee34b08de12028eb42942"

Scarb.toml

@@ -0,0 +1,18 @@
+[package]
+name = "gmx_referral_cairo"
+version = "0.1.0"
+
+[lib]
+
+
+[[target.starknet-contract]]
+sierra = true
+casm = true
+casm-add-pythonic-hints = true
+allowed-libfuncs-list.name = "all"
+build-external-contracts = ["openzeppelin::presets::erc20::ERC20"]
+
+[dependencies]
+starknet = "2.4.3"
+openzeppelin = { git = "https://github.com/OpenZeppelin/cairo-contracts.git", tag = "v0.11.0" }
+snforge_std = { git = "https://github.com/foundry-rs/starknet-foundry.git", tag = "v0.20.1" }

snfoundry.toml

@@ -0,0 +1,4 @@
+[sncast.default]
+account = "my_imported_account"
+accounts-file = "~/my_accounts.json"
+url = "https://api-starknet-sepolia.dwellir.com/dd28e566-3260-4d8d-8180-6ef1a161e41c"

src/lib.cairo

@@ -0,0 +1,7 @@
+mod referral_storage;
+
+mod test_contracts {
+    mod referral_storage_v2;
+}
+
+

src/referral_storage.cairo

@@ -0,0 +1,110 @@
+use starknet::{ContractAddress,ClassHash};
+use zeroable::Zeroable;
+
+
+#[starknet::interface]
+trait IReferralStorage<TContractState> {
+    fn get_trader_referral_code(
+         self: @TContractState,
+        _account: ContractAddress,
+    ) -> ContractAddress;
+
+    fn set_trader_referral_code(
+        ref self: TContractState,
+        _code: ContractAddress,
+    );
+
+    fn upgrade(ref self: TContractState, new_class_hash: ClassHash);
+}
+
+
+#[starknet::contract]
+mod ReferralStorage {
+    use starknet::ContractAddress;
+    use starknet::get_caller_address;
+    use zeroable::Zeroable;
+    use starknet::ClassHash;
+    use openzeppelin::access::ownable::OwnableComponent;
+    use openzeppelin::upgrades::UpgradeableComponent;
+    use openzeppelin::upgrades::interface::IUpgradeable;
+
+
+    component!(path: UpgradeableComponent, storage: upgradeable_storage, event: UpgradeableEvent);
+    component!(path: OwnableComponent, storage: ownable, event: OwnableEvent);
+
+    #[abi(embed_v0)]
+    /// Ownable
+    impl OwnableImpl = OwnableComponent::OwnableImpl<ContractState>;
+    impl OwnableInternalImpl = OwnableComponent::InternalImpl<ContractState>;
+
+    /// Upgradeable
+    impl UpgradeableInternalImpl = UpgradeableComponent::InternalImpl<ContractState>;
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    enum Event {
+        SetTraderReferralCode: SetTraderReferralCode,
+        #[flat]
+        OwnableEvent: OwnableComponent::Event,
+        #[flat]
+        UpgradeableEvent: UpgradeableComponent::Event
+    }
+
+
+    #[derive(Drop, starknet::Event)]
+    struct SetTraderReferralCode {
+        account: ContractAddress,
+        code: ContractAddress,
+    }
+
+    #[storage]
+    struct Storage {
+        // @notice Maps the trader to the referee code
+        trader_referral_codes: LegacyMap::<ContractAddress, ContractAddress>,
+
+        #[substorage(v0)]
+        ownable: OwnableComponent::Storage,
+        #[substorage(v0)]
+        upgradeable_storage: UpgradeableComponent::Storage,
+    }
+
+    #[constructor]
+    fn constructor(ref self: ContractState, owner: ContractAddress) {
+        self.ownable.initializer(owner);
+    }
+
+    #[abi(embed_v0)]
+    impl ReferralStorage of super::IReferralStorage<ContractState> {
+
+        fn get_trader_referral_code(
+             self: @ContractState,
+            _account: ContractAddress,
+        ) -> ContractAddress {
+             self.trader_referral_codes.read(_account)
+        }
+
+        // @notice Sets the code to be referred by the trader
+        // @param _code The code to set
+        // @dev Trader can set the code only once
+        // @dev Code owner cannot set the code for himself
+        fn set_trader_referral_code(
+            ref self: ContractState,
+            _code: ContractAddress,
+        ){
+            assert!(_code != get_caller_address(), "ReferralStorage: code owner cannot set code for himself");
+
+            if(!self.trader_referral_codes.read(get_caller_address()).is_non_zero()){
+                let _account = get_caller_address();
+                self.trader_referral_codes.write(_account, _code);
+                self.emit(SetTraderReferralCode{account:_account, code: _code});              
+            }
+        }
+
+        fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
+            self.ownable.assert_only_owner();
+
+            self.upgradeable_storage._upgrade(new_class_hash);
+        }
+    }
+}
+

src/test_contracts/referral_storage_v2.cairo

@@ -0,0 +1,159 @@
+use starknet::{ContractAddress,ClassHash};
+use zeroable::Zeroable;
+
+
+#[starknet::interface]
+trait IReferralStorageV2<TContractState> {
+    fn get_trader_referral_code(
+        ref self: TContractState,
+        _account: ContractAddress,
+    ) -> felt252;
+
+    fn get_code_owner(
+        ref self: TContractState,
+        _code: felt252,
+    ) -> ContractAddress;
+
+    fn get_code_from_owner(
+        ref self: TContractState,
+        _account: ContractAddress,
+    ) -> felt252;
+
+    fn set_trader_referral_code(
+        ref self: TContractState,
+        _code: felt252,
+    );
+
+    fn register_code(
+        ref self: TContractState,
+        _code: felt252,
+    );
+
+    fn upgrade(ref self: TContractState, new_class_hash: ClassHash);
+
+}
+
+
+#[starknet::contract]
+mod ReferralStorageV2 {
+    use starknet::ContractAddress;
+    use starknet::get_caller_address;
+    use zeroable::Zeroable;
+    use starknet::ClassHash;
+    use openzeppelin::access::ownable::OwnableComponent;
+    use openzeppelin::upgrades::UpgradeableComponent;
+    use openzeppelin::upgrades::interface::IUpgradeable;
+
+
+    component!(path: UpgradeableComponent, storage: upgradeable_storage, event: UpgradeableEvent);
+    component!(path: OwnableComponent, storage: ownable, event: OwnableEvent);
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    enum Event {
+        SetTraderReferralCode: SetTraderReferralCode,
+        RegisterCode: RegisterCode,
+        #[flat]
+        OwnableEvent: OwnableComponent::Event,
+        #[flat]
+        UpgradeableEvent: UpgradeableComponent::Event
+    }
+
+
+    #[derive(Drop, starknet::Event)]
+    struct SetTraderReferralCode {
+        account: ContractAddress,
+        code: felt252,
+    }
+
+    #[derive(Drop, starknet::Event)]
+    struct RegisterCode {
+        code: felt252,
+        account: ContractAddress,
+    }
+
+    #[storage]
+    struct Storage {
+        code_owner: LegacyMap::<felt252, ContractAddress>,
+        owner_to_code: LegacyMap::<ContractAddress, felt252>,
+        trader_referral_codes: LegacyMap::<ContractAddress, felt252>,
+        #[substorage(v0)]
+        ownable: OwnableComponent::Storage,
+        #[substorage(v0)]
+        upgradeable_storage: UpgradeableComponent::Storage,
+    }
+
+    #[constructor]
+    fn constructor(ref self: ContractState, owner: ContractAddress) {
+        self.ownable.initializer(owner);
+    }
+
+    /// Ownable
+    #[abi(embed_v0)]
+    impl OwnableImpl = OwnableComponent::OwnableImpl<ContractState>;
+    impl OwnableInternalImpl = OwnableComponent::InternalImpl<ContractState>;
+
+    /// Upgradeable
+    impl UpgradeableInternalImpl = UpgradeableComponent::InternalImpl<ContractState>;
+
+    #[abi(embed_v0)]
+    impl ReferralStorageV2 of super::IReferralStorageV2<ContractState> {
+
+        fn get_trader_referral_code(
+            ref self: ContractState,
+            _account: ContractAddress,
+        ) -> felt252 {
+             self.trader_referral_codes.read(_account)
+        }
+
+        fn get_code_owner(
+            ref self: ContractState,
+            _code: felt252,
+        ) -> ContractAddress {
+            self.code_owner.read(_code)
+        }
+
+        fn get_code_from_owner(
+            ref self: ContractState,
+            _account: ContractAddress,
+        ) -> felt252 {
+            self.owner_to_code.read(_account)
+        }
+
+        fn set_trader_referral_code(
+            ref self: ContractState,
+            _code: felt252,
+        ){
+            assert!(self.code_owner.read(_code).is_non_zero(), "ReferralStorage: code not found");
+            assert!(self.code_owner.read(_code) != get_caller_address(), "ReferralStorage: code owner cannot set code for himself");
+
+            let _account = get_caller_address();
+            self.trader_referral_codes.write(_account, _code);
+            self.emit(SetTraderReferralCode{account:_account, code: _code});  
+        }
+
+        fn register_code(
+            ref self: ContractState,
+            _code: felt252,
+        ){
+            assert!(self.owner_to_code.read(get_caller_address()).is_non_zero(), "ReferralStorage: user already has a code");
+
+            assert!(!self.code_owner.read(_code).is_non_zero(), "ReferralStorage: code already registered");
+
+            self.code_owner.write(_code, get_caller_address());
+            self.owner_to_code.write(get_caller_address(), _code);
+
+            self.emit(RegisterCode{code:_code, account: get_caller_address()});
+        }
+
+        fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
+            self.ownable.assert_only_owner();
+
+            self.upgradeable_storage._upgrade(new_class_hash);
+        }
+    }
+
+
+
+}
+

tests/lib.cairo

@@ -0,0 +1 @@
+mod referral_storage_test;

tests/referral_storage_test.cairo

@@ -0,0 +1,62 @@
+use starknet::{ContractAddress,get_caller_address,contract_address_try_from_felt252};
+use zeroable::Zeroable;
+use gmx_referral_cairo::referral_storage::{ReferralStorage,IReferralStorageDispatcher,IReferralStorageDispatcherTrait};
+use snforge_std::{
+    declare, ContractClassTrait, start_prank, stop_prank, CheatTarget, spy_events,
+    SpyOn, EventSpy, EventFetcher, Event, EventAssertions
+};
+use gmx_referral_cairo::test_contracts::referral_storage_v2::{ReferralStorageV2,IReferralStorageV2Dispatcher,IReferralStorageV2DispatcherTrait};
+
+fn setup_referral_storage_dispatcher() -> (ContractAddress,IReferralStorageDispatcher,ContractAddress) {
+    let contract = declare("ReferralStorage");
+
+    let mut contract_constructor_calldata = Default::default();
+    let owner = contract_address_try_from_felt252('owner').unwrap();
+
+    Serde::serialize(@owner, ref contract_constructor_calldata);
+    let contract_address = contract.deploy(@contract_constructor_calldata).unwrap();
+
+    (contract_address,IReferralStorageDispatcher { contract_address },owner)
+}
+
+#[test]
+fn test_set_trader_referral_code() {
+    let (contract_address,dispatcher,_owner) = setup_referral_storage_dispatcher();
+
+    let user = 123.try_into().unwrap();
+
+    let user2 = 124.try_into().unwrap();
+    start_prank(CheatTarget::One(contract_address), user2);
+    dispatcher.set_trader_referral_code(user);
+
+    let user2_referral_code = dispatcher.get_trader_referral_code(user2);
+    assert_eq!(user2_referral_code, user);
+}
+
+#[test]
+#[should_panic(expected: ("ReferralStorage: code owner cannot set code for himself",))]
+fn test_set_trader_referral_code_self() {
+    let (contract_address,dispatcher,_owner) = setup_referral_storage_dispatcher();
+
+    let user = 123.try_into().unwrap();
+    start_prank(CheatTarget::One(contract_address), user);
+    dispatcher.set_trader_referral_code(user);
+}
+
+#[test]
+#[should_panic]
+fn test_upgrade_referral_storage_by_not_onwer() {
+    let (_contract_address,dispatcher,_owner) = setup_referral_storage_dispatcher();
+
+    let new_pool_class_hash = declare("ReferralStorageV2").class_hash;
+    dispatcher.upgrade(new_pool_class_hash);
+}
+
+#[test]
+fn test_upgrade_referral_storage() {
+    let (contract_address,dispatcher,owner) = setup_referral_storage_dispatcher();
+
+    let new_pool_class_hash = declare("ReferralStorageV2").class_hash;
+    start_prank(CheatTarget::One(contract_address), owner);
+    dispatcher.upgrade(new_pool_class_hash);
+}