Unified: Add locations for blocks

Blocks are slightly awkward, since we destructure them (as codeBlock) in
the query, and then reconstruct them (as `block`) in the rule body,
meaning the location of the block is (by default) assigned to the wrong
place by the heuristic.

To get around this, I added a helper function that updates the location
appropriately, including handling cases where we only optionally match a
block. (Also, in some cases we can fix this by just not destructuring
`codeBlock`s in the first place -- there's already a rule that maps
`codeBlock` to `block`.)

Author: tausbn

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

@@ -1,5 +1,5 @@
 use codeql_extractor::extractor::desugaring;
-use yeast::{ConcreteDesugarer, DesugaringConfig, PhaseKind, Rule, rule, tree};
+use yeast::{ConcreteDesugarer, DesugaringConfig, PhaseKind, Rule, rule, tree, tree_at};
 
 /// User context propagated from outer rules down to the inner rules that
 /// emit the corresponding output declarations, so that each emitted node
@@ -60,6 +60,17 @@ impl SwiftContext {
     }
 }
 
+fn block_with_anchor(
+    ctx: &mut yeast::build::BuildCtx<'_, SwiftContext>,
+    statements: Vec<yeast::Id>,
+    anchor: Option<yeast::Id>,
+) -> yeast::Id {
+    match anchor {
+        Some(anchor) => tree_at!(ctx, anchor, (block stmt: {statements})),
+        None => tree!(ctx, (block stmt: {statements})),
+    }
+}
+
 /// Build a freshly-created `chained_declaration` modifier node if
 /// `ctx.is_chained`, else `None`. Used by inner declaration rules to
 /// emit the chained tag for non-first children of a flattening outer
@@ -301,14 +312,14 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
                 bindings: (patternBinding
                     pattern: (identifierPattern identifier: @@name)
                     typeAnnotation: (typeAnnotation type: @ty)
-                    accessorBlock: (accessorBlock accessors: (codeBlockItem)+ @body)))
+                    accessorBlock: (accessorBlock accessors: (codeBlockItem)+ @body) @@accessor_block))
             =>
             (accessor_declaration
                 modifier: (modifier #{spec})
                 name_node: (identifier #{name})
                 type: {ty}
-                accessor_kind: (accessor_kind "get")
-                body: (block stmt: {body}))
+                accessor_kind: {ctx.literal_at_start_of("accessor_kind", "get", accessor_block)}
+                body: {block_with_anchor(&mut ctx, body, Some(accessor_block))})
         ),
         // A property with an explicit accessor block. swift-syntax makes both
         // shapes plain `accessorDecl`s, so they are told apart by the presence
@@ -564,15 +575,15 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
                 signature: (functionSignature
                     parameterClause: (functionParameterClause parameters: _* @params)
                     returnClause: (returnClause type: @ret)?)
-                body: (codeBlock statements: _* @body))
+                body: (codeBlock) @body)
             =>
             (function_declaration
                 modifier: {mods}
                 name_node: (identifier #{name})
                 type_parameter: {type_params}
                 parameter: {params}
                 return_type: {ret}
-                body: (block stmt: {body}))
+                body: {body})
         ),
         rule!(
             (functionDecl
@@ -747,14 +758,14 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
                     capture: (closureCaptureClause items: _* @captures)?
                     parameterClause: _* @params
                     returnClause: (returnClause type: @ret)?)?
-                statements: _* @body)
+                statements: _* @body) @@closure
             =>
             (function_expr
                 modifier: {attrs}
                 capture_declaration: {captures}
                 parameter: {params}
                 return_type: {ret}
-                body: (block stmt: {body}))
+                body: {block_with_anchor(&mut ctx, body, Some(closure))})
         ),
         // A closure capture (`[weak self]`, `[x]`, `[y = expr]`). The optional
         // ownership specifier (`weak`/`unowned`) becomes a modifier; the
@@ -1277,22 +1288,33 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
                 modifiers: _* @mods
                 signature: (functionSignature
                     parameterClause: (functionParameterClause parameters: _* @params))
-                body: (codeBlock statements: _* @body_stmts)?)
+                body: (codeBlock) @body)
             =>
             (constructor_declaration
                 modifier: {mods}
                 parameter: {params}
-                body: (block stmt: {body_stmts}))
+                body: {body})
+        ),
+        rule!(
+            (initializerDecl
+                modifiers: _* @mods
+                signature: (functionSignature
+                    parameterClause: (functionParameterClause parameters: _* @params)))
+            =>
+            (constructor_declaration
+                modifier: {mods}
+                parameter: {params}
+                body: (block))
         ),
         // Deinit declaration → destructor_declaration. Body statements optional.
         rule!(
             (deinitializerDecl
                 modifiers: _* @mods
-                body: (codeBlock statements: _* @body_stmts))
+                body: (codeBlock) @body)
             =>
             (destructor_declaration
                 modifier: {mods}
-                body: (block stmt: {body_stmts}))
+                body: {body})
         ),
         // Typealias declaration
         rule!(

unified/extractor/tests/corpus/swift/expressions/unsafe-expression.output

@@ -47,5 +47,5 @@ top_level
       stmt:
         function_declaration
           name_node: identifier "doWork"
-          body: block
+          body: block "{}"
         unsupported_node "unsafe doWork()"

unified/extractor/tests/corpus/swift/types/class-function.output

@@ -54,4 +54,4 @@ top_level
             function_declaration
               modifier: modifier "class"
               name_node: identifier "make"
-              body: block
+              body: block "{}"

unified/extractor/tests/corpus/swift/types/constructor-with-parameters.output

@@ -77,4 +77,4 @@ top_level
                   external_name_node: identifier "height"
                   type: identifier "Int"
                   pattern: identifier "h"
-              body: block
+              body: block "{}"

unified/extractor/tests/corpus/swift/types/function-type-with-convention-attribute.output

@@ -68,4 +68,4 @@ top_level
           type: unsupported_node "@convention(c) () -> Void"
           value:
             function_expr
-              body: block
+              body: block "{}"

unified/extractor/tests/corpus/swift/types/function-type-with-sendable-attribute.output

@@ -61,4 +61,4 @@ top_level
           type: unsupported_node "@Sendable () -> Void"
           value:
             function_expr
-              body: block
+              body: block "{}"

unified/extractor/tests/corpus/swift/types/static-function.output

@@ -54,4 +54,4 @@ top_level
             function_declaration
               modifier: modifier "static"
               name_node: identifier "make"
-              body: block
+              body: block "{}"

unified/extractor/tests/location_tests.rs

@@ -71,6 +71,15 @@ fn generic_type_children_have_local_ranges() {
     assert_has_span(&ast, source, "identifier", Some("Foo"), "Foo");
 }
 
+#[test]
+fn declaration_and_operator_tokens_keep_precise_ranges() {
+    let source = "func f() { return x }";
+    let ast = desugar(source);
+
+    assert_has_span(&ast, source, "block", None, "{ return x }");
+    assert_has_span(&ast, source, "return_expr", None, "return x");
+}
+
 #[test]
 fn synthetic_type_and_modifier_nodes_use_empty_scope_start_ranges() {
     let source = "let array: [T]\nlet optional: T?\nenum E { case a, b }";