yeast: Add explicit location macros

Adds two new macros, both variants of the existing `tree!` macro.

First, `tree_at!` takes an extra argument, and sets the location of the
root of the constructed tree to be that of the argument in question.

Secondly, `tree_spanning!` does the same construction, but accepts an
iterable of nodes instead. It then makes it so that the location of the
root node of the constructed tree is the smallest span that contains all
of the locations given by the iterable.

Author: tausbn

shared/yeast/doc/yeast.md

@@ -256,6 +256,22 @@ elided keywords or delimiters without assigning the same broad range to every
 synthetic descendant. A transform that simply returns a translated capture
 does not widen that capture to the wrapper's range.
 
+When the desired range belongs to another node, `tree_at!` assigns that range
+to the template's root. Nested nodes still derive their own locations normally:
+
+```rust
+let synthetic = tree_at!(ctx, source, (synthetic_node child: (nested value: {child})));
+```
+
+`tree_spanning!` similarly assigns the union of several node ranges:
+
+```rust
+let synthetic = tree_spanning!(ctx, nodes, (synthetic_node child: {child}));
+```
+
+For literals, `ctx.literal_at_start_of(...)` creates an empty range at another
+node's start.
+
 For reviewing locations, `DumpOptions::show_abridged_source` prints each node's
 source range with every direct child replaced by its field name in Unicode
 angle brackets. This keeps delimiters and other parent-owned syntax visible

shared/yeast/src/build.rs

@@ -157,6 +157,16 @@ impl<'a, C> BuildCtx<'a, C> {
         results
     }
 
