From bb3a29b05fdaacd94ce46e7e666718bce8a30085 Mon Sep 17 00:00:00 2001
From: Karol Sewilo <karol.sewilo@swmansion.com>
Date: Thu, 12 Sep 2024 12:53:04 +0200
Subject: [PATCH] Enhance User Project Identification in Cairo Coverage

commit-id:77eb0f76
---
 CHANGELOG.md                              |  1 +
 crates/cairo-coverage/src/input/filter.rs | 36 +++++++++++++++++------
 crates/cairo-coverage/src/main.rs         |  4 +--
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2dbf724..ce80ee3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 #### Fixed
 
 - Bug where hit count was not correctly calculated for functions declared at same line
+- Bug where functions not from user project were included in coverage report if `SCARB_CACHE` was set
 
 #### Changed
 
diff --git a/crates/cairo-coverage/src/input/filter.rs b/crates/cairo-coverage/src/input/filter.rs
index 7f33192..97d4a09 100644
--- a/crates/cairo-coverage/src/input/filter.rs
+++ b/crates/cairo-coverage/src/input/filter.rs
@@ -1,6 +1,8 @@
 use crate::cli::IncludedComponent;
 use crate::data_loader::LoadedData;
 use crate::input::sierra_to_cairo_map::StatementOrigin;
+use anyhow::{Context, Result};
+use camino::Utf8PathBuf;
 use regex::Regex;
 use std::collections::HashSet;
 use std::iter::once;
@@ -8,6 +10,7 @@ use std::sync::LazyLock;
 
 pub static VIRTUAL_FILE_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\[.*?]").unwrap());
 const SNFORGE_TEST_EXECUTABLE: &str = "snforge_internal_test_executable";
+const SNFORGE_SIERRA_DIR: &str = ".snfoundry_versioned_programs";
 
 #[derive(Eq, PartialEq, Hash)]
 enum StatementCategory {
@@ -27,12 +30,17 @@ impl From<IncludedComponent> for StatementCategory {
 }
 
 pub struct Filter {
+    user_project_path: String,
     allowed_statement_categories: HashSet<StatementCategory>,
     test_functions: HashSet<String>,
 }
 
 impl Filter {
-    pub fn new(included_component: &[IncludedComponent], loaded_data: &LoadedData) -> Self {
+    pub fn new(
+        source_sierra_path: &str,
+        included_component: &[IncludedComponent],
+        loaded_data: &LoadedData,
+    ) -> Result<Self> {
         let test_functions = loaded_data
             .debug_info
             .executables
@@ -49,10 +57,12 @@ impl Filter {
             .chain(once(StatementCategory::UserFunction))
             .collect();
 
-        Self {
+        let user_project_path = find_user_project_path(source_sierra_path)?;
+        Ok(Self {
+            user_project_path,
             allowed_statement_categories,
             test_functions,
-        }
+        })
     }
 
     pub fn should_include(&self, statement_origin: &StatementOrigin) -> bool {
@@ -76,15 +86,23 @@ impl Filter {
             labels.insert(StatementCategory::Macro);
         }
 
-        // TODO(#55)
-        // TODO: We should probably filter by path to user project not by path to cache
-        // TODO: Can get this from source_sierra_path in call trace
-        if file_location.contains("com.swmansion.scarb") || file_location.contains(".cache/scarb") {
-            labels.insert(StatementCategory::NonUserFunction);
-        } else {
+        if file_location.contains(&self.user_project_path) {
             labels.insert(StatementCategory::UserFunction);
+        } else {
+            labels.insert(StatementCategory::NonUserFunction);
         }
 
         labels
     }
 }
+
+fn find_user_project_path(source_sierra_path: &str) -> Result<String> {
+    Utf8PathBuf::from(source_sierra_path)
+        .parent()
+        .filter(|parent| parent.ends_with(SNFORGE_SIERRA_DIR))
+        .and_then(|parent| parent.parent())
+        .map(ToString::to_string)
+        .context(format!(
+            "Source sierra path should be in the format: <project_root>/{SNFORGE_SIERRA_DIR}/<file>.sierra.json, got: {source_sierra_path}"
+        ))
+}
diff --git a/crates/cairo-coverage/src/main.rs b/crates/cairo-coverage/src/main.rs
index bcfdaea..ce2966d 100644
--- a/crates/cairo-coverage/src/main.rs
+++ b/crates/cairo-coverage/src/main.rs
@@ -27,8 +27,8 @@ fn main() -> Result<()> {
         .context(format!("Failed to open output file at path: {output_path}"))?;
 
     let loaded_data = LoadedDataMap::load(&cli.trace_files)?;
-    for (_, loaded_data) in loaded_data.iter() {
-        let filter = Filter::new(&cli.include, loaded_data);
+    for (source_sierra_path, loaded_data) in loaded_data.iter() {
+        let filter = Filter::new(source_sierra_path, &cli.include, loaded_data)?;
         let input_data = InputData::new(loaded_data, &filter)?;
         let coverage_data = create_files_coverage_data_with_hits(&input_data);
         let output_data = LcovFormat::from(coverage_data);