get_class_hash cheatcode (#441)

<!-- Reference any GitHub issues resolved by this PR -->

Closes #335 

## Introduced changes

Add a cheatcode to get the class hash associated to a given address:

```rust
let class_hash = declare(...);
// Prepare
let contract_address = deploy(...);
assert(get_class_hash(contract_address) == class_hash, 'Wrong class hash');
```

## Breaking changes

<!-- List of all breaking changes, if applicable -->

## Checklist

<!-- Make sure all of these are complete -->

- [X] Linked relevant issue
- [X] Updated relevant documentation
- [X] Added relevant tests
- [X] Performed self-review of the code
- [X] Added changes to `CHANGELOG.md`

---------

Co-authored-by: Piotr Magiera <[email protected]>

Author: glihm

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Documentation to public methods
 - Information sections to documentation about importing `snforge_std`
 - Added print support for basic numeric data types
+- `get_class_hash` cheatcode
 
 #### Changed
 

crates/cheatnet/src/cheatcodes.rs

@@ -10,6 +10,7 @@ use thiserror::Error;
 
 pub mod declare;
 pub mod deploy;
+pub mod get_class_hash;
 pub mod prank;
 pub mod roll;
 pub mod warp;

crates/cheatnet/src/cheatcodes/get_class_hash.rs

@@ -0,0 +1,19 @@
+use crate::{
+    cheatcodes::{CheatcodeError, EnhancedHintError},
+    CheatnetState,
+};
+use blockifier::state::state_api::StateReader;
+use starknet_api::core::{ClassHash, ContractAddress};
+
+impl CheatnetState {
+    /// Gets the class hash at the given address.
+    pub fn get_class_hash(
+        &mut self,
+        contract_address: ContractAddress,
+    ) -> Result<ClassHash, CheatcodeError> {
+        match self.blockifier_state.get_class_hash_at(contract_address) {
+            Ok(class_hash) => Ok(class_hash),
+            Err(e) => Err(CheatcodeError::Unrecoverable(EnhancedHintError::State(e))),
+        }
+    }
+}

crates/forge/src/cheatcodes_hint_processor.rs

@@ -4,7 +4,8 @@ use std::path::PathBuf;
 
 use crate::scarb::StarknetContractArtifacts;
 use anyhow::{anyhow, Result};
-use blockifier::execution::execution_utils::stark_felt_to_felt;
+use blockifier::execution::execution_utils::{felt_to_stark_felt, stark_felt_to_felt};
+
 use cairo_felt::Felt252;
 use cairo_vm::hint_processor::hint_processor_definition::HintProcessorLogic;
 use cairo_vm::hint_processor::hint_processor_definition::HintReference;
@@ -18,7 +19,6 @@ use cheatnet::{
     cheatcodes::{CheatcodeError, ContractArtifacts, EnhancedHintError},
     CheatnetState,
 };
-use num_traits::Num;
 use num_traits::ToPrimitive;
 use serde::Deserialize;
 use starknet_api::core::{ContractAddress, PatriciaKey};
@@ -237,8 +237,7 @@ impl CairoHintProcessor<'_> {
                         .collect(),
                 ) {
                     Ok(class_hash) => {
-                        let felt_class_hash =
-                            felt252_from_hex_string(&class_hash.to_string()).unwrap();
+                        let felt_class_hash = stark_felt_to_felt(class_hash.0);
 
                         buffer
                             .write(Felt252::from(0))
@@ -287,6 +286,22 @@ impl CairoHintProcessor<'_> {
                 print(inputs);
                 Ok(())
             }
+            "get_class_hash" => {
+                let contract_address = contract_address_from_felt252(&inputs[0])?;
+
+                match self.cheatnet_state.get_class_hash(contract_address) {
+                    Ok(class_hash) => {
+                        let felt_class_hash = stark_felt_to_felt(class_hash.0);
+
+                        buffer
+                            .write(felt_class_hash)
+                            .expect("Failed to insert contract class hash");
+                        Ok(())
+                    }
+                    Err(CheatcodeError::Recoverable(_)) => unreachable!(),
+                    Err(CheatcodeError::Unrecoverable(err)) => Err(err),
+                }
+            }
             _ => Err(anyhow!("Unknown cheatcode selector: {selector}")).map_err(Into::into),
         }?;
 
@@ -376,12 +391,6 @@ fn print(inputs: Vec<Felt252>) {
     }
 }
 