+    /// Assign an explicit source range to a newly-built result root.
+    #[doc(hidden)]
+    pub fn set_node_source_range(&mut self, node: Id, source_range: Option<Range>) -> Id {
+        if let Some(source_range) = source_range {
+            self.ast.set_source_range(node, source_range);
+            self.created_nodes.remove(&node);
+        }
+        node
+    }
+
     /// Look up a capture variable, returning its node Id.
     pub fn capture(&self, name: &str) -> Id {
         self.captures
@@ -180,6 +190,18 @@ impl<'a, C> BuildCtx<'a, C> {
         self.ast.source_text(id)
     }
 
+    /// Return the source range of a parsed or synthetic node.
+    fn source_range_of(&self, id: Id) -> Option<Range> {
+        self.ast.get_node(id).and_then(|node| node.source_range())
+    }
+
+    /// Return an empty range between two non-overlapping nodes.
+    pub fn empty_source_range_between(&self, left: Id, right: Id) -> Option<Range> {
+        let left = self.source_range_of(left)?;
+        let right = self.source_range_of(right)?;
+        (left.end_byte <= right.start_byte).then(|| left.empty_at_end())
+    }
+
     /// Create a named AST node with the given kind and fields.
     pub fn node(&mut self, kind: &str, fields: Vec<(&str, Vec<Id>)>) -> Id {
         let kind_id = self
@@ -217,6 +239,12 @@ impl<'a, C> BuildCtx<'a, C> {
         self.create_named_token_with_range(kind, value.to_string(), source_range)
     }
 
+    /// Create a literal with an empty range at another node's start.
+    pub fn literal_at_start_of(&mut self, kind: &'static str, value: &str, source: Id) -> Id {
+        let source_range = self.source_range_of(source).map(Range::empty_at_start);
+        self.literal_with_source_range(kind, value, source_range)
+    }
+
     /// Create a leaf node with an auto-generated unique name.
     pub fn fresh(&mut self, kind: &'static str, name: &str) -> Id {
         let generated = self.fresh.resolve(name);

shared/yeast/src/lib.rs

@@ -18,6 +18,40 @@ mod visitor;
 pub use range::{Point, Range};
 pub use yeast_macros::{query, rule, rules, tree, trees};
 
+/// Build a single AST node whose root uses another node's source range.
+///
+/// Nested nodes in the template are built normally and derive their locations
+/// from their own children.
+#[macro_export]
+macro_rules! tree_at {
+    ($ctx:ident, $source:expr, ($($tree:tt)*)) => {{
+        let __yeast_source: $crate::Id = $source;
+        let __yeast_source_range = $ctx
+            .ast
+            .get_node(__yeast_source)
+            .and_then(|node| node.source_range());
+        let __yeast_node: $crate::Id = $crate::tree!($ctx, ($($tree)*));
+        $ctx.set_node_source_range(__yeast_node, __yeast_source_range)
+    }};
+}
+
+/// Build a single AST node whose root spans a collection of nodes.
+///
+/// Nested nodes in the template are built normally and derive their locations
+/// from their own children.
+#[macro_export]
+macro_rules! tree_spanning {
+    ($ctx:ident, $sources:expr, ($($tree:tt)*)) => {{
+        let __yeast_source_range = ::std::iter::IntoIterator::into_iter($sources)
+            .filter_map(|source: $crate::Id| {
+                $ctx.ast.get_node(source).and_then(|node| node.source_range())
+            })
+            .reduce($crate::Range::union);
+        let __yeast_node: $crate::Id = $crate::tree!($ctx, ($($tree)*));
+        $ctx.set_node_source_range(__yeast_node, __yeast_source_range)
+    }};
+}
+
 use captures::Captures;
 use query::QueryNode;
 
@@ -633,6 +667,18 @@ impl Ast {
         });
     }
 
+    /// Replace a synthetic node's source range.
+    pub(crate) fn set_source_range(&mut self, id: Id, source_range: Range) {
+        let node = self
+            .nodes
+            .get_mut(id.0)
+            .unwrap_or_else(|| panic!("set_source_range: invalid node id {}", id.0));
+        if matches!(node.content, NodeContent::Range(_)) {
+            panic!("set_source_range: cannot modify a parsed node");
+        }
+        node.source_range = Some(source_range);
+    }
+
     /// Register a named node kind, returning its id (idempotent). Lets callers
     /// build an AST in a single pass, registering kinds as nodes are created
     /// rather than pre-populating the schema.

shared/yeast/tests/test.rs

@@ -1799,6 +1799,77 @@ fn test_explicit_recursive_translation_keeps_nested_rule_location() {
     assert_eq!(identifier.byte_range(), 4..7);
 }
 
+/// `tree_at!` assigns the captured node's range only to the template root.
+#[test]
+fn test_tree_at_assigns_capture_range_to_root_only() {
+    let rule: Rule = rule!(
+        (call
+            method: (identifier) @name
+            receiver: (identifier) @recv
+        ) @@source
+        =>
+        call {
+            let arguments = tree_at!(ctx, source, (argument_list argument: (integer "0")));
+            tree!((call method: {name} receiver: {recv} arguments: {arguments}))
+        }
+    );
+
+    let ast = run_and_ast("foo.bar()", vec![rule]);
+    let arguments = ast
+        .reachable_node_ids()
+        .into_iter()
+        .filter_map(|id| ast.get_node(id))
+        .find(|node| node.kind_name() == "argument_list")
+        .expect("argument list exists");
+    assert_eq!(arguments.byte_range(), 0..9);
+
+    let integer = ast
+        .reachable_node_ids()
+        .into_iter()
+        .filter_map(|id| ast.get_node(id))
+        .find(|node| node.kind_name() == "integer")
+        .expect("integer exists");
+    assert_eq!(integer.byte_range(), 0..0);
+}
+
+/// `tree_spanning!` assigns the union of the captured node ranges only to the
+/// template root.
+#[test]
+fn test_tree_spanning_assigns_union_to_root_only() {
+    let rule: Rule = rule!(
+        (call
+            method: (identifier) @name
+            receiver: (identifier) @recv
+        )
+        =>
+        call {
+            let arguments = tree_spanning!(
+                ctx,
+                [recv, name],
+                (argument_list argument: (integer "0"))
+            );
+            tree!((call method: {name} receiver: {recv} arguments: {arguments}))
+        }
+    );
+
+    let ast = run_and_ast("foo.bar()", vec![rule]);
+    let arguments = ast
+        .reachable_node_ids()
+        .into_iter()
+        .filter_map(|id| ast.get_node(id))
+        .find(|node| node.kind_name() == "argument_list")
+        .expect("argument list exists");
+    assert_eq!(arguments.byte_range(), 0..7);
+
+    let integer = ast
+        .reachable_node_ids()
+        .into_iter()
+        .filter_map(|id| ast.get_node(id))
+        .find(|node| node.kind_name() == "integer")
+        .expect("integer exists");
+    assert_eq!(integer.byte_range(), 0..0);
+}
+
 /// Explicit ranges on multiple results must not be widened to the range of the
 /// input node matched by the rule.
 #[test]