Rewrite Registry & Account as Components and add upgradable to Registry

src/account.cairo

@@ -1 +1,3 @@
 mod account;
+
+use account::AccountComponent;

src/account/account.cairo

@@ -1,8 +1,8 @@
 ////////////////////////////////
-// Account contract
+// Account Component
 ////////////////////////////////
-#[starknet::contract]
-mod Account {
+#[starknet::component]
+mod AccountComponent {
     use starknet::{
         get_tx_info, get_caller_address, get_contract_address, get_block_timestamp, ContractAddress,
         account::Call, call_contract_syscall, replace_class_syscall, ClassHash, SyscallResultTrait
@@ -14,23 +14,19 @@ mod Account {
     use zeroable::Zeroable;
     use token_bound_accounts::interfaces::IERC721::{IERC721DispatcherTrait, IERC721Dispatcher};
     use token_bound_accounts::interfaces::IAccount::IAccount;
-
-    // SRC5 interface for token bound accounts
-    const TBA_INTERFACE_ID: felt252 =
-        0x539036932a2ab9c4734fbfd9872a1f7791a3f577e45477336ae0fd0a00c9ff;
+    use token_bound_accounts::interfaces::IAccount::{TBA_INTERFACE_ID};
 
     #[storage]
     struct Storage {
-        _token_contract: ContractAddress, // contract address of NFT
-        _token_id: u256, // token ID of NFT
-        _unlock_timestamp: u64, // time to unlock account when locked
+        Account_token_contract: ContractAddress, // contract address of NFT
+        Account_token_id: u256, // token ID of NFT
+        Account_unlock_timestamp: u64, // time to unlock account when locked
     }
 
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
         AccountCreated: AccountCreated,
-        AccountUpgraded: AccountUpgraded,
         AccountLocked: AccountLocked,
         TransactionExecuted: TransactionExecuted
     }
@@ -53,15 +49,6 @@ mod Account {
         response: Span<Span<felt252>>
     }
 
-    /// @notice Emitted when the account upgrades to a new implementation
-    /// @param account tokenbound account to be upgraded
-    /// @param implementation the upgraded account class hash
-    #[derive(Drop, starknet::Event)]
-    struct AccountUpgraded {
-        account: ContractAddress,
-        implementation: ClassHash
-    }
-
     /// @notice Emitted when the account is locked
     /// @param account tokenbound account who's lock function was triggered
     /// @param locked_at timestamp at which the lock function was triggered
@@ -74,51 +61,54 @@ mod Account {
         duration: u64,
     }
 
-    #[constructor]
-    fn constructor(ref self: ContractState, token_contract: ContractAddress, token_id: u256) {
-        self._token_contract.write(token_contract);
-        self._token_id.write(token_id);
-
-        let owner = self._get_owner(token_contract, token_id);
-        self.emit(AccountCreated { owner });
+    mod Errors {
+        const LOCKED_ACCOUNT: felt252 = 'Account: account is locked!';
+        const INV_TX_VERSION:  felt252 = 'Account: invalid tx version';
+        const UNAUTHORIZED: felt252 = 'Account: unauthorized';
+        const INV_SIG_LEN: felt252 = 'Account: invalid sig length';
+        const INV_SIGNATURE: felt252 = 'Account: invalid signature';
     }
 
-    #[external(v0)]
-    impl IAccountImpl of IAccount<ContractState> {
+    #[embeddable_as(AccountImpl)]
+    impl Account<
+        TContractState,
+        +HasComponent<TContractState>,
+        +Drop<TContractState>
+        > of IAccount<ComponentState<TContractState>> {
         /// @notice used for signature validation
         /// @param hash The message hash 
         /// @param signature The signature to be validated
         fn is_valid_signature(
-            self: @ContractState, hash: felt252, signature: Span<felt252>
+            self: @ComponentState<TContractState>, hash: felt252, signature: Span<felt252>
         ) -> felt252 {
             self._is_valid_signature(hash, signature)
         }
 
         fn __validate_deploy__(
-            self: @ContractState, class_hash: felt252, contract_address_salt: felt252,
+            self: @ComponentState<TContractState>, class_hash: felt252, contract_address_salt: felt252,
         ) -> felt252 {
             self._validate_transaction()
         }
 
-        fn __validate_declare__(self: @ContractState, class_hash: felt252) -> felt252 {
+        fn __validate_declare__(self: @ComponentState<TContractState>, class_hash: felt252) -> felt252 {
             self._validate_transaction()
         }
 
         /// @notice validate an account transaction
         /// @param calls an array of transactions to be executed
-        fn __validate__(ref self: ContractState, mut calls: Array<Call>) -> felt252 {
+        fn __validate__(ref self: ComponentState<TContractState>, mut calls: Array<Call>) -> felt252 {
             self._validate_transaction()
         }
 
         /// @notice executes a transaction
         /// @param calls an array of transactions to be executed
-        fn __execute__(ref self: ContractState, mut calls: Array<Call>) -> Array<Span<felt252>> {
+        fn __execute__(ref self: ComponentState<TContractState>, mut calls: Array<Call>) -> Array<Span<felt252>> {
             self._assert_only_owner();
             let (lock_status, _) = self._is_locked();
-            assert(!lock_status, 'Account: account is locked!');
+            assert(!lock_status, Errors::LOCKED_ACCOUNT);
 
             let tx_info = get_tx_info().unbox();
-            assert(tx_info.version != 0, 'invalid tx version');
+            assert(tx_info.version != 0, Errors::INV_TX_VERSION);
 
             let retdata = self._execute_calls(calls.span());
             let hash = tx_info.transaction_hash;
@@ -131,36 +121,25 @@ mod Account {
         /// @param token_contract the contract address of the NFT
         /// @param token_id the token ID of the NFT
         fn owner(
-            self: @ContractState, token_contract: ContractAddress, token_id: u256
+            self: @ComponentState<TContractState>, token_contract: ContractAddress, token_id: u256
         ) -> ContractAddress {
             self._get_owner(token_contract, token_id)
         }
 
         /// @notice returns the contract address and token ID of the NFT
-        fn token(self: @ContractState) -> (ContractAddress, u256) {
+        fn token(self: @ComponentState<TContractState>) -> (ContractAddress, u256) {
             self._get_token()
         }
 
-        /// @notice ugprades an account implementation
-        /// @param implementation the new class_hash
-        fn upgrade(ref self: ContractState, implementation: ClassHash) {
-            self._assert_only_owner();
-            let (lock_status, _) = self._is_locked();
-            assert(!lock_status, 'Account: account is locked!');
-            assert(!implementation.is_zero(), 'Invalid class hash');
-            replace_class_syscall(implementation).unwrap_syscall();
-            self.emit(AccountUpgraded { account: get_contract_address(), implementation, });
-        }
-
         // @notice protection mechanism for selling token bound accounts. can't execute when account is locked
         // @param duration for which to lock account
-        fn lock(ref self: ContractState, duration: u64) {
+        fn lock(ref self: ComponentState<TContractState>, duration: u64) {
             self._assert_only_owner();
             let (lock_status, _) = self._is_locked();
-            assert(!lock_status, 'Account: account already locked');
+            assert(!lock_status, Errors::LOCKED_ACCOUNT);
             let current_timestamp = get_block_timestamp();
             let unlock_time = current_timestamp + duration;
-            self._unlock_timestamp.write(unlock_time);
+            self.Account_unlock_timestamp.write(unlock_time);
             self
                 .emit(
                     AccountLocked {
@@ -170,13 +149,13 @@ mod Account {
         }
 
         // @notice returns account lock status and time left until account unlocks
-        fn is_locked(self: @ContractState) -> (bool, u64) {
+        fn is_locked(self: @ComponentState<TContractState>) -> (bool, u64) {
             return self._is_locked();
         }
 
         // @notice check that account supports TBA interface
         // @param interface_id interface to be checked against
-        fn supports_interface(self: @ContractState, interface_id: felt252) -> bool {
+        fn supports_interface(self: @ComponentState<TContractState>, interface_id: felt252) -> bool {
             if (interface_id == TBA_INTERFACE_ID) {
                 return true;
             } else {
@@ -186,20 +165,35 @@ mod Account {
     }
 
     #[generate_trait]
-    impl internalImpl of InternalTrait {
+    impl InternalImpl<
+        TContractState,
+        +HasComponent<TContractState>,
+        +Drop<TContractState>
+    > of InternalTrait<TContractState> {
+
+        /// @notice initializes the account by setting the initial token conrtact and token id
+        fn initializer(ref self: ComponentState<TContractState>, token_contract: ContractAddress, token_id: u256) {
+
+        self.Account_token_contract.write(token_contract);
+        self.Account_token_id.write(token_id);
+
+        let owner = self._get_owner(token_contract, token_id);
+        self.emit(AccountCreated { owner });
+    }
+
         /// @notice check that caller is the token bound account
-        fn _assert_only_owner(ref self: ContractState) {
+        fn _assert_only_owner(ref self: ComponentState<TContractState>) {
             let caller = get_caller_address();
-            let owner = self._get_owner(self._token_contract.read(), self._token_id.read());
-            assert(caller == owner, 'Account: unathorized');
+            let owner = self._get_owner(self.Account_token_contract.read(), self.Account_token_id.read());
+            assert(caller == owner, Errors::UNAUTHORIZED);
         }
 
         /// @notice internal function for getting NFT owner
         /// @param token_contract contract address of NFT
         // @param token_id token ID of NFT
         // NB: This function aims for compatibility with all contracts (snake or camel case) but do not work as expected on mainnet as low level calls do not return err at the moment. Should work for contracts which implements CamelCase but not snake_case until starknet v0.15.
         fn _get_owner(
-            self: @ContractState, token_contract: ContractAddress, token_id: u256
+            self: @ComponentState<TContractState>, token_contract: ContractAddress, token_id: u256
         ) -> ContractAddress {
             let mut calldata: Array<felt252> = ArrayTrait::new();
             Serde::serialize(@token_id, ref calldata);
@@ -214,15 +208,15 @@ mod Account {
         }
 
         /// @notice internal transaction for returning the contract address and token ID of the NFT
-        fn _get_token(self: @ContractState) -> (ContractAddress, u256) {
-            let contract = self._token_contract.read();
-            let tokenId = self._token_id.read();
+        fn _get_token(self: @ComponentState<TContractState>) -> (ContractAddress, u256) {
+            let contract = self.Account_token_contract.read();
+            let tokenId = self.Account_token_id.read();
             (contract, tokenId)
         }
 
         // @notice protection mechanism for TBA trading. Returns the lock-status (true or false), and the remaning time till account unlocks.
-        fn _is_locked(self: @ContractState) -> (bool, u64) {
-            let unlock_timestamp = self._unlock_timestamp.read();
+        fn _is_locked(self: @ComponentState<TContractState>) -> (bool, u64) {
+            let unlock_timestamp = self.Account_unlock_timestamp.read();
             let current_time = get_block_timestamp();
             if (current_time < unlock_timestamp) {
                 let time_until_unlocks = unlock_timestamp - current_time;
@@ -233,26 +227,26 @@ mod Account {
         }
 
         /// @notice internal function for tx validation
-        fn _validate_transaction(self: @ContractState) -> felt252 {
+        fn _validate_transaction(self: @ComponentState<TContractState>) -> felt252 {
             let tx_info = get_tx_info().unbox();
             let tx_hash = tx_info.transaction_hash;
             let signature = tx_info.signature;
             assert(
                 self._is_valid_signature(tx_hash, signature) == starknet::VALIDATED,
-                'Account: invalid signature'
+                Errors::INV_SIGNATURE
             );
             starknet::VALIDATED
         }
 
         /// @notice internal function for signature validation
         fn _is_valid_signature(
-            self: @ContractState, hash: felt252, signature: Span<felt252>
+            self: @ComponentState<TContractState>, hash: felt252, signature: Span<felt252>
         ) -> felt252 {
             let signature_length = signature.len();
-            assert(signature_length == 2_u32, 'Account: invalid sig length');
+            assert(signature_length == 2_u32, Errors::INV_SIG_LEN);
 
             let caller = get_caller_address();
-            let owner = self._get_owner(self._token_contract.read(), self._token_id.read());
+            let owner = self._get_owner(self.Account_token_contract.read(), self.Account_token_id.read());
             if (caller == owner) {
                 return starknet::VALIDATED;
             } else {
@@ -262,7 +256,7 @@ mod Account {
 
         /// @notice internal function for executing transactions
         /// @param calls An array of transactions to be executed
-        fn _execute_calls(ref self: ContractState, mut calls: Span<Call>) -> Array<Span<felt252>> {
+        fn _execute_calls(ref self: ComponentState<TContractState>, mut calls: Span<Call>) -> Array<Span<felt252>> {
             let mut result: Array<Span<felt252>> = ArrayTrait::new();
             let mut calls = calls;
 

src/interfaces.cairo

@@ -1,3 +1,4 @@
 mod IAccount;
 mod IERC721;
 mod IRegistry;
+mod IUpgradeable;

src/interfaces/IAccount.cairo

@@ -2,6 +2,9 @@ use starknet::ContractAddress;
 use starknet::ClassHash;
 use starknet::account::Call;
 
+// SRC5 interface for token bound accounts
+const TBA_INTERFACE_ID: felt252 = 0x539036932a2ab9c4734fbfd9872a1f7791a3f577e45477336ae0fd0a00c9ff;
+
 #[starknet::interface]
 trait IAccount<TContractState> {
     fn is_valid_signature(
@@ -17,7 +20,6 @@ trait IAccount<TContractState> {
     fn owner(
         self: @TContractState, token_contract: ContractAddress, token_id: u256
     ) -> ContractAddress;
-    fn upgrade(ref self: TContractState, implementation: ClassHash);
     fn lock(ref self: TContractState, duration: u64);
     fn is_locked(self: @TContractState) -> (bool, u64);
     fn supports_interface(self: @TContractState, interface_id: felt252) -> bool;

src/interfaces/IUpgradeable.cairo

@@ -0,0 +1,6 @@
+use starknet::ClassHash;
+
+#[starknet::interface]
+trait IUpgradeable<TContractState> {
+    fn upgrade(ref self: TContractState, new_class_hash: ClassHash);
+}

src/lib.cairo

@@ -1,4 +1,6 @@
 mod registry;
 mod account;
 mod interfaces;
+mod presets;
 mod test_helper;
+mod upgradeable;

src/presets.cairo

@@ -0,0 +1,5 @@
+mod account;
+mod registry;
+
+use account::Account;
+use registry::Registry;

src/presets/account.cairo

@@ -0,0 +1,54 @@
+////////////////////////////////
+// Account contract
+////////////////////////////////
+#[starknet::contract]
+mod Account {
+    use starknet::ContractAddress;
+    use starknet::ClassHash;
+    use token_bound_accounts::account::AccountComponent;
+    use token_bound_accounts::upgradeable::UpgradeableComponent;
+    use token_bound_accounts::interfaces::IUpgradeable::IUpgradeable;
+
+    component!(path: AccountComponent, storage: account, event: AccountEvent);
+    component!(path: UpgradeableComponent, storage: upgradeable, event: UpgradeableEvent);
+
+    // Account
+    #[abi(embed_v0)]
+    impl AccountImpl = AccountComponent::AccountImpl<ContractState>;
+    impl AccountInternalImpl = AccountComponent::InternalImpl<ContractState>;
+
+    // Upgradeable
+    impl UpgradeableInternalImpl = UpgradeableComponent::InternalImpl<ContractState>;
+
+    #[storage]
+    struct Storage {
+        #[substorage(v0)]
+        account: AccountComponent::Storage,
+        #[substorage(v0)]
+        upgradeable: UpgradeableComponent::Storage
+    }
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    enum Event {
+        #[flat]
+        AccountEvent: AccountComponent::Event,
+        #[flat]
+        UpgradeableEvent: UpgradeableComponent::Event
+    }
+
+    #[constructor]
+    fn constructor(ref self: ContractState, token_contract: ContractAddress, token_id: u256) {
+        self.account.initializer(token_contract, token_id);
+    }
+
+    #[external(v0)]
+    impl UpgradeableImpl of IUpgradeable<ContractState> {
+        fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
+            self.account._assert_only_owner();
+            let (lock_status, _) = self.account._is_locked();
+            assert(!lock_status, AccountComponent::Errors::LOCKED_ACCOUNT);
+            self.upgradeable._upgrade(new_class_hash);
+        }
+    }   
+}