unified: Emit locations for comments

tausbn · tausbn · commit ddab6ccb3569 · 2026-07-10T11:36:37.000Z
Also gathers these into a separate side-channel during the initial AST
construction. That way, we don't encounter these as weird extra nodes
while running yeast.
diff --git a/unified/swift-syntax-rs/README.md b/unified/swift-syntax-rs/README.md
@@ -177,15 +177,20 @@ rewrite rules operate on — via [`yeast_adapter::json_to_ast`](src/yeast_adapte
 
 ```rust
 let json = swift_syntax_rs::parse_to_json("let x = 1")?;
-let ast = swift_syntax_rs::yeast_adapter::json_to_ast(&json)?;
+let adapted = swift_syntax_rs::yeast_adapter::json_to_ast(&json)?;
+let ast = adapted.ast; // the yeast::Ast
+let comments = adapted.trivia; // side-channel comment/unexpectedText tokens
 ```
 
 The adapter mirrors tree-sitter's node model, which is what yeast expects:
 layout nodes and varying tokens (identifiers, literals, operators) become
 **named** nodes; fixed tokens (keywords, punctuation) become **anonymous**
-nodes keyed by their text. It preserves swift-syntax's own kind/field names —
-aligning them with the tree-sitter-swift schema so the existing rewrite rules
-fire is a separate, later step.
+nodes keyed by their text. Comments (and `unexpectedText`) are harvested into a
+side channel (`adapted.trivia`) during the same traversal rather than embedded
+in the tree, matching how the extractor treats tree-sitter `extra` nodes. It
+preserves swift-syntax's own kind/field names — aligning them with the
+tree-sitter-swift schema so the existing rewrite rules fire is a separate,
+later step.
 
 ## Layout
 
diff --git a/unified/swift-syntax-rs/src/yeast_adapter.rs b/unified/swift-syntax-rs/src/yeast_adapter.rs
@@ -25,6 +25,29 @@ use serde_json::Value;
 use yeast::schema::Schema;
 use yeast::{Ast, Id, NodeContent, Point, Range};
 
+/// A comment (or `unexpectedText`) recovered from the syntax tree's trivia.
+///
+/// These are collected into a side channel rather than embedded in the
+/// [`yeast::Ast`], mirroring how the extractor treats tree-sitter `extra`
+/// nodes: they carry a location and text but are not attached to a parent.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct TriviaToken {
+    /// The trivia kind (e.g. `lineComment`, `blockComment`, `docLineComment`,
+    /// `docBlockComment`, `unexpectedText`).
+    pub kind: String,
+    /// The verbatim source text of the piece (e.g. `// comment`).
+    pub text: String,
+    /// The source range the piece occupies.
+    pub range: Range,
+}
+
+/// The result of adapting a swift-syntax JSON tree: the [`yeast::Ast`] plus the
+/// comment/`unexpectedText` trivia harvested from it (in source order).
+pub struct AdaptedTree {
+    pub ast: Ast,
+    pub trivia: Vec<TriviaToken>,
+}
+
 /// swift-syntax `TokenKind` cases whose text is *not* determined by the kind
 /// (i.e. `TokenKind.defaultText == nil`). These carry varying information and
 /// are modelled as named leaf nodes; every other token is a fixed
