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
@@ -304,14 +315,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
@@ -567,15 +578,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
@@ -750,14 +761,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/closures/closure-with-capture-list.output

@@ -68,13 +68,13 @@ top_level source="⟨body⟩"
           modifier: modifier "let" source="let"
           pattern: identifier "f" source="f"
           value:
-            function_expr source="{ [⟨capture_declaration⟩] in ⟨body⟩ }"
+            function_expr source="⟨body⟩⟨capture_declaration⟩"
               capture_declaration:
                 variable_declaration source="⟨modifier⟩ ⟨pattern⟩"
                   modifier: modifier "weak" source="weak"
                   pattern: identifier "self" source="self"
               body:
-                block source="⟨stmt⟩"
+                block source="{ [weak self] in ⟨stmt⟩ }"
                   stmt:
                     call_expr source="⟨callee⟩()"
                       callee:

unified/extractor/tests/corpus/swift/closures/closure-with-explicit-parameters.output

@@ -70,14 +70,14 @@ top_level source="⟨body⟩"
           modifier: modifier "let" source="let"
           pattern: identifier "f" source="f"
           value:
-            function_expr source="{ (⟨parameter⟩) -> ⟨return_type⟩ in ⟨body⟩ }"
+            function_expr source="⟨body⟩⟨parameter⟩⟨return_type⟩"
               parameter:
                 parameter source="⟨pattern⟩: ⟨type⟩"
                   type: identifier "Int" source="Int"
                   pattern: identifier "x" source="x"
               return_type: identifier "Int" source="Int"
               body:
-                block source="⟨stmt⟩"
+                block source="{ (x: Int) -> Int in ⟨stmt⟩ }"
                   stmt:
                     binary_expr source="⟨left⟩ ⟨operator⟩ ⟨right⟩"
                       left: identifier "x" source="x"

unified/extractor/tests/corpus/swift/closures/closure-with-shorthand-parameters.output

@@ -47,9 +47,9 @@ top_level source="⟨body⟩"
           modifier: modifier "let" source="let"
           pattern: identifier "f" source="f"
           value:
-            function_expr source="{ ⟨body⟩ }"
+            function_expr source="⟨body⟩"
               body:
-                block source="⟨stmt⟩"
+                block source="{ ⟨stmt⟩ }"
                   stmt:
                     binary_expr source="⟨left⟩ ⟨operator⟩ ⟨right⟩"
                       left: identifier "$0" source="$0"

unified/extractor/tests/corpus/swift/closures/multi-statement-closure.output

@@ -101,14 +101,14 @@ top_level source="⟨body⟩"
           modifier: modifier "let" source="let"
           pattern: identifier "f" source="f"
           value:
-            function_expr source="{ (⟨parameter⟩) -> ⟨return_type⟩ in\n  ⟨body⟩\n}"
+            function_expr source="⟨body⟩⟨parameter⟩⟨return_type⟩"
               parameter:
                 parameter source="⟨pattern⟩: ⟨type⟩"
                   type: identifier "Int" source="Int"
                   pattern: identifier "x" source="x"
               return_type: identifier "Int" source="Int"
               body:
-                block source="⟨stmt⟩\n  ⟨stmt⟩"
+                block source="{ (x: Int) -> Int in\n  ⟨stmt⟩\n  ⟨stmt⟩\n}"
                   stmt:
                     variable_declaration source="⟨modifier⟩ ⟨pattern⟩ = ⟨value⟩"
                       modifier: modifier "let" source="let"

unified/extractor/tests/corpus/swift/closures/trailing-closure.output

@@ -51,9 +51,9 @@ top_level source="⟨body⟩"
           argument:
             argument source="⟨value⟩"
               value:
-                function_expr source="{ ⟨body⟩ }"
+                function_expr source="⟨body⟩"
                   body:
-                    block source="⟨stmt⟩"
+                    block source="{ ⟨stmt⟩ }"
                       stmt:
                         binary_expr source="⟨left⟩ ⟨operator⟩ ⟨right⟩"
                           left: identifier "$0" source="$0"

unified/extractor/tests/corpus/swift/control-flow/defer-statement.output

@@ -79,10 +79,10 @@ top_level source="⟨body⟩"
   body:
     block source="⟨stmt⟩"
       stmt:
-        function_declaration source="func ⟨name_node⟩() {\n  ⟨body⟩\n}"
+        function_declaration source="func ⟨name_node⟩() ⟨body⟩"
           name_node: identifier "withCleanup" source="withCleanup"
           body:
-            block source="⟨stmt⟩\n  ⟨stmt⟩"
+            block source="{\n  ⟨stmt⟩\n  ⟨stmt⟩\n}"
               stmt:
                 unsupported_node "defer { print(\"cleanup\") }" source="defer { print(\"cleanup\") }"
                 call_expr source="⟨callee⟩(⟨argument⟩)"

unified/extractor/tests/corpus/swift/control-flow/discard-statement.output

@@ -74,9 +74,9 @@ top_level source="⟨body⟩"
             base_type source="⟨type⟩"
               type: unsupported_node "~Copyable" source="~Copyable"
           member:
-            function_declaration source="⟨modifier⟩ func ⟨name_node⟩() {\n    ⟨body⟩\n  }"
+            function_declaration source="⟨modifier⟩ func ⟨name_node⟩() ⟨body⟩"
               modifier: modifier "consuming" source="consuming"
               name_node: identifier "close" source="close"
               body:
-                block source="⟨stmt⟩"
+                block source="{\n    ⟨stmt⟩\n  }"
                   stmt: unsupported_node "discard self" source="discard self"

unified/extractor/tests/corpus/swift/control-flow/fallthrough.output

@@ -86,15 +86,15 @@ top_level source="⟨body⟩"
   body:
     block source="⟨stmt⟩"
       stmt:
-        function_declaration source="func ⟨name_node⟩(⟨parameter⟩) {\n  ⟨body⟩\n}"
+        function_declaration source="func ⟨name_node⟩(⟨parameter⟩) ⟨body⟩"
           name_node: identifier "classify" source="classify"
           parameter:
             parameter source="⟨external_name_node⟩ ⟨pattern⟩: ⟨type⟩"
               external_name_node: identifier "_" source="_"
               type: identifier "Int" source="Int"
               pattern: identifier "x" source="x"
           body:
-            block source="⟨stmt⟩"
+            block source="{\n  ⟨stmt⟩\n}"
               stmt:
                 switch_expr source="switch ⟨value⟩ {\n  ⟨case⟩\n  ⟨case⟩\n  }"
                   value: identifier "x" source="x"

unified/extractor/tests/corpus/swift/expressions/array-type-constructor.output

@@ -158,14 +158,14 @@ top_level source="⟨body⟩"
                   value: int_literal "1" source="1"
                 argument source="⟨value⟩"
                   value:
-                    function_expr source="{ ⟨parameter⟩ ⟨parameter⟩ in\n\t⟨body⟩\n}"
+                    function_expr source="⟨body⟩⟨parameter⟩⟨parameter⟩"
                       parameter:
                         parameter source="⟨pattern⟩,"
                           pattern: identifier "_" source="_"
                         parameter source="⟨pattern⟩"
                           pattern: identifier "count" source="count"
                       body:
-                        block source="⟨stmt⟩"
+                        block source="{ _, count in\n\t⟨stmt⟩\n}"
                           stmt:
                             binary_expr source="⟨left⟩ ⟨operator⟩ ⟨right⟩"
                               left: identifier "count" source="count"