Add sierra path to contracts data

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/cheatcodes/declare.rs

@@ -18,7 +18,7 @@ pub fn declare(
  contract_name: &str,
  contracts_data: &ContractsData,
 ) -> Result<ClassHash, CheatcodeError> {
- let contract_artifact = contracts_data.contracts.get(contract_name).with_context(|| {
+ let contract_artifact = contracts_data.get_artifacts_for_contract(contract_name).with_context(|| {
  format!("Failed to get contract artifact for name = {contract_name}. Make sure starknet target is correctly defined in Scarb.toml file.")
  }).map_err::<EnhancedHintError, _>(From::from)?;
 
@@ -27,8 +27,7 @@ pub fn declare(
  let contract_class = BlockifierContractClass::V1(contract_class);
 
  let class_hash = *contracts_data
- .class_hashes
- .get_by_left(contract_name)
+ .get_class_hash_for_contract(contract_name)
  .expect("Failed to get class hash");
 
  match state.get_compiled_contract_class(class_hash) {

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/contracts_data.rs

@@ -1,6 +1,7 @@
 use super::cheatcodes::declare::get_class_hash;
 use anyhow::Result;
 use bimap::BiMap;
+use camino::Utf8PathBuf;
 use conversions::IntoConv;
 use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
 use scarb_api::StarknetContractArtifacts;
@@ -11,22 +12,49 @@ use std::collections::HashMap;
 
 #[derive(Debug, Clone)]
 pub struct ContractsData {
- pub contracts: HashMap<String, StarknetContractArtifacts>,
- pub class_hashes: BiMap<String, ClassHash>,
- pub selectors: HashMap<EntryPointSelector, String>,
+ contracts: HashMap<String, Contract>,
+ class_hash_index: BiMap<String, ClassHash>,
+ selectors: HashMap<EntryPointSelector, String>,
+}
+
+#[derive(Debug, Clone)]
+struct Contract {
+ artifacts: StarknetContractArtifacts,
+ class_hash: ClassHash,
+ _sierra_source_path: Utf8PathBuf,
 }
 
 impl ContractsData {
- pub fn try_from(contracts: HashMap<String, StarknetContractArtifacts>) -> Result<Self> {
+ pub fn try_from(
+ contracts: HashMap<String, (StarknetContractArtifacts, Utf8PathBuf)>,
+ ) -> Result<Self> {
  let parsed_contracts: HashMap<String, SierraClass> = contracts
  .par_iter()
- .map(|(name, artifact)| Ok((name.clone(), serde_json::from_str(&artifact.sierra)?)))
+ .map(|(name, (artifact, _))| {
+ Ok((name.clone(), serde_json::from_str(&artifact.sierra)?))
+ })
  .collect::<Result<_>>()?;
 
  let class_hashes: Vec<(String, ClassHash)> = parsed_contracts
  .par_iter()
  .map(|(name, sierra_class)| Ok((name.clone(), get_class_hash(sierra_class)?)))
  .collect::<Result<_>>()?;
+ let class_hash_index = BiMap::from_iter(class_hashes);
+
+ let contracts = contracts
+ .into_iter()
+ .map(|(name, (artifacts, sierra_source_path))| {
+ let class_hash = *class_hash_index.get_by_left(&name).unwrap();
+ (
+ name,
+ Contract {
+ artifacts,
+ class_hash,
+ _sierra_source_path: sierra_source_path,
+ },
+ )
+ })
+ .collect();
 
  let selectors = parsed_contracts
  .into_par_iter()
@@ -36,10 +64,35 @@ impl ContractsData {
 
  Ok(ContractsData {
  contracts,
- class_hashes: BiMap::from_iter(class_hashes),
+ class_hash_index,
  selectors,
  })
  }
+
+ #[must_use]
+ pub fn get_artifacts_for_contract(&self, name: &str) -> Option<&StarknetContractArtifacts> {
+ self.contracts.get(name).map(|contract| &contract.artifacts)
+ }
+
+ #[must_use]
+ pub fn get_class_hash_for_contract(&self, name: &str) -> Option<&ClassHash> {
+ self.contracts
+ .get(name)
+ .map(|contract| &contract.class_hash)
+ }
+
+ #[must_use]
+ pub fn get_contract_name_from_class_hash(&self, class_hash: &ClassHash) -> Option<&String> {
+ self.class_hash_index.get_by_right(class_hash)
+ }
+
+ #[must_use]
+ pub fn get_function_name_from_entry_point_selector(
+ &self,
+ entry_point_selector: &EntryPointSelector,
+ ) -> Option<&String> {
+ self.selectors.get(entry_point_selector)
+ }
 }
 
 fn build_name_selector_map(abi: Vec<AbiEntry>) -> HashMap<EntryPointSelector, String> {

crates/cheatnet/tests/cheatcodes/declare.rs

@@ -1,21 +1,7 @@
 use crate::common::{get_contracts, state::create_cached_state};
-use cheatnet::runtime_extensions::forge_runtime_extension::cheatcodes::declare::{
- declare, get_class_hash,
-};
+use cheatnet::runtime_extensions::forge_runtime_extension::cheatcodes::declare::declare;
 use cheatnet::runtime_extensions::forge_runtime_extension::cheatcodes::CheatcodeError;
 use runtime::EnhancedHintError;
-use scarb_api::StarknetContractArtifacts;
-use starknet_api::core::ClassHash;
-use std::collections::HashMap;
-
-fn get_contract_class_hash(
- contract_name: &str,
- contracts: &HashMap<String, StarknetContractArtifacts>,
-) -> ClassHash {
- let contract = contracts.get(contract_name).unwrap();
- let sierra_class = serde_json::from_str(&contract.sierra).unwrap();
- get_class_hash(&sierra_class).unwrap()
-}
 
 #[test]
 fn declare_simple() {
@@ -26,9 +12,11 @@ fn declare_simple() {
  let contracts_data = get_contracts();
 
  let class_hash = declare(&mut cached_state, contract_name, &contracts_data).unwrap();
- let expected_class_hash = get_contract_class_hash(contract_name, &contracts_data.contracts);
+ let expected_class_hash = contracts_data
+ .get_class_hash_for_contract(contract_name)
+ .unwrap();
 
- assert_eq!(class_hash, expected_class_hash);
+ assert_eq!(class_hash, *expected_class_hash);
 }
 
 #[test]
@@ -41,8 +29,10 @@ fn declare_multiple() {
 
  for contract_name in contract_names {
  let class_hash = declare(&mut cached_state, contract_name, &contracts_data).unwrap();
- let expected_class_hash = get_contract_class_hash(contract_name, &contracts_data.contracts);
- assert_eq!(class_hash, expected_class_hash);
+ let expected_class_hash = contracts_data
+ .get_class_hash_for_contract(contract_name)
+ .unwrap();
+ assert_eq!(class_hash, *expected_class_hash);
  }
 }
 
@@ -55,8 +45,10 @@ fn declare_same_contract() {
  let contracts_data = get_contracts();
 
  let class_hash = declare(&mut cached_state, contract_name, &contracts_data).unwrap();
- let expected_class_hash = get_contract_class_hash(contract_name, &contracts_data.contracts);
- assert_eq!(class_hash, expected_class_hash);
+ let expected_class_hash = contracts_data
+ .get_class_hash_for_contract(contract_name)
+ .unwrap();
+ assert_eq!(class_hash, *expected_class_hash);
 
  let output = declare(&mut cached_state, contract_name, &contracts_data);
 

crates/cheatnet/tests/common/mod.rs

@@ -24,7 +24,7 @@ use cheatnet::state::CheatnetState;
 use conversions::IntoConv;
 use runtime::starknet::context::build_context;
 use scarb_api::metadata::MetadataCommandExt;
-use scarb_api::{get_contracts_map, ScarbCommand};
+use scarb_api::{get_contracts_artifacts_and_source_sierra_paths, ScarbCommand};
 use starknet::core::utils::get_selector_from_name;
 use starknet_api::core::PatriciaKey;
 use starknet_api::core::{ClassHash, ContractAddress};
@@ -76,7 +76,10 @@ pub fn get_contracts() -> ContractsData {
 
  let package = scarb_metadata.packages.first().unwrap();
 
- ContractsData::try_from(get_contracts_map(&scarb_metadata, &package.id, None).unwrap()).unwrap()
+ let contracts =
+ get_contracts_artifacts_and_source_sierra_paths(&scarb_metadata, &package.id, None)
+ .unwrap();
+ ContractsData::try_from(contracts).unwrap()
 }
 
 pub fn deploy_contract(

crates/forge-runner/src/build_trace_data.rs

@@ -108,9 +108,11 @@ pub fn build_profiler_call_entry_point(
  } = value;
 
  let mut contract_name = class_hash
- .and_then(|c| contracts_data.class_hashes.get_by_right(&c))
+ .and_then(|c| contracts_data.get_contract_name_from_class_hash(&c))
+ .cloned();
+ let mut function_name = contracts_data
+ .get_function_name_from_entry_point_selector(&entry_point_selector)
  .cloned();
- let mut function_name = contracts_data.selectors.get(&entry_point_selector).cloned();
 
  if entry_point_selector.0
  == get_selector_from_name(TEST_ENTRY_POINT_SELECTOR)

crates/forge/src/compiled_raw.rs

@@ -4,14 +4,14 @@ use forge_runner::expected_result::ExpectedTestResult;
 use serde::Deserialize;
 
 #[derive(Debug, Clone, Deserialize)]
-pub(crate) struct CompiledTestCrateRaw {
+pub struct CompiledTestCrateRaw {
  pub sierra_program: VersionedProgram,
  pub test_cases: Vec<TestCaseRaw>,
  pub tests_location: CrateLocation,
 }
 
 #[derive(Debug, PartialEq, Clone, Deserialize)]
-pub(crate) struct TestCaseRaw {
+pub struct TestCaseRaw {
  pub name: String,
  pub available_gas: Option<usize>,
  pub ignored: bool,
@@ -22,15 +22,15 @@ pub(crate) struct TestCaseRaw {
 }
 
 #[derive(Debug, PartialEq, Clone, Copy, Deserialize)]
-pub(crate) enum CrateLocation {
+pub enum CrateLocation {
  /// Main crate in a package
  Lib,
  /// Crate in the `tests/` directory
  Tests,
 }
 
 #[derive(Debug, Clone, PartialEq, Deserialize)]
-pub(crate) enum RawForkConfig {
+pub enum RawForkConfig {
  Id(String),
  Params(RawForkParams),
 }

crates/forge/src/lib.rs

@@ -1,10 +1,8 @@
-use anyhow::{anyhow, Context, Result};
-use camino::Utf8Path;
+use anyhow::{anyhow, Result};
 use warn::{
  warn_if_available_gas_used_with_incompatible_scarb_version, warn_if_incompatible_rpc_version,
 };
 
-use crate::scarb::load_test_artifacts;
 use forge_runner::test_case_summary::AnyTestCaseSummary;
 use std::sync::Arc;
 
@@ -99,19 +97,20 @@ async fn to_runnable(
 /// * `fork_target` - A configuration of forks used in tests
 #[allow(clippy::implicit_hasher)]
 pub async fn run(
+ compiled_test_crates: Vec<CompiledTestCrateRaw>,
  package_name: &str,
- snforge_target_dir_path: &Utf8Path,
  tests_filter: &TestsFilter,
  runner_config: Arc<RunnerConfig>,
  runner_params: Arc<RunnerParams>,
  fork_targets: &[ForkTarget],
  block_number_map: &mut BlockNumberMap,
 ) -> Result<Vec<TestCrateSummary>> {
- let test_crates = load_test_artifacts(snforge_target_dir_path, package_name)
- .context("Failed to load test artifacts, make sure to use scarb >=2.5.4")?;
- let all_tests: usize = test_crates.iter().map(|tc| tc.test_cases.len()).sum();
+ let all_tests: usize = compiled_test_crates
+ .iter()
+ .map(|tc| tc.test_cases.len())
+ .sum();
 
- let test_crates = test_crates
+ let test_crates = compiled_test_crates
  .into_iter()
  .map(|tc| tests_filter.filter_tests(tc))
  .collect::<Result<Vec<CompiledTestCrateRaw>>>()?;

crates/forge/src/main.rs

@@ -3,7 +3,10 @@ use camino::Utf8Path;
 use clap::{Parser, Subcommand, ValueEnum};
 use configuration::load_package_config;
 use forge::scarb::config::ForgeConfig;
-use forge::scarb::{build_contracts_with_scarb, build_test_artifacts_with_scarb};
+use forge::scarb::{
+ build_contracts_with_scarb, build_test_artifacts_with_scarb, get_test_artifacts_path,
+ load_test_artifacts,
+};
 use forge::shared_cache::{clean_cache, set_cached_failed_tests_names};
 use forge::test_filter::TestsFilter;
 use forge::{pretty_printing, run};
@@ -12,7 +15,7 @@ use forge_runner::test_crate_summary::TestCrateSummary;
 use forge_runner::{RunnerConfig, RunnerParams, CACHE_DIR};
 use rand::{thread_rng, RngCore};
 use scarb_api::{
- get_contracts_map,
+ get_contracts_artifacts_and_source_sierra_paths,
  metadata::{Metadata, MetadataCommandExt, PackageMetadata},
  package_matches_version_requirement, target_dir_for_workspace, ScarbCommand,
 };
@@ -265,13 +268,19 @@ fn test_workspace(args: TestArgs) -> Result<bool> {
  for package in &packages {
  env::set_current_dir(&package.root)?;
 
- let forge_config =
- load_package_config::<ForgeConfig>(&scarb_metadata, &package.id)?;
- let contracts =
- get_contracts_map(&scarb_metadata, &package.id, None).unwrap_or_default();
+ let test_artifacts_path =
+ get_test_artifacts_path(&snforge_target_dir_path, &package.name);
+ let compiled_test_crates = load_test_artifacts(&test_artifacts_path)?;
 
+ let contracts = get_contracts_artifacts_and_source_sierra_paths(
+ &scarb_metadata,
+ &package.id,
+ None,
+ )?;
  let contracts_data = ContractsData::try_from(contracts)?;
 
+ let forge_config =
+ load_package_config::<ForgeConfig>(&scarb_metadata, &package.id)?;
  let runner_config = Arc::new(combine_configs(
  &workspace_root,
  args.exit_first,
@@ -287,8 +296,8 @@ fn test_workspace(args: TestArgs) -> Result<bool> {
  Arc::new(RunnerParams::new(contracts_data, env::vars().collect()));
 
  let tests_file_summaries = run(
+ compiled_test_crates,
  &package.name,
- &snforge_target_dir_path,
  &TestsFilter::from_flags(
  args.test_filter.clone(),
  args.exact,

crates/forge/src/scarb.rs

@@ -1,7 +1,7 @@
 use crate::compiled_raw::CompiledTestCrateRaw;
 use crate::scarb::config::{ForgeConfig, RawForgeConfig};
 use anyhow::{Context, Result};
-use camino::Utf8Path;
+use camino::{Utf8Path, Utf8PathBuf};
 use configuration::PackageConfig;
 use scarb_api::ScarbCommand;
 use scarb_ui::args::PackagesFilter;
@@ -43,16 +43,18 @@ pub fn build_test_artifacts_with_scarb(filter: PackagesFilter) -> Result<()> {
  Ok(())
 }
 
-pub(crate) fn load_test_artifacts(
+#[must_use]
+pub fn get_test_artifacts_path(
  snforge_target_dir_path: &Utf8Path,
  package_name: &str,
-) -> Result<Vec<CompiledTestCrateRaw>> {
- let snforge_test_artifact_path =
- snforge_target_dir_path.join(format!("{package_name}.snforge_sierra.json"));
- let test_crates = serde_json::from_str::<Vec<CompiledTestCrateRaw>>(&std::fs::read_to_string(
- snforge_test_artifact_path,
- )?)?;
- Ok(test_crates)
+) -> Utf8PathBuf {
+ snforge_target_dir_path.join(format!("{package_name}.snforge_sierra.json"))
+}
+
+pub fn load_test_artifacts(test_artifacts_path: &Utf8PathBuf) -> Result<Vec<CompiledTestCrateRaw>> {
+ Ok(serde_json::from_str::<Vec<CompiledTestCrateRaw>>(
+ &std::fs::read_to_string(test_artifacts_path)?,
+ )?)
 }
 
 #[cfg(test)]