@@ -142,16 +165,19 @@ fn children_of(value: &Value) -> Vec<&Value> {
 ///
 /// This is a single traversal: each node's kind and field names are registered
 /// in the schema on the fly, immediately before the node is created. Children
-/// are built first so a parent's field lists reference existing ids.
-fn build(node: &Value, ast: &mut Ast) -> Result<Id, String> {
+/// are built first so a parent's field lists reference existing ids. Any
+/// comment/`unexpectedText` trivia carried by a token is harvested into
+/// `trivia` during the same pass rather than embedded in the tree.
+fn build(node: &Value, ast: &mut Ast, trivia: &mut Vec<TriviaToken>) -> Result<Id, String> {
     let info = classify(node)?;
+    collect_trivia(node, trivia);
 
     let mut fields: BTreeMap<u16, Vec<Id>> = BTreeMap::new();
     for (field, value) in field_entries(node) {
         let field_id = ast.register_field(field);
         let mut ids = Vec::new();
         for child in children_of(value) {
-            ids.push(build(child, ast)?);
+            ids.push(build(child, ast, trivia)?);
         }
         fields.insert(field_id, ids);
     }
@@ -171,6 +197,35 @@ fn build(node: &Value, ast: &mut Ast) -> Result<Id, String> {
     ))
 }
 
+/// Harvest a token's `leadingTrivia`/`trailingTrivia` pieces (each already
+/// filtered to comments/`unexpectedText` upstream) into `out`. Non-token nodes
+/// have no trivia keys, so this is a no-op for them.
+fn collect_trivia(node: &Value, out: &mut Vec<TriviaToken>) {
+    for key in ["leadingTrivia", "trailingTrivia"] {
+        let Some(Value::Array(pieces)) = node.get(key) else {
+            continue;
+        };
+        for piece in pieces {
+            let (Some(kind), Some(range)) = (
+                piece.get("kind").and_then(Value::as_str),
+                parse_range(piece),
+            ) else {
+                continue;
+            };
+            let text = piece
+                .get("text")
+                .and_then(Value::as_str)
+                .unwrap_or("")
+                .to_string();
+            out.push(TriviaToken {
+                kind: kind.to_string(),
+                text,
+                range,
+            });
+        }
+    }
+}
+
 /// Parse a node's `range` into a [`yeast::Range`].
 ///
 /// The JSON carries, for `start` and `end`, a 0-based UTF-8 file byte `offset`,
@@ -201,14 +256,20 @@ fn parse_range(node: &Value) -> Option<Range> {
 }
 
 /// Convert a swift-syntax JSON tree (as produced by [`crate::parse_to_json`])
-/// into a [`yeast::Ast`].
-pub fn json_to_ast(json: &str) -> Result<Ast, String> {
+/// into a [`yeast::Ast`] plus the comment/`unexpectedText` trivia harvested
+/// from it. Both are produced in a single traversal.
+pub fn json_to_ast(json: &str) -> Result<AdaptedTree, String> {
     let root: Value = serde_json::from_str(json).map_err(|e| format!("invalid JSON: {e}"))?;
 
     let mut ast = Ast::with_schema(Schema::new());
-    let root_id = build(&root, &mut ast)?;
+    let mut trivia = Vec::new();
+    let root_id = build(&root, &mut ast, &mut trivia)?;
     ast.set_root(root_id);
-    Ok(ast)
+
+    // Emit trivia in source order (the traversal visits nodes bottom-up).
+    trivia.sort_by_key(|t| t.range.start_byte);
+
+    Ok(AdaptedTree { ast, trivia })
 }
 
 #[cfg(test)]
@@ -245,7 +306,9 @@ mod tests {
 
     #[test]
     fn builds_ast_from_json() {
-        let ast = json_to_ast(sample_json()).expect("adapter should succeed");
+        let ast = json_to_ast(sample_json())
+            .expect("adapter should succeed")
+            .ast;
         let root = ast.get_root();
         let root_node = ast.get_node(root).expect("root exists");
         assert_eq!(root_node.kind_name(), "sourceFile");
@@ -254,7 +317,9 @@ mod tests {
 
     #[test]
     fn classifies_named_and_anonymous_tokens() {
-        let ast = json_to_ast(sample_json()).expect("adapter should succeed");
+        let ast = json_to_ast(sample_json())
+            .expect("adapter should succeed")
+            .ast;
         // Walk all nodes and collect (kind_name, is_named) for the two tokens.
         let mut let_named = None;
         let mut ident_named = None;
@@ -273,7 +338,9 @@ mod tests {
 
     #[test]
     fn preserves_leaf_text() {
-        let ast = json_to_ast(sample_json()).expect("adapter should succeed");
+        let ast = json_to_ast(sample_json())
+            .expect("adapter should succeed")
+            .ast;
         let ident = ast
             .nodes()
             .iter()
@@ -286,7 +353,9 @@ mod tests {
 
     #[test]
     fn maps_source_locations() {
-        let ast = json_to_ast(sample_json()).expect("adapter should succeed");
+        let ast = json_to_ast(sample_json())
+            .expect("adapter should succeed")
+            .ast;
         let ident = ast
             .nodes()
             .iter()
@@ -300,13 +369,55 @@ mod tests {
         assert_eq!(ident.end_position(), Point::new(0, 5));
     }
 
+    #[test]
+    fn collects_trivia_into_side_channel() {
+        // A token carrying a trailing line comment in its trivia.
+        let json = r#"{
+            "kind": "sourceFile",
+            "range": {"start":{"offset":0,"line":1,"column":1},"end":{"offset":14,"line":1,"column":15}},
+            "value": {
+                "kind": "token",
+                "tokenKind": "integerLiteral(\"1\")",
+                "text": "1",
+                "range": {"start":{"offset":0,"line":1,"column":1},"end":{"offset":1,"line":1,"column":2}},
+                "trailingTrivia": [
+                    {
+                        "kind": "lineComment",
+                        "text": "// c",
+                        "range": {"start":{"offset":2,"line":1,"column":3},"end":{"offset":6,"line":1,"column":7}}
+                    }
+                ]
+            }
+        }"#;
+        let adapted = json_to_ast(json).expect("adapter should succeed");
+
+        // The comment is in the side channel, with its text and location.
+        assert_eq!(adapted.trivia.len(), 1);
+        let comment = &adapted.trivia[0];
+        assert_eq!(comment.kind, "lineComment");
+        assert_eq!(comment.text, "// c");
+        assert_eq!(comment.range.start_byte, 2);
+        assert_eq!(comment.range.end_byte, 6);
+
+        // It is not embedded in the AST as a node.
+        assert!(
+            adapted
+                .ast
+                .nodes()
+                .iter()
+                .all(|n| n.kind_name() != "lineComment"),
+            "comment should not appear as an AST node"
+        );
+    }
+
     /// End-to-end: real Swift source parsed by the shim, then adapted into a
     /// `yeast::Ast`. Requires the Swift toolchain (like the crate's FFI tests).
     #[test]
     fn end_to_end_from_swift_source() {
-        let json = crate::parse_to_json("func f(n: Int) -> Int { return n }")
+        let json = crate::parse_to_json("func f(n: Int) -> Int { return n } // trailing")
             .expect("parsing should succeed");
-        let ast = json_to_ast(&json).expect("adapter should succeed");
+        let adapted = json_to_ast(&json).expect("adapter should succeed");
+        let ast = &adapted.ast;
 
         let root = ast.get_node(ast.get_root()).expect("root exists");
         assert_eq!(root.kind_name(), "sourceFile");
@@ -323,5 +434,15 @@ mod tests {
             kinds.contains(&"func"),
             "expected an anonymous `func` token, got kinds: {kinds:?}"
         );
+
+        // The trailing comment is recovered into the side channel.
+        assert!(
+            adapted
+                .trivia
+                .iter()
+                .any(|t| t.kind == "lineComment" && t.text == "// trailing"),
+            "expected the trailing comment in the trivia side channel, got: {:?}",
+            adapted.trivia
+        );
     }
 }
diff --git a/unified/swift-syntax-rs/swift/Sources/SwiftSyntaxFFI/SwiftSyntaxFFI.swift b/unified/swift-syntax-rs/swift/Sources/SwiftSyntaxFFI/SwiftSyntaxFFI.swift
@@ -39,20 +39,39 @@ private let keptTriviaKinds: Set<String> = [
     "unexpectedText",
 ]
 
-/// Serialize a trivia collection into an array of `{ kind, text }` pieces,
-/// keeping only the kinds in `keptTriviaKinds`.
-private func serializeTrivia(_ trivia: Trivia) -> [Any] {
-    trivia.pieces.compactMap { piece -> [String: Any]? in
+/// Serialize a trivia collection into an array of `{ kind, text, range }`
+/// pieces, keeping only the kinds in `keptTriviaKinds`.
+///
+/// `start` is the absolute position of the first piece (a token's leading
+/// trivia starts at `token.position`; its trailing trivia at
+/// `token.endPositionBeforeTrailingTrivia`). Each piece's range is derived by
+/// accumulating piece lengths, so kept pieces carry an exact source location.
+private func serializeTrivia(
+    _ trivia: Trivia,
+    startingAt start: AbsolutePosition,
+    _ converter: SourceLocationConverter
+) -> [Any] {
+    var result: [Any] = []
+    var offset = start.utf8Offset
+    for piece in trivia.pieces {
+        let length = piece.sourceLength.utf8Length
         // The label of an enum case mirror is the case name (e.g. "spaces",
         // "lineComment"), which gives us a stable kind without an exhaustive
         // switch over every `TriviaPiece` case.
         let kind = Mirror(reflecting: piece).children.first?.label ?? "\(piece)"
-        guard keptTriviaKinds.contains(kind) else { return nil }
-        return [
-            "kind": kind,
-            "text": Trivia(pieces: [piece]).description,
-        ]
+        if keptTriviaKinds.contains(kind) {
+            result.append([
+                "kind": kind,
+                "text": Trivia(pieces: [piece]).description,
+                "range": [
+                    "start": location(AbsolutePosition(utf8Offset: offset), converter),
+                    "end": location(AbsolutePosition(utf8Offset: offset + length), converter),
+                ],
+            ])
+        }
+        offset += length
     }
+    return result
 }
 
 /// Recursively convert a SwiftSyntax node into a JSON-serializable value.
@@ -95,11 +114,17 @@ private func serialize(
             "range": range,
         ]
         // Only emit trivia when present; after filtering, most tokens have none.
-        let leading = serializeTrivia(token.leadingTrivia)
+        // Leading trivia starts at the token's own position; trailing trivia
+        // starts just after the token's content.
+        let leading = serializeTrivia(
+            token.leadingTrivia, startingAt: token.position, converter)
         if !leading.isEmpty {
             result["leadingTrivia"] = leading
         }
-        let trailing = serializeTrivia(token.trailingTrivia)
+        let trailing = serializeTrivia(
+            token.trailingTrivia,
+            startingAt: token.endPositionBeforeTrailingTrivia,
+            converter)
         if !trailing.isEmpty {
             result["trailingTrivia"] = trailing
         }
