feat: add root path and pass file path explicitely

Author: nohehf

languages/tree-sitter-stack-graphs-python/src/stack-graphs.tsg

@@ -13,6 +13,7 @@
 ;; ^^^^^^^^^^^^^^^^
 
 global FILE_PATH
+global ROOT_PATH
 global ROOT_NODE
 global JUMP_TO_SCOPE_NODE
 
@@ -272,7 +273,9 @@ inherit .parent_module
   node grandparent_module_ref_node
   var grandparent_module_ref = grandparent_module_ref_node
 
-  scan FILE_PATH {
+  ; get the file path relative to the root path
+  let rel_path = (replace FILE_PATH ROOT_PATH "")
+  scan rel_path {
     "([^/]+)/"
     {
       node def_dot

languages/tree-sitter-stack-graphs-python/test/root_path.py

@@ -0,0 +1,21 @@
+# ------ path: foo/bar/module.py -----------#
+# ------ global: ROOT_PATH=foo/bar  -----------#
+
+foo = 42
+
+# ------ path: foo/bar/baz/module.py -----------#
+# ------ global: ROOT_PATH=foo/bar  -----------#
+
+bar = "hello"
+
+# ------ path: foo/bar/main.py -------------#
+# ------ global: ROOT_PATH=foo/bar  -----------#
+
+from module import foo
+from baz.module import bar
+
+print(foo)
+#     ^ defined: 4, 14
+
+print(bar)
+#     ^ defined: 9, 15

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

@@ -428,10 +428,19 @@ impl<'a> Indexer<'a> {
         cancellation_flag: &dyn CancellationFlag,
     ) -> std::result::Result<(), BuildErrorWithSource<'b>> {
         let relative_source_path = source_path.strip_prefix(source_root).unwrap();
+        // here the file should also have stripped the source_root from its path
         if let Some(lc) = lcs.primary {
             let globals = Variables::new();
             lc.sgl
-                .build_stack_graph_into(graph, file, source, &globals, cancellation_flag)
+                .build_stack_graph_into(
+                    graph,
+                    file,
+                    source,
+                    source_path,
+                    source_root,
+                    &globals,
+                    cancellation_flag,
+                )
                 .map_err(|inner| BuildErrorWithSource {
                     inner,
                     source_path: source_path.to_path_buf(),

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

@@ -321,6 +321,8 @@ impl TestArgs {
                     &mut test.graph,
                     test_fragment.file,
                     &test_fragment.source,
+                    &test_fragment.path,
+                    &test_fragment.root_path,
                     &globals,
                     cancellation_flag.as_ref(),
                 )

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

@@ -438,6 +438,7 @@ static PRECEDENCE_ATTR: &'static str = "precedence";
 static ROOT_NODE_VAR: &'static str = "ROOT_NODE";
 static JUMP_TO_SCOPE_NODE_VAR: &'static str = "JUMP_TO_SCOPE_NODE";
 static FILE_PATH_VAR: &'static str = "FILE_PATH";
+static ROOT_PATH_VAR: &'static str = "ROOT_PATH";
 
 /// Holds information about how to construct stack graphs for a particular language.
 pub struct StackGraphLanguage {
@@ -557,10 +558,12 @@ impl StackGraphLanguage {
         stack_graph: &'a mut StackGraph,
         file: Handle<File>,
         source: &'a str,
+        source_path: &Path,
+        source_root: &Path,
         globals: &'a Variables<'a>,
         cancellation_flag: &'a dyn CancellationFlag,
     ) -> Result<(), BuildError> {
-        self.builder_into_stack_graph(stack_graph, file, source)
+        self.builder_into_stack_graph(stack_graph, file, source, source_path, source_root)
             .build(globals, cancellation_flag)
     }
 
@@ -573,8 +576,10 @@ impl StackGraphLanguage {
         stack_graph: &'a mut StackGraph,
         file: Handle<File>,
         source: &'a str,
+        source_path: &'a Path,
+        source_root: &'a Path,
     ) -> Builder<'a> {
-        Builder::new(self, stack_graph, file, source)
+        Builder::new(self, stack_graph, file, source, source_path, source_root)
     }
 }
 
@@ -583,6 +588,8 @@ pub struct Builder<'a> {
     stack_graph: &'a mut StackGraph,
     file: Handle<File>,
     source: &'a str,
+    source_path: &'a Path,
+    source_root: &'a Path,
     graph: Graph<'a>,
     remapped_nodes: HashMap<usize, NodeID>,
     injected_node_count: usize,
@@ -595,13 +602,17 @@ impl<'a> Builder<'a> {
         stack_graph: &'a mut StackGraph,
         file: Handle<File>,
         source: &'a str,
+        source_path: &'a Path,
+        source_root: &'a Path,
     ) -> Self {
         let span_calculator = SpanCalculator::new(source);
         Builder {
             sgl,
             stack_graph,
             file,
             source,
+            source_path,
+            source_root,
             graph: Graph::new(),
             remapped_nodes: HashMap::new(),
             injected_node_count: 0,
@@ -635,23 +646,33 @@ impl<'a> Builder<'a> {
         let tree = parse_errors.into_tree();
 
         let mut globals = Variables::nested(globals);
+
         if globals.get(&ROOT_NODE_VAR.into()).is_none() {
             let root_node = self.inject_node(NodeID::root());
             globals
                 .add(ROOT_NODE_VAR.into(), root_node.into())
                 .expect("Failed to set ROOT_NODE");
         }
+
         let jump_to_scope_node = self.inject_node(NodeID::jump_to());
         globals
             .add(JUMP_TO_SCOPE_NODE_VAR.into(), jump_to_scope_node.into())
             .expect("Failed to set JUMP_TO_SCOPE_NODE");
+
         if globals.get(&FILE_PATH_VAR.into()).is_none() {
-            let file_name = self.stack_graph[self.file].to_string();
+            let file_name = self.source_path.to_str().unwrap().to_string();
             globals
                 .add(FILE_PATH_VAR.into(), file_name.into())
                 .expect("Failed to set FILE_PATH");
         }
 
+        if globals.get(&ROOT_PATH_VAR.into()).is_none() {
+            let root_path = self.source_root.to_str().unwrap().to_string();
+            globals
+                .add(ROOT_PATH_VAR.into(), root_path.into())
+                .expect("Failed to set ROOT_PATH");
+        }
+
         let mut config = ExecutionConfig::new(&self.sgl.functions, &globals)
             .lazy(true)
             .debug_attributes(

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

@@ -79,10 +79,14 @@ impl LanguageConfiguration {
                 Loader::load_globals_from_config_str(builtins_config, &mut builtins_globals)?;
             }
             let file = builtins.add_file("<builtins>").unwrap();
+            let builtins_path_var = Path::new("<builtins>");
+            let builtins_root = Path::new("");
             sgl.build_stack_graph_into(
                 &mut builtins,
                 file,
                 builtins_source,
+                builtins_path_var,
+                builtins_root,
                 &builtins_globals,
                 cancellation_flag,
             )
@@ -325,17 +329,28 @@ impl Loader {
         graph: &mut StackGraph,
         cancellation_flag: &dyn CancellationFlag,
     ) -> Result<(), LoadError<'a>> {
-        let file = graph.add_file(&path.to_string_lossy()).unwrap();
+        let file_name = path.to_string_lossy();
+        let file: stack_graphs::arena::Handle<stack_graphs::graph::File> =
+            graph.add_file(&file_name).unwrap();
+        let builtins_root = Path::new("");
         let mut globals = Variables::new();
         Self::load_globals_from_config_str(&config, &mut globals)?;
-        sgl.build_stack_graph_into(graph, file, &source, &globals, cancellation_flag)
-            .map_err(|err| LoadError::Builtins {
-                inner: err,
-                source_path: path.to_path_buf(),
-                source,
-                tsg_path: sgl.tsg_path.to_path_buf(),
-                tsg: sgl.tsg_source.clone(),
-            })?;
+        sgl.build_stack_graph_into(
+            graph,
+            file,
+            &source,
+            path,
+            builtins_root,
+            &globals,
+            cancellation_flag,
+        )
+        .map_err(|err| LoadError::Builtins {
+            inner: err,
+            source_path: path.to_path_buf(),
+            source,
+            tsg_path: sgl.tsg_path.to_path_buf(),
+            tsg: sgl.tsg_source.clone(),
+        })?;
         return Ok(());
     }
 

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

@@ -157,6 +157,7 @@ pub struct Test {
 pub struct TestFragment {
     pub file: Handle<File>,
     pub path: PathBuf,
+    pub root_path: PathBuf,
     pub source: String,
     pub assertions: Vec<Assertion>,
     pub globals: HashMap<String, String>,
@@ -180,6 +181,7 @@ impl Test {
         let mut prev_source = String::new();
         let mut line_files = Vec::new();
         let mut line_count = 0;
+        let default_root_path = PathBuf::from("");
         for (current_line_number, current_line) in
             PositionedSubstring::lines_iter(source).enumerate()
         {
@@ -202,6 +204,7 @@ impl Test {
                     fragments.push(TestFragment {
                         file,
                         path: current_path,
+                        root_path: default_root_path.clone(),
                         source: current_source,
                         assertions: Vec::new(),
                         globals: current_globals,
@@ -254,6 +257,7 @@ impl Test {
             fragments.push(TestFragment {
                 file,
                 path: current_path,
+                root_path: default_root_path.clone(),
                 source: current_source,
                 assertions: Vec::new(),
                 globals: current_globals,

tree-sitter-stack-graphs/tests/it/builder.rs

@@ -5,6 +5,8 @@
 // Please see the LICENSE-APACHE or LICENSE-MIT files in this distribution for license details.
 // ------------------------------------------------------------------------------------------------
 
+use std::path::Path;
+
 use stack_graphs::graph::StackGraph;
 use tree_sitter_graph::Variables;
 use tree_sitter_stack_graphs::NoCancellation;
@@ -22,15 +24,27 @@ fn can_support_preexisting_nodes() {
     "#;
     let python = "pass";
 
+    let file_name = "test.py";
+    let source_path = Path::new(file_name);
+    let source_root = Path::new("");
+
     let mut graph = StackGraph::new();
-    let file = graph.get_or_create_file("test.py");
+    let file = graph.get_or_create_file(file_name);
     let node_id = graph.new_node_id(file);
     let _preexisting_node = graph.add_scope_node(node_id, true).unwrap();
 
     let globals = Variables::new();
     let language = StackGraphLanguage::from_str(tree_sitter_python::language(), tsg).unwrap();
     language
-        .build_stack_graph_into(&mut graph, file, python, &globals, &NoCancellation)
+        .build_stack_graph_into(
+            &mut graph,
+            file,
+            python,
+            source_path,
+            source_root,
+            &globals,
+            &NoCancellation,
+        )
         .expect("Failed to build graph");
 }
 
@@ -45,13 +59,18 @@ fn can_support_injected_nodes() {
     "#;
     let python = "pass";
 
+    let file_name = "test.py";
+    let source_path = Path::new(file_name);
+    let source_root = Path::new("");
+
     let mut graph = StackGraph::new();
-    let file = graph.get_or_create_file("test.py");
+    let file = graph.get_or_create_file(file_name);
     let node_id = graph.new_node_id(file);
     let _preexisting_node = graph.add_scope_node(node_id, true).unwrap();
 
     let language = StackGraphLanguage::from_str(tree_sitter_python::language(), tsg).unwrap();
-    let mut builder = language.builder_into_stack_graph(&mut graph, file, python);
+    let mut builder =
+        language.builder_into_stack_graph(&mut graph, file, python, source_path, source_root);
 
     let mut globals = Variables::new();
     globals

tree-sitter-stack-graphs/tests/it/main.rs

@@ -5,6 +5,8 @@
 // Please see the LICENSE-APACHE or LICENSE-MIT files in this distribution for license details.
 // ------------------------------------------------------------------------------------------------
 
+use std::path::Path;
+
 use stack_graphs::arena::Handle;
 use stack_graphs::graph::File;
 use stack_graphs::graph::StackGraph;
@@ -23,11 +25,22 @@ pub(self) fn build_stack_graph(
     python_source: &str,
     tsg_source: &str,
 ) -> Result<(StackGraph, Handle<File>), BuildError> {
+    let file_name = "test.py";
+    let source_path = Path::new(file_name);
+    let source_root = Path::new("");
     let language =
         StackGraphLanguage::from_str(tree_sitter_python::language(), tsg_source).unwrap();
     let mut graph = StackGraph::new();
-    let file = graph.get_or_create_file("test.py");
+    let file = graph.get_or_create_file(file_name);
     let globals = Variables::new();
-    language.build_stack_graph_into(&mut graph, file, python_source, &globals, &NoCancellation)?;
+    language.build_stack_graph_into(
+        &mut graph,
+        file,
+        python_source,
+        source_path,
+        source_root,
+        &globals,
+        &NoCancellation,
+    )?;
     Ok((graph, file))
 }

tree-sitter-stack-graphs/tests/it/test.rs

@@ -68,12 +68,22 @@ fn build_stack_graph_into(
     graph: &mut StackGraph,
     file: Handle<File>,
     python_source: &str,
+    source_path: &Path,
+    source_root: &Path,
     tsg_source: &str,
     globals: &Variables,
 ) -> Result<(), BuildError> {
     let language =
         StackGraphLanguage::from_str(tree_sitter_python::language(), tsg_source).unwrap();
-    language.build_stack_graph_into(graph, file, python_source, globals, &NoCancellation)?;
+    language.build_stack_graph_into(
+        graph,
+        file,
+        python_source,
+        source_path,
+        source_root,
+        globals,
+        &NoCancellation,
+    )?;
     Ok(())
 }
 
@@ -102,6 +112,8 @@ fn check_test(
             &mut test.graph,
             fragments.file,
             &fragments.source,
+            &fragments.path,
+            &fragments.root_path,
             tsg_source,
             &globals,
         )