Merge pull request #22531 from github/copilot/update-ast-transformation

Include elided tokens in yeast AST node locations

Author: tausbn

shared/yeast/src/lib.rs

@@ -591,11 +591,17 @@ impl Ast {
         let source_range = match &content {
             // Parsed nodes already carry an exact source range in their content.
             NodeContent::Range(_) => source_range,
-            // Synthesized nodes derive location from children when possible,
-            // and fall back to the inherited rule-match range otherwise.
+            // Synthesized nodes derive location from both their children and
+            // the inherited rule-match range, so tokens matched by a rule but
+            // elided from its output still contribute to the replacement range.
             _ => self
                 .union_source_range_of_children(&fields)
-                .or(source_range),
+                .map_or(source_range, |child_range| {
+                    Some(match source_range {
+                        Some(source_range) => union_source_ranges(child_range, source_range),
+                        None => child_range,
+                    })
+                }),
         };
         let id = self.nodes.len();
         self.nodes.push(Node {
@@ -786,6 +792,25 @@ impl Ast {
     }
 }
 
+fn union_source_ranges(first: Range, second: Range) -> Range {
+    let (start_byte, start_point) = if first.start_byte <= second.start_byte {
+        (first.start_byte, first.start_point)
+    } else {
+        (second.start_byte, second.start_point)
+    };
+    let (end_byte, end_point) = if first.end_byte >= second.end_byte {
+        (first.end_byte, first.end_point)
+    } else {
+        (second.end_byte, second.end_point)
+    };
+    Range {
+        start_byte,
+        end_byte,
+        start_point,
+        end_point,
+    }
+}
+
 /// A node in our AST
 #[derive(PartialEq, Eq, Debug, Clone, Serialize)]
 pub struct Node {

shared/yeast/tests/test.rs

@@ -1663,6 +1663,39 @@ fn test_hash_brace_uses_capture_location_for_leaf() {
     assert_eq!(bar.end_byte(), 7);
 }
 
+/// Regression test: tokens matched by a rule but elided from the output still
+/// contribute to the source location of the synthesized replacement node.
+#[test]
+fn test_elided_tokens_contribute_to_replacement_location() {
+    let rule: Rule = rule!(
+        (call
+            method: (identifier) @name
+            receiver: (identifier) @recv
+        )
+        =>
+        (call
+            method: {name}
+        )
+    );
+
+    let ast = run_and_ast("foo.bar()", vec![rule]);
+    let call_ids: Vec<yeast::Id> = ast
+        .reachable_node_ids()
+        .into_iter()
+        .filter(|&id| {
+            ast.get_node(id)
+                .is_some_and(|node| node.kind_name() == "call")
+        })
+        .collect();
+
+    assert_eq!(call_ids.len(), 1, "expected exactly one reachable call");
+    let call_id = call_ids[0];
+    let call = ast.get_node(call_id).unwrap();
+
+    assert_eq!(call.start_byte(), 0);
+    assert_eq!(call.end_byte(), 9);
+}
+
 // ---- `rules!` macro tests (compile-time type-checking) ----
 
 /// `rules!` should accept well-typed rules using the bare-rule-body

unified/extractor/src/languages/swift/swift.rs

@@ -115,15 +115,21 @@ fn member_chain(
     ctx: &mut yeast::build::BuildCtx<'_, SwiftContext>,
     parts: Vec<yeast::Id>,
 ) -> yeast::Id {
+    // `member_chain` builds the imported expression inside the larger import
+    // declaration rule. The imported expression should span the import path,
+    // not the whole declaration including the `import` keyword.
+    let source_range = ctx.source_range.take();
     let mut iter = parts.into_iter();
     let first = iter
         .next()
         .expect("identifier with `part:` must have at least one part");
     let init = tree!((name_expr identifier: (identifier #{first})));
-    iter.fold(
+    let result = iter.fold(
         init,
         |acc, elem| tree!((member_access_expr base: {acc} member: (identifier #{elem}))),
-    )
+    );
+    ctx.source_range = source_range;
+    result
 }
 
 /// Compound-assignment operator spellings (`+=`, `<<=`, ...). Used to tell a

unified/extractor/tests/corpus_tests.rs

@@ -98,9 +98,8 @@ fn collect_corpus_stems(dir: &Path, out: &mut Vec<std::path::PathBuf>) {
 
 #[cfg(bazel)]
 fn corpus_dir() -> std::path::PathBuf {
-    let base = std::path::PathBuf::from(
-         std::env::var("RUNFILES_DIR").expect("RUNFILES_DIR not set"),
-    );
+    let base =
+        std::path::PathBuf::from(std::env::var("RUNFILES_DIR").expect("RUNFILES_DIR not set"));
     std::fs::read_dir(&base)
         .expect("failed to read RUNFILES_DIR")
         .filter_map(Result::ok)

unified/extractor/tests/swift_syntax_pipeline.rs

@@ -16,6 +16,39 @@ mod languages;
 /// A real `swift-syntax-rs` JSON dump of the Swift source `let x = 1`.
 const LET_X_JSON: &str = include_str!("fixtures/let_x.swiftsyntax.json");
 
+const IMPORT_FOUNDATION_JSON: &str = r#"{
+    "kind": "sourceFile",
+    "range": {"start":{"offset":0,"line":1,"column":1},"end":{"offset":17,"line":1,"column":18}},
+    "statements": [
+        {
+            "kind": "codeBlockItem",
+            "range": {"start":{"offset":0,"line":1,"column":1},"end":{"offset":17,"line":1,"column":18}},
+            "item": {
+                "kind": "importDecl",
+                "range": {"start":{"offset":0,"line":1,"column":1},"end":{"offset":17,"line":1,"column":18}},
+                "importKeyword": {
+                    "kind": "token",
+                    "tokenKind": "keyword(SwiftSyntax.Keyword.import)",
+                    "text": "import",
+                    "range": {"start":{"offset":0,"line":1,"column":1},"end":{"offset":6,"line":1,"column":7}}
+                },
+                "path": [
+                    {
+                        "kind": "importPathComponent",
+                        "range": {"start":{"offset":7,"line":1,"column":8},"end":{"offset":17,"line":1,"column":18}},
+                        "name": {
+                            "kind": "token",
+                            "tokenKind": "identifier(\"Foundation\")",
+                            "text": "Foundation",
+                            "range": {"start":{"offset":7,"line":1,"column":8},"end":{"offset":17,"line":1,"column":18}}
+                        }
+                    }
+                ]
+            }
+        }
+    ]
+}"#;
+
 #[test]
 fn swift_syntax_json_runs_through_the_desugarer() {
     let lang = languages::all_language_specs()
@@ -47,3 +80,37 @@ fn swift_syntax_json_runs_through_the_desugarer() {
     assert!(dump.contains("top_level"), "unexpected dump: {dump}");
     assert!(dump.contains("block"), "unexpected dump: {dump}");
 }
+
+#[test]
+fn import_name_expr_location_excludes_import_keyword() {
+    let lang = languages::all_language_specs()
+        .into_iter()
+        .find(|l| l.file_globs.iter().any(|g| g.contains("swift")))
+        .expect("swift language spec");
+    let desugarer = lang.desugarer.as_ref();
+    let adapted = languages::swift_adapter::json_to_ast(IMPORT_FOUNDATION_JSON)
+        .expect("adapter should succeed");
+
+    let desugared = desugarer
+        .run_from_ast(adapted.ast)
+        .expect("desugaring an import should not error");
+
+    let name_expr_ids: Vec<yeast::Id> = desugared
+        .reachable_node_ids()
+        .into_iter()
+        .filter(|&id| {
+            desugared
+                .get_node(id)
+                .is_some_and(|node| node.kind_name() == "name_expr")
+        })
+        .collect();
+    assert_eq!(
+        name_expr_ids.len(),
+        1,
+        "expected exactly one reachable name_expr"
+    );
+
+    let name_expr = desugared.get_node(name_expr_ids[0]).unwrap();
+    assert_eq!(name_expr.start_byte(), 7);
+    assert_eq!(name_expr.end_byte(), 17);
+}