Make compile_test_prepared_db API more configurable

crates/bin/cairo-run/Cargo.toml

@@ -15,3 +15,4 @@ cairo-lang-diagnostics = { path = "../../cairo-lang-diagnostics", version = "~2.
 cairo-lang-runner = { path = "../../cairo-lang-runner", version = "~2.8.2" }
 cairo-lang-sierra-generator = { path = "../../cairo-lang-sierra-generator", version = "~2.8.2" }
 cairo-lang-starknet = { path = "../../cairo-lang-starknet", version = "~2.8.2" }
+cairo-lang-utils = { path = "../../cairo-lang-utils", version = "~2.8.2" }

crates/bin/cairo-run/src/main.rs

@@ -14,7 +14,8 @@ use cairo_lang_runner::{ProfilingInfoCollectionConfig, SierraCasmRunner, Starkne
 use cairo_lang_sierra_generator::db::SierraGenGroup;
 use cairo_lang_sierra_generator::program_generator::SierraProgramWithDebug;
 use cairo_lang_sierra_generator::replace_ids::{DebugReplacer, SierraIdReplacer};
-use cairo_lang_starknet::contract::get_contracts_info;
+use cairo_lang_starknet::contract::{find_contracts, get_contracts_info};
+use cairo_lang_utils::Upcast;
 use clap::Parser;
 
 /// Compiles a Cairo project and runs the function `main`.
@@ -75,7 +76,8 @@ fn main() -> anyhow::Result<()> {
  anyhow::bail!("Program requires gas counter, please provide `--available-gas` argument.");
  }
 
- let contracts_info = get_contracts_info(db, main_crate_ids, &replacer)?;
+ let contracts = find_contracts((*db).upcast(), &main_crate_ids);
+ let contracts_info = get_contracts_info(db, contracts, &replacer)?;
  let sierra_program = replacer.apply(&sierra_program);
 
  let runner = SierraCasmRunner::new(

crates/cairo-lang-starknet/src/contract.rs

@@ -41,6 +41,7 @@ use crate::plugin::consts::{ABI_ATTR, ABI_ATTR_EMBED_V0_ARG};
 mod test;
 
 /// Represents a declaration of a contract.
+#[derive(Clone)]
 pub struct ContractDeclaration {
  /// The id of the module that defines the contract.
  pub submodule_id: SubmoduleId,
@@ -293,10 +294,9 @@ pub struct ContractInfo {
 /// Returns the list of functions in a given module.
 pub fn get_contracts_info<T: SierraIdReplacer>(
  db: &dyn SierraGenGroup,
- main_crate_ids: Vec<CrateId>,
+ contracts: Vec<ContractDeclaration>,
  replacer: &T,
 ) -> Result<OrderedHashMap<Felt252, ContractInfo>, anyhow::Error> {
- let contracts = find_contracts(db.upcast(), &main_crate_ids);
  let mut contracts_info = OrderedHashMap::default();
  for contract in contracts {
  let (class_hash, contract_info) = analyze_contract(db, &contract, replacer)?;

crates/cairo-lang-test-plugin/src/lib.rs

@@ -7,6 +7,7 @@ use cairo_lang_compiler::diagnostics::DiagnosticsReporter;
 use cairo_lang_compiler::get_sierra_program_for_functions;
 use cairo_lang_debug::DebugWithDb;
 use cairo_lang_defs::ids::{FreeFunctionId, FunctionWithBodyId, ModuleItemId};
+use cairo_lang_filesystem::db::FilesGroup;
 use cairo_lang_filesystem::ids::CrateId;
 use cairo_lang_lowering::ids::ConcreteFunctionWithBodyId;
 use cairo_lang_semantic::db::SemanticGroup;
@@ -23,7 +24,8 @@ use cairo_lang_sierra_generator::program_generator::SierraProgramWithDebug;
 use cairo_lang_sierra_generator::replace_ids::DebugReplacer;
 use cairo_lang_sierra_generator::statements_locations::StatementsLocations;
 use cairo_lang_starknet::contract::{
- find_contracts, get_contract_abi_functions, get_contracts_info, ContractInfo,
+ find_contracts, get_contract_abi_functions, get_contracts_info, ContractDeclaration,
+ ContractInfo,
 };
 use cairo_lang_starknet::plugin::consts::{CONSTRUCTOR_MODULE, EXTERNAL_MODULE, L1_HANDLER_MODULE};
 use cairo_lang_starknet_classes::casm_contract_class::ENTRY_POINT_COST;
@@ -52,6 +54,19 @@ pub struct TestsCompilationConfig {
  /// Adds the starknet contracts to the compiled tests.
  pub starknet: bool,
 
+ /// Contracts to compile.
+ /// If defined, only this contacts will be available in tests.
+ /// If not, all contracts from `contract_crate_ids` will be compiled.
+ pub contract_declarations: Option<Vec<ContractDeclaration>>,
+
+ /// Crates to be searched for contracts.
+ /// If not defined, all crates will be searched.
+ pub contract_crate_ids: Option<Vec<CrateId>>,
+
+ /// Crates to be searched for executable attributes.
+ /// If not defined, test crates will be searched.
+ pub executable_crate_ids: Option<Vec<CrateId>>,
+
  /// Adds mapping used by [cairo-profiler](https://github.com/software-mansion/cairo-profiler) to
  /// [Annotations] in [DebugInfo] in the compiled tests.
  pub add_statements_functions: bool,
@@ -75,12 +90,17 @@ pub struct TestsCompilationConfig {
 pub fn compile_test_prepared_db(
  db: &RootDatabase,
  tests_compilation_config: TestsCompilationConfig,
- main_crate_ids: Vec<CrateId>,
  test_crate_ids: Vec<CrateId>,
  diagnostics_reporter: DiagnosticsReporter<'_>,
 ) -> Result<TestCompilation> {
+ let contracts = tests_compilation_config.contract_declarations.unwrap_or_else(|| {
+ find_contracts(
+ db,
+ &tests_compilation_config.contract_crate_ids.unwrap_or_else(|| db.crates()),
+ )
+ });
  let all_entry_points = if tests_compilation_config.starknet {
- find_contracts(db, &main_crate_ids)
+ contracts
  .iter()
  .flat_map(|contract| {
  chain!(
@@ -96,7 +116,10 @@ pub fn compile_test_prepared_db(
  vec![]
  };
 
- let executable_functions = find_executable_function_ids(db, main_crate_ids.clone());
+ let executable_functions = find_executable_function_ids(
+ db,
+ tests_compilation_config.executable_crate_ids.unwrap_or_else(|| test_crate_ids.clone()),
+ );
  let all_tests = find_all_tests(db, test_crate_ids.clone());
 
  let func_ids = chain!(
@@ -157,7 +180,7 @@ pub fn compile_test_prepared_db(
  )
  })
  .collect_vec();
- let contracts_info = get_contracts_info(db, main_crate_ids.clone(), &replacer)?;
+ let contracts_info = get_contracts_info(db, contracts, &replacer)?;
  let sierra_program = ProgramArtifact::stripped(sierra_program).with_debug_info(DebugInfo {
  executables,
  annotations,

crates/cairo-lang-test-runner/src/lib.rs

@@ -71,6 +71,9 @@ impl TestRunner {
  starknet,
  add_statements_functions: config.run_profiler == RunProfilerConfig::Cairo,
  add_statements_code_locations: false,
+ contract_declarations: None,
+ contract_crate_ids: None,
+ executable_crate_ids: None,
  },
  )?;
  Ok(Self { compiler, config })
@@ -262,7 +265,6 @@ impl TestCompiler {
  compile_test_prepared_db(
  &self.db,
  self.config.clone(),
- self.main_crate_ids.clone(),
  self.test_crate_ids.clone(),
  diag_reporter,
  )

crates/cairo-lang-test-runner/src/test.rs

@@ -20,6 +20,9 @@ fn test_compiled_serialization() {
  starknet: true,
  add_statements_functions: false,
  add_statements_code_locations: false,
+ contract_declarations: None,
+ contract_crate_ids: None,
+ executable_crate_ids: None,
  },
  )
  .unwrap();