Ensure special files are also indexed

tree-sitter-stack-graphs/src/cli/index.rs

@@ -7,6 +7,8 @@
 
 use clap::Args;
 use clap::ValueHint;
+use stack_graphs::arena::Handle;
+use stack_graphs::graph::File;
 use stack_graphs::graph::StackGraph;
 use stack_graphs::partial::PartialPaths;
 use stack_graphs::storage::FileStatus;
@@ -18,6 +20,7 @@ use std::time::Duration;
 use thiserror::Error;
 use tree_sitter_graph::Variables;
 
+use crate::loader::FileLanguageConfigurations;
 use crate::loader::FileReader;
 use crate::loader::Loader;
 use crate::BuildError;
@@ -29,6 +32,7 @@ use super::util::duration_from_seconds_str;
 use super::util::iter_files_and_directories;
 use super::util::sha1;
 use super::util::wait_for_input;
+use super::util::BuildErrorWithSource;
 use super::util::ConsoleLogger;
 use super::util::ExistingPathBufValueParser;
 use super::util::FileLogger;
@@ -219,23 +223,24 @@ impl<'a> Indexer<'a> {
         }
 
         let mut file_reader = FileReader::new();
-        let lc = match self
+        let lcs = match self
             .loader
             .load_for_file(source_path, &mut file_reader, &NoCancellation)
         {
-            Ok(Some(sgl)) => sgl,
-            Ok(None) => {
+            Ok(lcs) if !lcs.has_some() => {
                 if missing_is_error {
                     file_status.failure("not supported", None);
                 }
                 return Ok(());
             }
+            Ok(lcs) => lcs,
             Err(crate::loader::LoadError::Cancelled(_)) => {
                 file_status.warning("language loading timed out", None);
                 return Ok(());
             }
             Err(e) => return Err(IndexError::LoadError(e)),
         };
+
         let source = file_reader.get(source_path)?;
         let tag = sha1(source);
 
@@ -264,64 +269,39 @@ impl<'a> Indexer<'a> {
         let mut graph = StackGraph::new();
         let file = graph
             .add_file(&source_path.to_string_lossy())
-            .expect("file not present in emtpy graph");
+            .expect("file not present in empty graph");
 
-        let relative_source_path = source_path.strip_prefix(source_root).unwrap();
-        let result = if let Some(fa) = source_path
-            .file_name()
-            .and_then(|f| lc.special_files.get(&f.to_string_lossy()))
-        {
-            fa.build_stack_graph_into(
-                &mut graph,
-                file,
-                &relative_source_path,
-                &source,
-                &mut std::iter::empty(),
-                &HashMap::new(),
-                &cancellation_flag,
-            )
-        } else {
-            let globals = Variables::new();
-            lc.sgl
-                .build_stack_graph_into(&mut graph, file, &source, &globals, &cancellation_flag)
-        };
-        match result {
-            Err(BuildError::Cancelled(_)) => {
-                file_status.warning("parsing timed out", None);
-                self.db
-                    .store_error_for_file(source_path, &tag, "parsing timed out")?;
-                return Ok(());
-            }
-            Err(err @ BuildError::ParseErrors(_)) => {
-                file_status.failure(
-                    "parsing failed",
-                    Some(&err.display_pretty(
-                        source_path,
-                        source,
-                        lc.sgl.tsg_path(),
-                        lc.sgl.tsg_source(),
-                    )),
-                );
-                self.db.store_error_for_file(
-                    source_path,
-                    &tag,
-                    &format!("parsing failed: {}", err),
-                )?;
-                return Ok(());
-            }
-            Err(err) => {
-                file_status.failure(
-                    "failed to build stack graph",
-                    Some(&err.display_pretty(
+        let result = Self::build_stack_graph(
+            &mut graph,
+            file,
+            source_root,
+            source_path,
+            &source,
+            lcs,
+            &cancellation_flag,
+        );
+        if let Err(err) = result {
+            match err.inner {
+                BuildError::Cancelled(_) => {
+                    file_status.warning("parsing timed out", None);
+                    self.db
+                        .store_error_for_file(source_path, &tag, "parsing timed out")?;
+                    return Ok(());
+                }
+                BuildError::ParseErrors { .. } => {
+                    file_status.failure("parsing failed", Some(&err.display_pretty()));
+                    self.db.store_error_for_file(
                         source_path,
-                        source,
-                        lc.sgl.tsg_path(),
-                        lc.sgl.tsg_source(),
-                    )),
-                );
-                return Err(IndexError::StackGraph);
+                        &tag,
+                        &format!("parsing failed: {}", err.inner),
+                    )?;
+                    return Ok(());
+                }
+                _ => {
+                    file_status.failure("failed to build stack graph", Some(&err.display_pretty()));
+                    return Err(IndexError::StackGraph);
+                }
             }
-            Ok(_) => true,
         };
 
         let mut partials = PartialPaths::new();
@@ -354,6 +334,49 @@ impl<'a> Indexer<'a> {
         Ok(())
     }
 
+    fn build_stack_graph<'b>(
+        graph: &mut StackGraph,
+        file: Handle<File>,
+        source_root: &Path,
+        source_path: &Path,
+        source: &'b str,
+        lcs: FileLanguageConfigurations<'b>,
+        cancellation_flag: &dyn CancellationFlag,
+    ) -> std::result::Result<(), BuildErrorWithSource<'b>> {
+        let relative_source_path = source_path.strip_prefix(source_root).unwrap();
+        if let Some(lc) = lcs.primary {
+            let globals = Variables::new();
+            lc.sgl
+                .build_stack_graph_into(graph, file, source, &globals, cancellation_flag)
+                .map_err(|inner| BuildErrorWithSource {
+                    inner,
+                    source_path: source_path.to_path_buf(),
+                    source_str: source,
+                    tsg_path: lc.sgl.tsg_path().to_path_buf(),
+                    tsg_str: &lc.sgl.tsg_source(),
+                })?;
+        }
+        for (_, fa) in lcs.secondary {
+            fa.build_stack_graph_into(
+                graph,
+                file,
+                &relative_source_path,
+                &source,
+                &mut std::iter::empty(),
+                &HashMap::new(),
+                cancellation_flag,
+            )
+            .map_err(|inner| BuildErrorWithSource {
+                inner,
+                source_path: source_path.to_path_buf(),
+                source_str: &source,
+                tsg_path: PathBuf::new(),
+                tsg_str: "",
+            })?;
+        }
+        Ok(())
+    }
+
     /// Determines if a path should be skipped because we have not seen the
     /// continue_from mark yet. If the mark is seen, it is cleared, after which
     /// all paths are accepted.

tree-sitter-stack-graphs/src/cli/match.rs

@@ -41,7 +41,10 @@ pub struct MatchArgs {
 impl MatchArgs {
     pub fn run(self, mut loader: Loader) -> anyhow::Result<()> {
         let mut file_reader = FileReader::new();
-        let lc = match loader.load_for_file(&self.source_path, &mut file_reader, &NoCancellation)? {
+        let lc = match loader
+            .load_for_file(&self.source_path, &mut file_reader, &NoCancellation)?
+            .primary
+        {
             Some(lc) => lc,
             None => return Err(anyhow!("No stack graph language found")),
         };

tree-sitter-stack-graphs/src/cli/test.rs

@@ -33,6 +33,7 @@ use crate::CancellationFlag;
 use crate::NoCancellation;
 
 use super::util::iter_files_and_directories;
+use super::util::BuildErrorWithSource;
 use super::util::ConsoleFileLogger;
 use super::util::FileLogger;
 
@@ -206,8 +207,11 @@ impl TestArgs {
             test_path.to_path_buf()
         };
         let mut file_reader = MappingFileReader::new(&load_path, test_path);
-        let lc = match loader.load_for_file(&load_path, &mut file_reader, cancellation_flag)? {
-            Some(sgl) => sgl,
+        let lc = match loader
+            .load_for_file(&load_path, &mut file_reader, cancellation_flag)?
+            .primary
+        {
+            Some(lc) => lc,
             None => return Ok(TestResult::new()),
         };
 
@@ -236,7 +240,7 @@ impl TestArgs {
             let result = if let Some(fa) = test_fragment
                 .path
                 .file_name()
-                .and_then(|f| lc.special_files.get(&f.to_string_lossy()))
+                .and_then(|file_name| lc.special_files.get(&file_name.to_string_lossy()))
             {
                 let mut all_paths = test.fragments.iter().map(|f| f.path.as_path());
                 fa.build_stack_graph_into(
@@ -248,19 +252,34 @@ impl TestArgs {
                     &test_fragment.globals,
                     cancellation_flag,
                 )
+                .map_err(|inner| BuildErrorWithSource {
+                    inner,
+                    source_path: test_path.to_path_buf(),
+                    source_str: &test_fragment.source,
+                    tsg_path: PathBuf::new(),
+                    tsg_str: "",
+                })
             } else if lc.matches_file(
                 &test_fragment.path,
                 &mut Some(test_fragment.source.as_ref()),
             )? {
                 globals.clear();
                 test_fragment.add_globals_to(&mut globals);
-                lc.sgl.build_stack_graph_into(
-                    &mut test.graph,
-                    test_fragment.file,
-                    &test_fragment.source,
-                    &globals,
-                    cancellation_flag,
-                )
+                lc.sgl
+                    .build_stack_graph_into(
+                        &mut test.graph,
+                        test_fragment.file,
+                        &test_fragment.source,
+                        &globals,
+                        cancellation_flag,
+                    )
+                    .map_err(|inner| BuildErrorWithSource {
+                        inner,
+                        source_path: test_path.to_path_buf(),
+                        source_str: &test_fragment.source,
+                        tsg_path: lc.sgl.tsg_path().to_path_buf(),
+                        tsg_str: &lc.sgl.tsg_source(),
+                    })
             } else {
                 return Err(anyhow!(
                     "Test fragment {} not supported by language of test file {}",
@@ -270,15 +289,7 @@ impl TestArgs {
             };
             match result {
                 Err(err) => {
-                    file_status.failure(
-                        "failed to build stack graph",
-                        Some(&err.display_pretty(
-                            test_path,
-                            source,
-                            lc.sgl.tsg_path(),
-                            lc.sgl.tsg_source(),
-                        )),
-                    );
+                    file_status.failure("failed to build stack graph", Some(&err.display_pretty()));
                     return Err(anyhow!("Failed to build graph for {}", test_path.display()));
                 }
                 Ok(_) => {}

tree-sitter-stack-graphs/src/cli/util.rs

@@ -546,3 +546,35 @@ pub fn wait_for_input() -> anyhow::Result<()> {
     std::io::stdin().read_line(&mut input)?;
     Ok(())
 }
+
+/// Wraps a build error with the relevant sources
+pub struct BuildErrorWithSource<'a> {
+    pub inner: crate::BuildError,
+    pub source_path: PathBuf,
+    pub source_str: &'a str,
+    pub tsg_path: PathBuf,
+    pub tsg_str: &'a str,
+}
+
+impl<'a> BuildErrorWithSource<'a> {
+    pub fn display_pretty(&'a self) -> impl std::fmt::Display + 'a {
+        DisplayBuildErrorPretty(self)
+    }
+}
+
+struct DisplayBuildErrorPretty<'a>(&'a BuildErrorWithSource<'a>);
+
+impl std::fmt::Display for DisplayBuildErrorPretty<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "{}",
+            self.0.inner.display_pretty(
+                &self.0.source_path,
+                self.0.source_str,
+                &self.0.tsg_path,
+                self.0.tsg_str,
+            )
+        )
+    }
+}

tree-sitter-stack-graphs/src/lib.rs

@@ -1148,6 +1148,9 @@ impl<'a> Builder<'a> {
 }
 
 pub trait FileAnalyzer {
+    /// Construct stack graph for the given file. Implementations must assume that nodes
+    /// for the given file may already exist, and make sure to prevent node id conflicts,
+    /// for example by using `StackGraph::new_node_id`.
     fn build_stack_graph_into<'a>(
         &self,
         stack_graph: &mut StackGraph,