yeast: Add special location field

When specified on a tree template node, this field is used to
override/set the location of the node it is attached to. The field
itself can contain an iterable of nodes, in which case the location will
be the smallest location spanning all of these nodes. The same behaviour
applies if there are multiple `location` fields on the same node.

Author: tausbn

shared/yeast-macros/src/ast.rs

@@ -108,7 +108,7 @@ pub(crate) struct Node {
     /// Optional leaf content such as `"foo"` or `#{name}` in
     /// `(identifier "foo")` or `(identifier #{name})`.
     pub(crate) content: Option<Content>,
-    /// Named fields such as `method: {name}` in `(call method: {name})`.
+    /// Named fields and source-location metadata.
     pub(crate) fields: Vec<Field>,
 }
 
@@ -131,4 +131,9 @@ pub(crate) enum Field {
         /// Whether the field value has a trailing `?`.
         optional: bool,
     },
+    /// Source-location metadata such as `location: {source}`.
+    Location {
+        /// A Rust block yielding the node or nodes whose locations should be combined.
+        block: syn::Block,
+    },
 }

shared/yeast-macros/src/lib.rs

@@ -65,6 +65,8 @@ pub fn query(input: TokenStream) -> TokenStream {
 /// field: {expr}                - extend a named field with `{expr}`'s ids
 /// field: (kind ...)?           - set the field only if every `#{expr}`
 ///                                beneath it has a value (see below)
+/// location: {expr}             - override this node's source range with the
+///                                span of one or more nodes; may be repeated
 /// ```
 ///
 /// # Optional fields

shared/yeast-macros/src/lower.rs

@@ -320,7 +320,7 @@ impl Node {
         let kind_str = self.kind.to_string();
 
         if let Some(content) = &self.content {
-            return match content {
+            let node = match content {
                 Content::Static(lit) => quote! { #ctx.literal(#kind_str, #lit) },
                 Content::Computed(block) => {
                     let block = rust_block(block);
@@ -362,6 +362,7 @@ impl Node {
                     }
                 }
             };
+            return self.with_location(node, &ctx);
         }
 
         let mut stmts = Vec::new();
@@ -372,7 +373,10 @@ impl Node {
                 name,
                 value,
                 optional,
-            } = field;
+            } = field
+            else {
+                continue;
+            };
             let field_name = name.to_string();
             let field_str = field_name
                 .strip_prefix("r#")
@@ -424,13 +428,42 @@ impl Node {
             }
         }
 
-        quote! {
+        let node = quote! {
             {
                 #(#stmts)*
                 let mut __fields: Vec<(&str, Vec<yeast::Id>)> = Vec::new();
                 #(#field_args)*
                 #ctx.node(#kind_str, __fields)
             }
+        };
+        self.with_location(node, &ctx)
+    }
+
+    fn with_location(&self, node: TokenStream, ctx: &Ident) -> TokenStream {
+        let locations = self.fields.iter().filter_map(|field| match field {
+            Field::Named { .. } => None,
+            Field::Location { block } => Some(rust_block(block)),
+        });
+        let locations: Vec<_> = locations.collect();
+        if locations.is_empty() {
+            return node;
+        }
+
+        quote! {
+            {
+                let mut __location_ids: Vec<yeast::Id> = Vec::new();
+                #(yeast::IntoFieldIds::extend_into(#locations, &mut __location_ids);)*
+                let __source_range = __location_ids
+                    .into_iter()
+                    .filter_map(|__source| {
+                        #ctx.ast
+                            .get_node(__source)
+                            .and_then(|__node| __node.source_range())
+                    })
+                    .reduce(yeast::Range::union);
+                let __node: yeast::Id = #node;
+                #ctx.set_node_source_range(__node, __source_range)
+            }
         }
     }
 }

shared/yeast-macros/src/template_parse.rs

@@ -70,13 +70,41 @@ fn parse_node(group: proc_macro2::Group) -> Result<Node> {
     let mut fields = Vec::new();
     while peek_is_field(&mut tokens) {
         let name = expect_ident(&mut tokens, "expected field name")?;
-        if content.is_some() {
+        let is_location = name == "location";
+        if content.is_some() && !is_location {
             return Err(syn::Error::new_spanned(
                 name,
-                "literal nodes cannot have named fields",
+                "literal nodes can only have the special `location` field",
             ));
         }
         expect_punct(&mut tokens, ':', "expected `:` after field name")?;
+        if is_location {
+            if !peek_is_group(&mut tokens, Delimiter::Brace) {
+                let error = match tokens.peek() {
+                    Some(token) => syn::Error::new_spanned(
+                        token.clone(),
+                        "`location` field value must be a Rust block",
+                    ),
+                    None => syn::Error::new(
+                        Span::call_site(),
+                        "`location` field value must be a Rust block",
+                    ),
+                };
+                return Err(error);
+            }
+            let block = parse_block(expect_group(&mut tokens, Delimiter::Brace)?)?;
+            if let Some(TokenTree::Punct(punctuation)) = tokens.peek() {
+                if punctuation.as_char() == '?' {
+                    return Err(syn::Error::new_spanned(
+                        punctuation.clone(),
+                        "`?` is not valid on a `location` field",
+                    ));
+                }
+            }
+            fields.push(Field::Location { block });
+            continue;
+        }
+
         let value = parse_one(&mut tokens)?;
         let optional = if matches!(value, Template::Splice(_)) {
             reject_stray_optional(&mut tokens, true)?;
@@ -240,6 +268,46 @@ mod tests {
         assert_eq!(node.fields.len(), 2);
     }
 
+    #[test]
+    fn parses_location_fields_on_literal_nodes() {
+        let template = parse_template(quote! {
+            (identifier "name" location: {source})
+        })
+        .unwrap();
+
+        let Template::Node(node) = template else {
+            panic!("expected node template");
+        };
+        assert_eq!(node.kind, "identifier");
+        assert!(node.content.is_some());
+        assert!(matches!(node.fields.as_slice(), [Field::Location { .. }]));
+    }
+
+    #[test]
+    fn rejects_non_block_location_fields() {
+        let result = parse_template(quote! {
+            (identifier "name" location: (identifier "source"))
+        });
+        let Err(error) = result else {
+            panic!("expected a non-block location value to be rejected");
+        };
+        assert_eq!(
+            error.to_string(),
+            "`location` field value must be a Rust block"
+        );
+    }
+
+    #[test]
+    fn rejects_optional_location_fields() {
+        let result = parse_template(quote! {
+            (identifier "name" location: {source}?)
+        });
+        let Err(error) = result else {
+            panic!("expected an optional location field to be rejected");
+        };
+        assert_eq!(error.to_string(), "`?` is not valid on a `location` field");
+    }
+
     #[test]
     fn parses_lists_and_skips_empty_nodes() {
         let templates = parse_templates(quote! {
@@ -254,13 +322,16 @@ mod tests {
     }
 
     #[test]
-    fn rejects_fields_on_literal_nodes() {
+    fn rejects_named_fields_on_literal_nodes() {
         let result = parse_template(quote! {
             (identifier "name" child: (identifier "nested"))
         });
         let Err(error) = result else {
             panic!("expected literal node fields to be rejected");
         };
-        assert_eq!(error.to_string(), "literal nodes cannot have named fields");
+        assert_eq!(
+            error.to_string(),
+            "literal nodes can only have the special `location` field"
+        );
     }
 }

shared/yeast/doc/yeast.md

@@ -271,6 +271,32 @@ 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.
 
+The special `location` field overrides the location of the node on which it
+appears. Its value must be a Rust block and uses the same node-splicing rules
+as an ordinary field block, but is metadata and is not added to the output AST:
+
+```rust
+rule!(
+    (wrapper
+        first: (_) @first
+        second: (_) @second
+        child: (_) @child)
+    =>
+    (synthetic_node
+        child: {child}
+        location: {[first, second]})
+)
+```
+
+If a `location` value yields several nodes, their locations are combined into
+the smallest range containing all of them. Multiple `location` fields on the
+same node are combined in the same way. The field can also be used on literal
+leaf nodes:
+
+```rust
+(identifier "synthetic" location: {source_node})
+```
+
 The following macros can be used to explicitly set the location associated
 with a newly-created node. They assign a location only to the root of their
 template; nested nodes still derive their locations normally.

shared/yeast/tests/test.rs

@@ -1708,6 +1708,75 @@ fn test_returned_capture_keeps_its_location() {
     assert_eq!(identifier.byte_range(), 4..7);
 }
 
+#[test]
+fn test_location_field_overrides_derived_location() {
+    let rule: Rule = rule!(
+        (assignment
+            left: (identifier) @left
+            right: (integer) @right)
+        =>
+        (call
+            method: {left}
+            location: {right})
+    );
+
+    let ast = run_and_ast("x = 1", vec![rule]);
+    let call = ast
+        .reachable_node_ids()
+        .into_iter()
+        .filter_map(|id| ast.get_node(id))
+        .find(|node| node.kind_name() == "call")
+        .expect("call exists");
+
+    assert_eq!(call.byte_range(), 4..5);
+}
+
+#[test]
+fn test_multiple_location_fields_combine_iterable_locations() {
+    let rule: Rule = rule!(
+        (assignment
+            left: (identifier) @left
+            right: (integer) @right)
+        =>
+        (call
+            method: {left}
+            arguments: (argument_list
+                location: {[left]}
+                location: {[right, left]}))
+    );
+
+    let ast = run_and_ast("x = 1", 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..5);
+}
+
+#[test]
+fn test_location_field_is_available_on_literal_nodes() {
+    let rule: Rule = rule!(
+        (assignment
+            left: (identifier) @left
+            right: (integer) @right)
+        =>
+        (identifier "synthetic" location: {right})
+    );
+
+    let ast = run_and_ast("x = 1", vec![rule]);
+    let identifier = ast
+        .reachable_node_ids()
+        .into_iter()
+        .filter_map(|id| ast.get_node(id))
+        .find(|node| node.kind_name() == "identifier")
+        .expect("identifier exists");
+
+    assert_eq!(identifier.byte_range(), 4..5);
+}
+
 #[test]
 fn test_ignored_location_field_is_excluded_from_rule_result_location() {
     let rule: Rule = rule!(