diff --git a/tree-sitter-stack-graphs/src/lib.rs b/tree-sitter-stack-graphs/src/lib.rs index e2377f853..e8bdd7492 100644 --- a/tree-sitter-stack-graphs/src/lib.rs +++ b/tree-sitter-stack-graphs/src/lib.rs @@ -160,6 +160,35 @@ //! } //! ``` //! +//! ### Annotating nodes with syntax type information +//! +//! You can annotate any stack graph node with information about its syntax type. To do this, add a `syntax_type` +//! attribute, whose value is a string indicating the syntax type. +//! +//! ``` skip +//! (function_definition name: (identifier) @id) @func { +//! node def +//! ; ... +//! attr (def) syntax_type = "function" +//! } +//! ``` +//! +//! ### Annotating definitions with definiens information +//! +//! You cannot annotate definitions with a definiens, which is the thing the definition covers. For example, for +//! a function definition, the definiens would be the function body. To do this, add a `definiens_node` attribute, +//! whose value is a syntax node that spans the definiens. +//! +//! ``` skip +//! (function_definition name: (identifier) @id body: (_) @body) @func { +//! node def +//! ; ... +//! attr (def) definiens_node = @body +//! } +//! ``` +//! +//! Definiens are optional and setting them to `#null` explicitly is allowed. +//! //! ### Connecting stack graph nodes with edges //! //! To connect two stack graph nodes, use the `edge` statement to add an edge between them: @@ -364,6 +393,7 @@ static SCOPE_TYPE: &'static str = "scope"; // Node attribute names static DEBUG_ATTR_PREFIX: &'static str = "debug_"; +static DEFINIENS_NODE: &'static str = "definiens_node"; static EMPTY_SOURCE_SPAN_ATTR: &'static str = "empty_source_span"; static IS_DEFINITION_ATTR: &'static str = "is_definition"; static IS_ENDPOINT_ATTR: &'static str = "is_endpoint"; @@ -372,13 +402,14 @@ static IS_REFERENCE_ATTR: &'static str = "is_reference"; static SCOPE_ATTR: &'static str = "scope"; static SOURCE_NODE_ATTR: &'static str = "source_node"; static SYMBOL_ATTR: &'static str = "symbol"; +static SYNTAX_TYPE: &'static str = "syntax_type"; static TYPE_ATTR: &'static str = "type"; // Expected attributes per node type static POP_SCOPED_SYMBOL_ATTRS: Lazy> = - Lazy::new(|| HashSet::from([TYPE_ATTR, SYMBOL_ATTR, IS_DEFINITION_ATTR])); + Lazy::new(|| HashSet::from([TYPE_ATTR, SYMBOL_ATTR, IS_DEFINITION_ATTR, DEFINIENS_NODE])); static POP_SYMBOL_ATTRS: Lazy> = - Lazy::new(|| HashSet::from([TYPE_ATTR, SYMBOL_ATTR, IS_DEFINITION_ATTR])); + Lazy::new(|| HashSet::from([TYPE_ATTR, SYMBOL_ATTR, IS_DEFINITION_ATTR, DEFINIENS_NODE])); static PUSH_SCOPED_SYMBOL_ATTRS: Lazy> = Lazy::new(|| HashSet::from([TYPE_ATTR, SYMBOL_ATTR, SCOPE_ATTR, IS_REFERENCE_ATTR])); static PUSH_SYMBOL_ATTRS: Lazy> = @@ -885,7 +916,7 @@ impl<'a> Builder<'a> { NodeType::PushSymbol => self.load_push_symbol(node_ref)?, NodeType::Scope => self.load_scope(node_ref)?, }; - self.load_span(node_ref, handle)?; + self.load_source_info(node_ref, handle)?; self.load_node_debug_info(node_ref, handle)?; } @@ -1004,10 +1035,14 @@ impl<'a> Builder<'a> { let id = self.node_id_for_graph_node(node_ref); let is_definition = self.load_flag(node, IS_DEFINITION_ATTR)?; self.verify_attributes(node, POP_SCOPED_SYMBOL_TYPE, &POP_SCOPED_SYMBOL_ATTRS); - Ok(self + let node_handle = self .stack_graph .add_pop_scoped_symbol_node(id, symbol, is_definition) - .unwrap()) + .unwrap(); + if is_definition { + self.load_definiens_info(node_ref, node_handle)?; + } + Ok(node_handle) } fn load_pop_symbol(&mut self, node_ref: GraphNodeRef) -> Result, BuildError> { @@ -1020,10 +1055,14 @@ impl<'a> Builder<'a> { let id = self.node_id_for_graph_node(node_ref); let is_definition = self.load_flag(node, IS_DEFINITION_ATTR)?; self.verify_attributes(node, POP_SYMBOL_TYPE, &POP_SYMBOL_ATTRS); - Ok(self + let node_handle = self .stack_graph .add_pop_symbol_node(id, symbol, is_definition) - .unwrap()) + .unwrap(); + if is_definition { + self.load_definiens_info(node_ref, node_handle)?; + } + Ok(node_handle) } fn load_push_scoped_symbol( @@ -1091,28 +1130,50 @@ impl<'a> Builder<'a> { } } - fn load_span( + fn load_source_info( &mut self, node_ref: GraphNodeRef, node_handle: Handle, ) -> Result<(), BuildError> { let node = &self.graph[node_ref]; - let source_node = match node.attributes.get(SOURCE_NODE_ATTR) { - Some(source_node) => &self.graph[source_node.as_syntax_node_ref()?], - None => return Ok(()), + if let Some(source_node) = node.attributes.get(SOURCE_NODE_ATTR) { + let source_node = &self.graph[source_node.as_syntax_node_ref()?]; + let mut span = self.span_calculator.for_node(source_node); + if match node.attributes.get(EMPTY_SOURCE_SPAN_ATTR) { + Some(empty_source_span) => empty_source_span.as_boolean()?, + None => false, + } { + span.end = span.start.clone(); + } + let containing_line = &self.source[span.start.containing_line.clone()]; + let containing_line = self.stack_graph.add_string(containing_line); + let source_info = self.stack_graph.source_info_mut(node_handle); + source_info.span = span; + source_info.containing_line = ControlledOption::some(containing_line); }; - let mut span = self.span_calculator.for_node(source_node); - if match node.attributes.get(EMPTY_SOURCE_SPAN_ATTR) { - Some(empty_source_span) => empty_source_span.as_boolean()?, - None => false, - } { - span.end = span.start.clone(); + if let Some(syntax_type) = node.attributes.get(SYNTAX_TYPE) { + let syntax_type = syntax_type.as_str()?; + let interned_string = self.stack_graph.add_string(syntax_type); + let source_info = self.stack_graph.source_info_mut(node_handle); + source_info.syntax_type = interned_string.into(); } - let containing_line = &self.source[span.start.containing_line.clone()]; - let containing_line = self.stack_graph.add_string(containing_line); + Ok(()) + } + + fn load_definiens_info( + &mut self, + node_ref: GraphNodeRef, + node_handle: Handle, + ) -> Result<(), BuildError> { + let node = &self.graph[node_ref]; + let definiens_node = match node.attributes.get(DEFINIENS_NODE) { + Some(Value::Null) => return Ok(()), + Some(definiens_node) => &self.graph[definiens_node.as_syntax_node_ref()?], + None => return Ok(()), + }; + let span = self.span_calculator.for_node(definiens_node); let source_info = self.stack_graph.source_info_mut(node_handle); - source_info.span = span; - source_info.containing_line = ControlledOption::some(containing_line); + source_info.definiens_span = span; Ok(()) } diff --git a/tree-sitter-stack-graphs/tests/it/nodes.rs b/tree-sitter-stack-graphs/tests/it/nodes.rs index 2befc6a95..c36a08744 100644 --- a/tree-sitter-stack-graphs/tests/it/nodes.rs +++ b/tree-sitter-stack-graphs/tests/it/nodes.rs @@ -284,3 +284,76 @@ fn can_calculate_spans() { let trimmed_line = &python[source_info.span.start.trimmed_line.clone()]; assert_eq!(trimmed_line, "a"); } + +#[test] +fn can_set_definiens() { + let tsg = r#" + (function_definition name:(_)@name body:(_)@body) { + node result + attr (result) type = "pop_symbol", symbol = (source-text @name), source_node = @name, is_definition + attr (result) definiens_node = @body + } + "#; + let python = r#" + def foo(): + pass + "#; + + let (graph, file) = build_stack_graph(python, tsg).unwrap(); + let node_handle = graph.nodes_for_file(file).next().unwrap(); + let source_info = graph.source_info(node_handle).unwrap(); + + let actual_span = format!( + "{}:{}-{}:{}", + source_info.definiens_span.start.line, + source_info.definiens_span.start.column.utf8_offset, + source_info.definiens_span.end.line, + source_info.definiens_span.end.column.utf8_offset, + ); + assert_eq!("2:8-2:12", actual_span) +} + +#[test] +fn can_set_null_definiens() { + let tsg = r#" + (function_definition name:(_)@name) { + node result + attr (result) type = "pop_symbol", symbol = (source-text @name), source_node = @name, is_definition + attr (result) definiens_node = #null + } + "#; + let python = r#" + def foo(): + pass + "#; + + let (graph, file) = build_stack_graph(python, tsg).unwrap(); + let node_handle = graph.nodes_for_file(file).next().unwrap(); + let source_info = graph.source_info(node_handle).unwrap(); + assert_eq!(lsp_positions::Span::default(), source_info.definiens_span) +} + +#[test] +fn can_set_syntax_type() { + let tsg = r#" + (function_definition) { + node result + attr (result) syntax_type = "function" + } + "#; + let python = r#" + def foo(): + pass + "#; + + let (graph, file) = build_stack_graph(python, tsg).unwrap(); + let node_handle = graph.nodes_for_file(file).next().unwrap(); + let source_info = graph.source_info(node_handle).unwrap(); + + let syntax_type = source_info + .syntax_type + .into_option() + .map(|s| &graph[s]) + .unwrap_or("MISSING"); + assert_eq!("function", syntax_type) +}