feat: replace account class (#1428)

<!--- Please provide a general summary of your changes in the title
above -->

<!-- Give an estimate of the time you spent on this PR in terms of work
days.
Did you spend 0.5 days on this PR or rather 2 days?  -->

Time spent on this PR:

## Pull request type

<!-- Please try to limit your pull request to one type,
submit multiple pull requests if needed. -->

Please check the type of change your PR introduces:

- [ ] Bugfix
- [ ] Feature
- [ ] Code style update (formatting, renaming)
- [ ] Refactoring (no functional changes, no api changes)
- [ ] Build related changes
- [ ] Documentation content changes
- [ ] Other (please describe):

## What is the current behavior?

<!-- Please describe the current behavior that you are modifying,
or link to a relevant issue. -->

Resolves #1424

## What is the new behavior?

<!-- Please describe the behavior or changes that are being added by
this PR. -->

- Adds a new entrypoint `replace_class` to the account contract to
replace it's class. This will be used when migrating from CairoZero to
Cairo
- Add a test to verify we can upgrade accounts arbitrarly from Kakarot.
-

<!-- Reviewable:start -->
- - -
This change is [<img src="https://reviewable.io/review_button.svg"
height="34" align="absmiddle"
alt="Reviewable"/>](https://reviewable.io/reviews/kkrt-labs/kakarot/1428)
<!-- Reviewable:end -->

src/kakarot/accounts/account_contract.cairo

@@ -13,6 +13,7 @@ from starkware.starknet.common.syscalls import (
     get_caller_address,
     get_block_timestamp,
     call_contract,
+    replace_class,
 )
 from starkware.cairo.common.bool import FALSE, TRUE
 
@@ -59,6 +60,16 @@ func get_evm_address{syscall_ptr: felt*, pedersen_ptr: HashBuiltin*, range_check
     return AccountContract.get_evm_address();
 }
 
+// @notice Sets the implementation of the account contract.
+// @param implementation_class The class to replace the current implementation with.
+@external
+func upgrade{syscall_ptr: felt*, pedersen_ptr: HashBuiltin*, range_check_ptr}(new_class: felt) {
+    // Access control check.
+    Ownable.assert_only_owner();
+    replace_class(new_class);
+    return ();
+}
+
 // @notice Checks if the account was initialized.
 // @return is_initialized: 1 if the account has been initialized 0 otherwise.
 @view

src/kakarot/interfaces/interfaces.cairo

@@ -39,7 +39,7 @@ namespace IAccount {
     func get_evm_address() -> (evm_address: felt) {
     }
 
-    func set_implementation(implementation: felt) {
+    func upgrade(new_class: felt) {
     }
 
     func bytecode_len() -> (len: felt) {

src/kakarot/library.cairo

@@ -358,7 +358,7 @@ namespace Kakarot {
     ) {
         alloc_locals;
         let starknet_address = Account.get_starknet_address(evm_address);
-        IAccount.set_implementation(starknet_address, new_class_hash);
+        IAccount.upgrade(starknet_address, new_class_hash);
         return ();
     }
 

tests/src/kakarot/accounts/test_account_contract.cairo

@@ -18,6 +18,7 @@ from kakarot.accounts.account_contract import (
     execute_starknet_call,
     set_code_hash,
     execute_from_outside,
+    upgrade,
 )
 from kakarot.accounts.model import OutsideExecution, CallArray
 
@@ -222,3 +223,11 @@ func test__set_code_hash{syscall_ptr: felt*, pedersen_ptr: HashBuiltin*, range_c
 
     return ();
 }
+
+func test__upgrade{syscall_ptr: felt*, pedersen_ptr: HashBuiltin*, range_check_ptr}() {
+    alloc_locals;
+    local new_class: felt;
+    %{ ids.new_class = program_input["new_class"] %}
+    upgrade(new_class);
+    return ();
+}

tests/src/kakarot/accounts/test_account_contract.py

@@ -152,6 +152,17 @@ def test_should_set_nonce(self, cairo_run):
                 value=1,
             )
 
+    class TestUpgrade:
+        @SyscallHandler.patch("Ownable_owner", 0xDEAD)
+        def test_should_assert_only_owner(self, cairo_run):
+            with cairo_error(message="Ownable: caller is not the owner"):
+                cairo_run("test__upgrade", new_class=0x00)
+
+        @SyscallHandler.patch("Ownable_owner", SyscallHandler.caller_address)
+        def test_should_set_new_class(self, cairo_run):
+            cairo_run("test__upgrade", new_class=0x1234)
+            SyscallHandler.mock_replace_class.assert_any_call(class_hash=0x1234)
+
     class TestJumpdests:
         class TestWriteJumpdests:
             @SyscallHandler.patch("Ownable_owner", 0xDEAD)

tests/src/kakarot/test_kakarot.cairo

@@ -18,6 +18,7 @@ from kakarot.kakarot import (
     set_authorized_cairo_precompile_caller,
     set_cairo1_helpers_class_hash,
     transfer_ownership,
+    upgrade_account,
 )
 from kakarot.model import model
 from kakarot.account import Account
@@ -215,3 +216,17 @@ func test__eth_chain_id{syscall_ptr: felt*, pedersen_ptr: HashBuiltin*, range_ch
     let (chain_id) = Kakarot.eth_chain_id();
     return chain_id;
 }
+
+func test__upgrade_account{syscall_ptr: felt*, pedersen_ptr: HashBuiltin*, range_check_ptr}() {
+    tempvar evm_address;
+    tempvar new_class_hash;
+
+    %{
+        ids.evm_address = program_input["evm_address"]
+        ids.new_class_hash = program_input["new_class_hash"]
+    %}
+
+    upgrade_account(evm_address, new_class_hash);
+
+    return ();
+}

tests/src/kakarot/test_kakarot.py

@@ -8,7 +8,10 @@
 from eth_utils.address import to_checksum_address
 from hypothesis import given
 from hypothesis.strategies import composite, integers
-from starkware.starknet.public.abi import get_storage_var_address
+from starkware.starknet.public.abi import (
+    get_selector_from_name,
+    get_storage_var_address,
+)
 from web3._utils.abi import map_abi_data
 from web3._utils.normalizers import BASE_RETURN_NORMALIZERS
 from web3.exceptions import NoABIFunctionsFound
@@ -292,6 +295,33 @@ def test_register_account_should_fail_caller_not_resolved_address(
             ):
                 cairo_run("test__register_account", evm_address=EVM_ADDRESS)
 
+        class TestUpgradeAccount:
+            @SyscallHandler.patch("Ownable_owner", 0xDEAD)
+            def test_should_assert_only_owner(self, cairo_run):
+                with cairo_error(message="Ownable: caller is not the owner"):
+                    cairo_run(
+                        "test__upgrade_account",
+                        evm_address=EVM_ADDRESS,
+                        new_class_hash=0x1234,
+                    )
+
+            @SyscallHandler.patch(
+                "Kakarot_evm_to_starknet_address", EVM_ADDRESS, 0x99999
+            )
+            @SyscallHandler.patch("Ownable_owner", SyscallHandler.caller_address)
+            @SyscallHandler.patch("IAccount.upgrade", lambda addr, data: [])
+            def test_upgrade_account_should_replace_class(self, cairo_run):
+                cairo_run(
+                    "test__upgrade_account",
+                    evm_address=EVM_ADDRESS,
+                    new_class_hash=0x1234,
+                )
+                SyscallHandler.mock_call.assert_called_with(
+                    contract_address=0x99999,
+                    function_selector=get_selector_from_name("upgrade"),
+                    calldata=[0x1234],
+                )
+
     class TestEthCall:
         @pytest.mark.slow
         @pytest.mark.SolmateERC20