Improve inference of the compile_data attribute

Author: mrkkrp

proto/messages.proto

@@ -31,8 +31,9 @@ message RustImportsResponse {
     repeated string imports = 2;
     repeated string test_imports = 3;
     repeated string extern_mods = 4;
-    bool success = 5;
-    string error_msg = 6;
+    repeated string compile_data = 5;
+    bool success = 6;
+    string error_msg = 7;
 }
 
 message LockfileCratesRequest {

rust_language/generate.go

@@ -431,10 +431,14 @@ func (l *rustLang) generateCargoRule(c *config.Config, args *language.GenerateAr
 	}
 
 	srcs := []string{}
+	compile_data := map[string]struct{}{"Cargo.toml": true}
 	responses := []*pb.RustImportsResponse{}
 
 	for src, response := range importsResponses {
 		srcs = append(srcs, src)
+		for _, f := range response.CompileData {
+			compile_data[f] = true
+		}
 		if response != nil {
 			responses = append(responses, response)
 		}
@@ -446,7 +450,7 @@ func (l *rustLang) generateCargoRule(c *config.Config, args *language.GenerateAr
 		newRule.SetAttr("srcs", srcs)
 	}
 	newRule.SetAttr("visibility", []string{"//visibility:public"})
-	newRule.SetAttr("compile_data", []string{"Cargo.toml"})
+	newRule.SetAttr("compile_data", setToVector(compile_data))
 
 	if targetName != crateName {
 		newRule.SetAttr("crate_name", crateName)
@@ -494,10 +498,14 @@ func (l *rustLang) generateBuildScript(c *config.Config, args *language.Generate
 	l.discoverModule(c, "build.rs", enabledFeatures, args, &importsResponses, true)
 
 	srcs := []string{}
+	compile_data := map[string]struct{}{"Cargo.toml": true}
 	responses := []*pb.RustImportsResponse{}
 
 	for src, response := range importsResponses {
 		srcs = append(srcs, src)
+		for _, f := range response.CompileData {
+			compile_data[f] = true
+		}
 		if response != nil {
 			responses = append(responses, response)
 		}
@@ -506,7 +514,7 @@ func (l *rustLang) generateBuildScript(c *config.Config, args *language.Generate
 	newRule := rule.NewRule("cargo_build_script", "build_script")
 	newRule.SetAttr("srcs", srcs)
 	newRule.SetAttr("visibility", []string{"//visibility:public"})
-	newRule.SetAttr("compile_data", []string{"Cargo.toml"})
+	newRule.SetAttr("compile_data", setToVector(compile_data))
 	newRule.SetAttr("crate_root", "build.rs")
 
 	cfg := l.GetConfig(args.Config)
@@ -589,3 +597,11 @@ func fileExists(path string, args *language.GenerateArgs) bool {
 	_, err := os.Stat(fullPath)
 	return err == nil
 }
+
+func setToVector(src map[string]bool) []string {
+	result := []string{}
+	for key := range src {
+		result = append(result, key)
+	}
+	return result
+}

rust_parser/main.rs

@@ -39,6 +39,7 @@ fn handle_rust_imports_request(
             response.imports = rust_imports.imports;
             response.test_imports = rust_imports.test_imports;
             response.extern_mods = rust_imports.extern_mods;
+            response.compile_data = rust_imports.compile_data;
         }
         Err(err) => {
             // Don't crash gazelle if we encounter an error, instead bubble it up so that we can

rust_parser/parser.rs

@@ -4,7 +4,7 @@ use std::collections::{HashSet, VecDeque};
 use std::error::Error;
 use std::fs::File;
 use std::io::Read;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 
 use syn::ext::IdentExt;
 use syn::parse_file;
@@ -16,6 +16,7 @@ pub struct RustImports {
     pub imports: Vec<String>,
     pub test_imports: Vec<String>,
     pub extern_mods: Vec<String>,
+    pub compile_data: Vec<String>,
 }
 
 #[derive(Debug, Default)]
@@ -44,15 +45,16 @@ pub fn parse_imports(
     let mut contents = String::new();
     file.read_to_string(&mut contents)?;
 
-    parse_imports_from_str(&contents, enabled_features)
+    parse_imports_from_str(&contents, enabled_features, path)
 }
 
 pub fn parse_imports_from_str(
     contents: &str,
     enabled_features: &[String],
+    path: PathBuf,
 ) -> Result<RustImports, Box<dyn Error>> {
     let ast = parse_file(contents)?;
-    let mut visitor = AstVisitor::new(enabled_features);
+    let mut visitor = AstVisitor::new(enabled_features, path);
     visitor.visit_file(&ast);
 
     let mut root_scope = visitor.mod_stack.pop_back().expect("no root scope");
@@ -70,7 +72,8 @@ pub fn parse_imports_from_str(
         hints: visitor.hints,
         imports: filter_imports(root_scope.imports),
         test_imports: filter_imports(root_scope.test_imports),
-        extern_mods: visitor.extern_mods,
+        extern_mods: visitor.extern_mods.into_iter().collect(),
+        compile_data: visitor.compile_data.into_iter().collect(),
     })
 }
 
@@ -160,33 +163,39 @@ impl Scope<'_> {
 
 #[derive(Debug)]
 struct AstVisitor<'ast> {
+    /// The path of the file we are visiting.
+    this_file: PathBuf,
     /// stack of mods in scope
     mod_stack: VecDeque<Scope<'ast>>,
     /// all mods that are currently in scope (including parent scopes)
     scope_mods: HashSet<Ident<'ast>>,
     /// collected hints
     hints: Hints,
     /// bare mods defined in external files
-    extern_mods: Vec<String>,
+    extern_mods: HashSet<String>,
     /// mods that are disallowed from being added to the current scope; this is currently only used
     /// for a hack, see below
     mod_denylist: HashSet<Ident<'ast>>,
     /// Enabled features
     enabled_features: HashSet<String>,
+    /// Files that are included via include_str! and include_bytes! macros.
+    compile_data: HashSet<String>,
 }
 
 impl AstVisitor<'_> {
-    fn new(enabled_features: &[String]) -> Self {
+    fn new(enabled_features: &[String], path: PathBuf) -> Self {
         let mut mod_stack = VecDeque::new();
         mod_stack.push_back(Scope::default());
 
         Self {
+            this_file: path,
             mod_stack,
             scope_mods: HashSet::default(),
             hints: Hints::default(),
-            extern_mods: Vec::default(),
+            extern_mods: HashSet::new(),
             mod_denylist: HashSet::new(),
             enabled_features: enabled_features.iter().cloned().collect(),
+            compile_data: HashSet::new(),
         }
     }
 }
@@ -535,7 +544,7 @@ impl<'ast> Visit<'ast> for AstVisitor<'ast> {
 
         if self.is_root_scope() && node.content.is_none() {
             // this mod is defined in a different file
-            self.extern_mods.push(node.ident.unraw().to_string());
+            self.extern_mods.insert(node.ident.unraw().to_string());
         }
 
         self.add_mod(&node.ident);
@@ -605,6 +614,48 @@ impl<'ast> Visit<'ast> for AstVisitor<'ast> {
         }
         visit::visit_item_macro(self, node);
     }
+
+    fn visit_expr_macro(&mut self, node: &'ast syn::ExprMacro) {
+        let macro_ident = node.mac.path.get_ident();
+
+        if let Some(ident) = macro_ident {
+            if ident == "include_str" || ident == "include_bytes" {
+                // Respect ignored scopes
+                if self.is_ignored_scope() {
+                    return;
+                }
+
+                // Parse macro input: expect a single string literal
+                if let Ok(syn::Expr::Lit(syn::ExprLit {
+                    lit: syn::Lit::Str(lit),
+                    ..
+                })) = syn::parse2::<syn::Expr>(node.mac.tokens.clone())
+                {
+                    let included_path = PathBuf::from(lit.value());
+
+                    if included_path.is_absolute() {
+                        panic!(
+                            "included paths must not be absolute: {}",
+                            included_path.display()
+                        )
+                    } else {
+                        let dir = self
+                            .this_file
+                            .parent()
+                            .unwrap_or_else(|| Path::new(""))
+                            .to_path_buf();
+                        let combined = dir.join(included_path);
+                        match normalize_path(&combined).to_str() {
+                            None => panic!("Invalid unicode in the path: {}", combined.display()),
+                            Some(x) => self.compile_data.insert(x.to_string()),
+                        };
+                    }
+                }
+            }
+        }
+
+        visit::visit_expr_macro(self, node);
+    }
 }
 
 fn parse_use_imports<'ast>(use_tree: &'ast syn::UseTree, imports: &mut HashSet<Ident<'ast>>) {
@@ -623,3 +674,30 @@ fn parse_use_imports<'ast>(use_tree: &'ast syn::UseTree, imports: &mut HashSet<I
         _ => (),
     }
 }
+
+/// Normalize a path by resolving `.` and `..` without touching the filesystem
+fn normalize_path(path: &Path) -> PathBuf {
+    // let components = path.components().peekable();
+    let mut stack = Vec::new();
+
+    for component in path.components() {
+        match component {
+            std::path::Component::ParentDir => {
+                if let Some(last) = stack.last() {
+                    // Only pop if last is a normal component, not RootDir
+                    if matches!(last, std::path::Component::Normal(_)) {
+                        stack.pop();
+                    } else {
+                        stack.push(component);
+                    }
+                } else {
+                    stack.push(component);
+                }
+            }
+            std::path::Component::CurDir => { /* skip `.` */ }
+            _ => stack.push(component),
+        }
+    }
+
+    stack.iter().collect()
+}