Support contracts (#64)

Closes #43

CHANGELOG.md

@@ -9,7 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 #### Added
 
+- Support for contracts
 - Option to not include macros in coverage report. To get the same behavior as before use `--include macros`
+- `--project-path` flag to specify the path to the project root directory. This useful when inference fails 
 
 #### Fixed
 

Cargo.toml

@@ -15,6 +15,7 @@ camino = "1.1.9"
 clap = { version = "4.5.17", features = ["derive"] }
 cairo-lang-sierra = "2.8.2"
 cairo-lang-sierra-to-casm = "2.8.2"
+cairo-lang-starknet-classes = "2.8.2"
 derived-deref = "2.1.0"
 itertools = "0.13.0"
 serde = "1.0.210"

README.md

@@ -9,10 +9,8 @@
 > 
 > We currently don't support:
 > - Branch coverage
-> - Contracts
 > 
 > Things that might not work as expected:
-> - Macros coverage
 > - Counters for how many times line was executed
 
 ## Installation

crates/cairo-coverage/Cargo.toml

@@ -8,15 +8,17 @@ anyhow.workspace = true
 camino.workspace = true
 cairo-lang-sierra.workspace = true
 cairo-lang-sierra-to-casm.workspace = true
+cairo-lang-starknet-classes.workspace = true
 clap.workspace = true
 derived-deref.workspace = true
 itertools.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 trace-data.workspace = true
 regex.workspace = true
+indoc.workspace = true
 
 [dev-dependencies]
 assert_fs.workspace = true
 snapbox.workspace = true
-indoc.workspace = true
+

crates/cairo-coverage/src/cli.rs

@@ -16,6 +16,10 @@ pub struct Cli {
     /// Include additional components in the coverage report.
     #[arg(long, short, num_args = 1..)]
     pub include: Vec<IncludedComponent>,
+
+    /// Path to the project directory. If not provided, the project directory is inferred from the trace.
+    #[arg(value_parser = parse_project_path, long)]
+    pub project_path: Option<Utf8PathBuf>,
 }
 
 #[derive(ValueEnum, Debug, Clone, Eq, PartialEq)]
@@ -38,3 +42,12 @@ fn parse_trace_file(path: &str) -> Result<Utf8PathBuf> {
 
     Ok(trace_file)
 }
+
+fn parse_project_path(path: &str) -> Result<Utf8PathBuf> {
+    let project_path = Utf8PathBuf::from(path);
+
+    ensure!(project_path.exists(), "Project path does not exist");
+    ensure!(project_path.is_dir(), "Project path is not a directory");
+
+    Ok(project_path)
+}

crates/cairo-coverage/src/data_loader/loaded_data.rs

@@ -1,65 +1,84 @@
-use anyhow::Context;
-use anyhow::Result;
+use crate::data_loader::sierra_program::{GetDebugInfos, SierraProgram};
+use anyhow::{Context, Result};
 use cairo_lang_sierra::debug_info::DebugInfo;
-use cairo_lang_sierra::program::{Program, ProgramArtifact, VersionedProgram};
+use cairo_lang_sierra_to_casm::compiler::CairoProgramDebugInfo;
 use camino::Utf8PathBuf;
 use derived_deref::Deref;
 use serde::de::DeserializeOwned;
 use std::collections::HashMap;
 use std::fs;
-use trace_data::CallTrace;
-
-type SourceSierraPath = String;
+use trace_data::{CallTrace, CallTraceNode, CasmLevelInfo};
 
 #[derive(Deref)]
-pub struct LoadedDataMap(HashMap<SourceSierraPath, LoadedData>);
+pub struct LoadedDataMap(HashMap<Utf8PathBuf, LoadedData>);
 
 pub struct LoadedData {
-    pub program: Program,
     pub debug_info: DebugInfo,
-    pub call_traces: Vec<CallTrace>,
+    pub casm_level_infos: Vec<CasmLevelInfo>,
+    pub casm_debug_info: CairoProgramDebugInfo,
 }
 
 impl LoadedDataMap {
-    pub fn load(call_trace_paths: &Vec<Utf8PathBuf>) -> Result<Self> {
-        let mut map: HashMap<SourceSierraPath, LoadedData> = HashMap::new();
-        for call_trace_path in call_trace_paths {
-            let call_trace: CallTrace = read_and_deserialize(call_trace_path)?;
+    pub fn load(call_trace_paths: &[Utf8PathBuf]) -> Result<Self> {
+        let execution_infos = call_trace_paths
+            .iter()
+            .map(read_and_deserialize)
+            .collect::<Result<Vec<_>>>()?
+            .into_iter()
+            .flat_map(load_nested_traces)
+            .filter_map(|call_trace| call_trace.cairo_execution_info)
+            .collect::<Vec<_>>();
 
-            let source_sierra_path = &call_trace
-                .cairo_execution_info
-                .as_ref()
-                .context("Missing key 'cairo_execution_info' in call trace. Perhaps you have outdated scarb?")?
-                .source_sierra_path;
+        // OPTIMIZATION:
+        // Group execution info by source Sierra path
+        // so that the same Sierra program does not need to be deserialized multiple times.
+        let execution_infos_by_sierra_path = execution_infos.into_iter().fold(
+            HashMap::new(),
+            |mut acc: HashMap<_, Vec<_>>, execution_info| {
+                acc.entry(execution_info.source_sierra_path)
+                    .or_default()
+                    .push(execution_info.casm_level_info);
+                acc
+            },
+        );
 
-            if let Some(loaded_data) = map.get_mut(&source_sierra_path.to_string()) {
-                loaded_data.call_traces.push(call_trace);
-            } else {
-                let VersionedProgram::V1 {
-                    program:
-                        ProgramArtifact {
-                            program,
+        Ok(Self(
+            execution_infos_by_sierra_path
+                .into_iter()
+                .map(|(source_sierra_path, casm_level_infos)| {
+                    read_and_deserialize::<SierraProgram>(&source_sierra_path)?
+                        .compile_and_get_debug_infos()
+                        .map(|(debug_info, casm_debug_info)| LoadedData {
                             debug_info,
-                        },
-                    ..
-                } = read_and_deserialize(source_sierra_path)?;
+                            casm_level_infos,
+                            casm_debug_info,
+                        })
+                        .context(format!(
+                            "Error occurred while loading program from: {source_sierra_path}"
+                        ))
+                        .map(|loaded_data| (source_sierra_path, loaded_data))
+                })
+                .collect::<Result<_>>()?,
+        ))
+    }
+}
 
-                map.insert(
-                    source_sierra_path.to_string(),
-                    LoadedData {
-                        program,
-                        debug_info: debug_info
-                            .context(format!("Debug info not found in: {source_sierra_path}"))?,
-                        call_traces: vec![call_trace],
-                    },
-                );
+fn load_nested_traces(call_trace: CallTrace) -> Vec<CallTrace> {
+    fn load_recursively(call_trace: CallTrace, acc: &mut Vec<CallTrace>) {
+        acc.push(call_trace.clone());
+        for call_trace_node in call_trace.nested_calls {
+            if let CallTraceNode::EntryPointCall(nested_call_trace) = call_trace_node {
+                load_recursively(nested_call_trace, acc);
             }
         }
-        Ok(Self(map))
     }
+
+    let mut call_traces = Vec::new();
+    load_recursively(call_trace, &mut call_traces);
+    call_traces
 }
 
-fn read_and_deserialize<T: DeserializeOwned>(file_path: &Utf8PathBuf) -> anyhow::Result<T> {
+fn read_and_deserialize<T: DeserializeOwned>(file_path: &Utf8PathBuf) -> Result<T> {
     fs::read_to_string(file_path)
         .context(format!("Failed to read file at path: {file_path}"))
         .and_then(|content| {

crates/cairo-coverage/src/data_loader/mod.rs

@@ -1,4 +1,5 @@
 mod loaded_data;
+mod sierra_program;
 mod types;
 
 pub use loaded_data::{LoadedData, LoadedDataMap};

crates/cairo-coverage/src/data_loader/sierra_program.rs

@@ -0,0 +1,85 @@
+use anyhow::{Context, Result};
+use cairo_lang_sierra::debug_info::DebugInfo;
+use cairo_lang_sierra::program::{Program, ProgramArtifact, VersionedProgram};
+use cairo_lang_sierra_to_casm::compiler::{CairoProgramDebugInfo, SierraToCasmConfig};
+use cairo_lang_sierra_to_casm::metadata::{calc_metadata, MetadataComputationConfig};
+use cairo_lang_starknet_classes::casm_contract_class::CasmContractClass;
+use cairo_lang_starknet_classes::contract_class::ContractClass;
+use serde::Deserialize;
+
+#[derive(Deserialize)]
+#[serde(untagged)]
+pub enum SierraProgram {
+    VersionedProgram(VersionedProgram),
+    ContractClass(ContractClass),
+}
+
+pub trait GetDebugInfos {
+    fn compile_and_get_debug_infos(self) -> Result<(DebugInfo, CairoProgramDebugInfo)>;
+}
+
+impl GetDebugInfos for VersionedProgram {
+    fn compile_and_get_debug_infos(self) -> Result<(DebugInfo, CairoProgramDebugInfo)> {
+        let VersionedProgram::V1 {
+            program:
+                ProgramArtifact {
+                    program,
+                    debug_info,
+                },
+            ..
+        } = self;
+
+        let debug_info = debug_info.context("Debug info not found in program")?;
+        let casm_debug_info = compile_program_to_casm_debug_info(&program)?;
+        Ok((debug_info, casm_debug_info))
+    }
+}
+
+impl GetDebugInfos for ContractClass {
+    fn compile_and_get_debug_infos(self) -> Result<(DebugInfo, CairoProgramDebugInfo)> {
+        let debug_info = self
+            .sierra_program_debug_info
+            .context("Debug info not found in contract")?;
+
+        // OPTIMIZATION:
+        // Debug info is unused in the compilation. This saves us a costly clone.
+        let casm_debug_info = compile_contract_class_to_casm_debug_info(ContractClass {
+            sierra_program_debug_info: None,
+            ..self
+        })?;
+
+        Ok((debug_info, casm_debug_info))
+    }
+}
+impl GetDebugInfos for SierraProgram {
+    fn compile_and_get_debug_infos(self) -> Result<(DebugInfo, CairoProgramDebugInfo)> {
+        match self {
+            SierraProgram::VersionedProgram(program) => program.compile_and_get_debug_infos(),
+            SierraProgram::ContractClass(contract_class) => {
+                contract_class.compile_and_get_debug_infos()
+            }
+        }
+    }
+}
+
+fn compile_program_to_casm_debug_info(program: &Program) -> Result<CairoProgramDebugInfo> {
+    cairo_lang_sierra_to_casm::compiler::compile(
+        program,
+        &calc_metadata(program, MetadataComputationConfig::default())
+            .context("Failed calculating Sierra variables")?,
+        SierraToCasmConfig {
+            gas_usage_check: false,
+            max_bytecode_size: usize::MAX,
+        },
+    )
+    .map(|casm| casm.debug_info)
+    .context("Failed to compile program to casm")
+}
+
+fn compile_contract_class_to_casm_debug_info(
+    contract_class: ContractClass,
+) -> Result<CairoProgramDebugInfo> {
+    CasmContractClass::from_contract_class_with_debug_info(contract_class, false, usize::MAX)
+        .map(|(_, casm_debug_info)| casm_debug_info)
+        .context("Failed to compile contract class to casm")
+}