-fn felt252_from_hex_string(value: &str) -> Result<Felt252> {
-    let stripped_value = value.replace("0x", "");
-    Felt252::from_str_radix(&stripped_value, 16)
-        .map_err(|_| anyhow!("Failed to convert value = {value} to Felt252"))
-}
-
 fn write_cheatcode_panic(buffer: &mut MemBuffer, panic_data: &[Felt252]) {
     buffer.write(1).expect("Failed to insert err code");
     buffer
@@ -392,30 +401,8 @@ fn write_cheatcode_panic(buffer: &mut MemBuffer, panic_data: &[Felt252]) {
         .expect("Failed to insert error in memory");
 }
 
-#[cfg(test)]
-mod test {
-    use super::*;
-
-    #[test]
-    fn felt_2525_from_prefixed_hex() {
-        assert_eq!(
-            felt252_from_hex_string("0x1234").unwrap(),
-            Felt252::from(0x1234)
-        );
-    }
-
-    #[test]
-    fn felt_2525_from_non_prefixed_hex() {
-        assert_eq!(
-            felt252_from_hex_string("1234").unwrap(),
-            Felt252::from(0x1234)
-        );
-    }
-
-    #[test]
-    fn felt_252_err_on_failed_conversion() {
-        let result = felt252_from_hex_string("yyyy");
-        let err = result.unwrap_err();
-        assert_eq!(err.to_string(), "Failed to convert value = yyyy to Felt252");
-    }
+fn contract_address_from_felt252(felt: &Felt252) -> Result<ContractAddress, EnhancedHintError> {
+    Ok(ContractAddress(PatriciaKey::try_from(felt_to_stark_felt(
+        felt,
+    ))?))
 }

crates/forge/tests/data/contracts/get_class_hash_checker.cairo

@@ -0,0 +1,32 @@
+use starknet::ClassHash;
+
+#[starknet::interface]
+trait IUpgradeable<T> {
+    fn upgrade(ref self: T, class_hash: ClassHash);
+}
+
+#[starknet::contract]
+mod GetClassHashCheckerUpg {
+
+    use starknet::ClassHash;
+    use result::ResultTrait;
+
+    #[storage]
+    struct Storage {
+        inner: felt252,
+    }
+
+    #[external(v0)]
+    impl IUpgradeableImpl of super::IUpgradeable<ContractState> {
+        fn upgrade(ref self: ContractState, class_hash: ClassHash) {
+            _upgrade(class_hash);
+        }
+    }
+
+    fn _upgrade(class_hash: ClassHash) {
+        match starknet::replace_class_syscall(class_hash) {
+            Result::Ok(()) => {},
+            Result::Err(e) => panic(e),
+        };
+    }
+}

crates/forge/tests/integration/get_class_hash.rs

@@ -0,0 +1,120 @@
+use crate::integration::common::corelib::{corelib_path, predeployed_contracts};
+use crate::integration::common::runner::Contract;
+use crate::{assert_passed, test_case};
+use camino::Utf8PathBuf;
+use forge::run;
+use indoc::indoc;
+use std::path::Path;
+
+#[test]
+fn get_class_hash() {
+    let test = test_case!(
+        indoc!(
+            r#"
+            use array::ArrayTrait;
+            use result::ResultTrait;
+            use snforge_std::{ declare, PreparedContract, deploy, start_prank, get_class_hash };
+
+            #[test]
+            fn test_get_class_hash() {
+                let class_hash = declare('GetClassHashCheckerUpg');
+                let prepared = PreparedContract { class_hash: class_hash, constructor_calldata: @ArrayTrait::new() };
+                let contract_address = deploy(prepared).unwrap();
+                assert(get_class_hash(contract_address) == class_hash, 'Incorrect class hash');
+            }
+        "#
+        ),
+        Contract::from_code_path(
+            "GetClassHashCheckerUpg".to_string(),
+            Path::new("tests/data/contracts/get_class_hash_checker.cairo"),
+        )
+        .unwrap()
+    );
+
+    let result = run(
+        &test.path().unwrap(),
+        &String::from("src"),
+        &test.path().unwrap().join("src/lib.cairo"),
+        &Some(test.linked_libraries()),
+        &Default::default(),
+        &corelib_path(),
+        &test.contracts(&corelib_path()).unwrap(),
+        &Utf8PathBuf::from_path_buf(predeployed_contracts().to_path_buf()).unwrap(),
+    )
+    .unwrap();
+
+    assert_passed!(result);
+}
+
+#[test]
+fn get_class_hash_replace_class() {
+    let test = test_case!(
+        indoc!(
+            r#"
+            use array::{ArrayTrait, SpanTrait};
+            use core::result::ResultTrait;
+            use starknet::ClassHash;
+            use snforge_std::{declare, deploy, get_class_hash, PreparedContract};
+
+            #[starknet::interface]
+            trait IUpgradeable<T> {
+                fn upgrade(ref self: T, class_hash: ClassHash);
+            }
+
+            #[starknet::interface]
+            trait IHelloStarknet<TContractState> {
+                fn increase_balance(ref self: TContractState, amount: felt252);
+                fn get_balance(self: @TContractState) -> felt252;
+                fn do_a_panic(self: @TContractState);
+                fn do_a_panic_with(self: @TContractState, panic_data: Array<felt252>);
+            }
+
+            #[test]
+            fn test_get_class_hash_replace_class() {
+                let class_hash = declare('GetClassHashCheckerUpg');
+
+                let prepared = PreparedContract {
+                    class_hash: class_hash,
+                    constructor_calldata: @ArrayTrait::new()
+                };
+
+                let contract_address = deploy(prepared).unwrap();
+
+                assert(get_class_hash(contract_address) == class_hash, 'Incorrect class hash');
+
+                let hsn_class_hash = declare('HelloStarknet');
+                IUpgradeableDispatcher { contract_address }.upgrade(hsn_class_hash);
+                assert(get_class_hash(contract_address) == hsn_class_hash, 'Incorrect upgrade class hash');
+
+                let hello_dispatcher = IHelloStarknetDispatcher { contract_address };
+                hello_dispatcher.increase_balance(42);
+                assert(hello_dispatcher.get_balance() == 42, 'Invalid balance');
+            }
+        "#
+        ),
+        Contract::from_code_path(
+            "GetClassHashCheckerUpg".to_string(),
+            Path::new("tests/data/contracts/get_class_hash_checker.cairo"),
+        )
+        .unwrap(),
+        Contract::from_code_path(
+            "HelloStarknet".to_string(),
+            Path::new("tests/data/contracts/hello_starknet.cairo"),
+        )
+        .unwrap()
+    );
+
+    let result = run(
+        &test.path().unwrap(),
+        &String::from("src"),
+        &test.path().unwrap().join("src/lib.cairo"),
+        &Some(test.linked_libraries()),
+        &Default::default(),
+        &corelib_path(),
+        &test.contracts(&corelib_path()).unwrap(),
+        &Utf8PathBuf::from_path_buf(predeployed_contracts().to_path_buf()).unwrap(),
+    )
+    .unwrap();
+
+    assert_passed!(result);
+}

crates/forge/tests/integration/mod.rs

@@ -2,6 +2,7 @@ pub(crate) mod common;
 mod declare;
 mod deploy;
 mod dispatchers;
+mod get_class_hash;
 mod prank;
 mod pure_cairo;
 mod roll;

docs/src/SUMMARY.md

@@ -44,6 +44,7 @@
   * [stop_roll](appendix/cheatcodes/stop_roll.md)
   * [start_warp](appendix/cheatcodes/start_warp.md)
   * [stop_warp](appendix/cheatcodes/stop_warp.md)
+  * [get_class_hash](appendix/cheatcodes/get_class_hash.md)
 * [Forge Library Functions References](appendix/forge-library.md)
     * [declare](appendix/forge-library/declare.md)
     * [deploy](appendix/forge-library/deploy.md)

docs/src/appendix/cheatcodes.md

@@ -6,6 +6,7 @@
 * [`stop_roll`](cheatcodes/stop_roll.md)
 * [`start_warp`](cheatcodes/start_warp.md)
 * [`stop_warp`](cheatcodes/stop_warp.md)
+* [`get_class_hash`](cheatcodes/get_class_hash.md)
 
 > ℹ️ **Info**
 > To use cheatcodes you need to add `snforge_std` package as a dependency in

docs/src/appendix/cheatcodes/get_class_hash.md

@@ -0,0 +1,44 @@
+# `get_class_hash`
+
+> `fn get_class_hash(contract_address: ContractAddress) -> ClassHash`
+
+Retrieves a class hash of a contract deployed under the given address.
+
+- `contract_address` - target contract address
+
+The main purpose of this cheatcode is to test upgradable contracts. For contract implementation:
+
+```rust
+    // ...
+    #[external(v0)]
+    impl IUpgradeableImpl of super::IUpgradeable<ContractState> {
+        fn upgrade(ref self: ContractState, class_hash: starknet::ClassHash) {
+            starknet::replace_class_syscall(class_hash).unwrap_syscall();
+        }
+    }
+    // ...
+```
+
+We can use `get_class_hash` to check if it upgraded properly:
+
+```rust
+    #[test]
+    fn test_get_class_hash() {
+        let class_hash = declare('Contract1');
+
+        let prepared = PreparedContract {
+            class_hash: class_hash,
+            constructor_calldata: @ArrayTrait::new()
+        };
+
+        let contract_address = deploy(prepared).unwrap();
+
+        assert(get_class_hash(contract_address) == class_hash, 'Incorrect class hash');
+
+        let other_class_hash = declare('OtherContract');
+
+        IUpgradeableDispatcher { contract_address }.upgrade(other_class_hash);
+
+        assert(get_class_hash(contract_address) == other_class_hash, 'Incorrect class hash upgrade');
+    }
+```