Unified: Fix locations of various tokens

Handles things like `try!` (which is represented as two separate tokens
-- we explicitly union their ranges) and "let" binding modifiers (where
we reuse the bindingSpecifier, getting its location and string value for
free).

Author: tausbn

shared/yeast/src/build.rs

@@ -33,9 +33,6 @@ pub struct BuildCtx<'a, C: 'a = ()> {
     pub ast: &'a mut Ast,
     pub captures: &'a Captures,
     pub fresh: &'a FreshScope,
-    /// Optional source range explicitly inherited by every synthetic node built
-    /// through this context.
-    pub source_range: Option<Range>,
     /// Source range of the node matched by the current rule.
     ///
     /// The `rule!` macro applies this range to locally-created result roots
@@ -66,26 +63,6 @@ impl<'a, C> BuildCtx<'a, C> {
             ast,
             captures,
             fresh,
-            source_range: None,
-            matched_source_range: None,
-            user_ctx,
-            translator: None,
-            created_nodes: BTreeSet::new(),
-        }
-    }
-
-    pub fn with_source_range(
-        ast: &'a mut Ast,
-        captures: &'a Captures,
-        fresh: &'a FreshScope,
-        source_range: Option<Range>,
-        user_ctx: &'a mut C,
-    ) -> Self {
-        Self {
-            ast,
-            captures,
-            fresh,
-            source_range,
             matched_source_range: None,
             user_ctx,
             translator: None,
@@ -107,7 +84,6 @@ impl<'a, C> BuildCtx<'a, C> {
             ast,
             captures,
             fresh,
-            source_range: None,
             matched_source_range: source_range,
             user_ctx,
             translator: Some(translator),
@@ -153,7 +129,7 @@ impl<'a, C> BuildCtx<'a, C> {
         fields: BTreeMap<FieldId, Vec<Id>>,
         is_named: bool,
     ) -> Id {
-        self.create_node_with_range(kind, content, fields, is_named, self.source_range)
+        self.create_node_with_range(kind, content, fields, is_named, None)
     }
 
     /// Create a named token and record it as constructed by this rule invocation.
@@ -179,7 +155,7 @@ impl<'a, C> BuildCtx<'a, C> {
 
     /// Create a named token using this context's explicit default source range.
     pub fn create_named_token(&mut self, kind: &'static str, content: String) -> Id {
-        self.create_named_token_with_range(kind, content, self.source_range)
+        self.create_named_token_with_range(kind, content, None)
     }
 
     /// Finish the current rule invocation by applying the matched source range
@@ -300,11 +276,7 @@ impl<'a, C> BuildCtx<'a, C> {
         value: &str,
         source_range: Option<Range>,
     ) -> Id {
-        self.create_named_token_with_range(
-            kind,
-            value.to_string(),
-            source_range.or(self.source_range),
-        )
+        self.create_named_token_with_range(kind, value.to_string(), source_range)
     }
 
     /// Create a literal with an empty range at another node's start.
@@ -385,7 +357,6 @@ impl<C: Clone> BuildCtx<'_, C> {
             ast: &mut *self.ast,
             captures: self.captures,
             fresh: self.fresh,
-            source_range: self.source_range,
             matched_source_range: self.matched_source_range,
             user_ctx: &mut child_user_ctx,
             translator: self.translator,

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

@@ -1,5 +1,7 @@
 use codeql_extractor::extractor::desugaring;
-use yeast::{ConcreteDesugarer, DesugaringConfig, PhaseKind, Rule, rule, tree, tree_at};
+use yeast::{
+    ConcreteDesugarer, DesugaringConfig, PhaseKind, Rule, rule, tree, tree_at, tree_spanning,
+};
 
 /// User context propagated from outer rules down to the inner rules that
 /// emit the corresponding output declarations, so that each emitted node
@@ -97,7 +99,9 @@ fn and_chain(
     conds
         .into_iter()
         .reduce(|acc, elem| {
-            tree!((binary_expr operator: (infix_operator "&&") left: {acc} right: {elem}))
+            let operator_range = ctx.empty_source_range_between(acc, elem);
+            let operator = ctx.literal_with_source_range("infix_operator", "&&", operator_range);
+            tree!((binary_expr operator: {operator} left: {acc} right: {elem}))
         })
         .expect("control-flow statement must have at least one condition")
 }
@@ -123,21 +127,15 @@ 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!((identifier #{first}));
-    let result = iter.fold(
+    iter.fold(
         init,
         |acc, elem| tree!((member_access_expr base: {acc} member_name_node: (identifier #{elem}))),
-    );
-    ctx.source_range = source_range;
-    result
+    )
 }
 
 /// Compound-assignment operator spellings (`+=`, `<<=`, ...). Used to tell a
@@ -477,14 +475,24 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
         // `enumCaseDecl` rule below) and are tagged `enum_case`, after any
         // `chained_declaration` tag.
         rule!(
-            (enumCaseElement name: @name parameterClause: (enumCaseParameterClause parameters: _* @params))
-            =>
-            (class_like_declaration
-                modifier: {ctx.outer_modifiers.clone()}
-                modifier: {chained_modifier(&mut ctx)}
-                modifier: (modifier "enum_case")
-                name_node: (identifier #{name})
-                member: (constructor_declaration parameter: {params} body: (block)))
+            (enumCaseElement
+                name: @name
+                parameterClause: (enumCaseParameterClause parameters: _* @params)) @@element
+            =>
+            class_like_declaration {
+                let body = tree!((block));
+                let constructor = tree_at!(
+                    ctx,
+                    element,
+                    (constructor_declaration parameter: {params} body: {body})
+                );
+                tree!((class_like_declaration
+                    modifier: {ctx.outer_modifiers.clone()}
+                    modifier: {chained_modifier(&mut ctx)}
+                    modifier: (modifier "enum_case")
+                    name_node: (identifier #{name})
+                    member: {constructor}))
+            }
         ),
         rule!(
             (enumCaseElement name: @name rawValue: (initializerClause value: @val))
@@ -684,12 +692,17 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
                 label: _? @@lbl
                 expression: (functionCallExpr
                     calledExpression: @constructor
-                    arguments: _* @elements))
+                    arguments: _* @elements) @@call)
             =>
             argument {
+                let value = tree_at!(
+                    ctx,
+                    call,
+                    (call_expr callee: {constructor} argument: {elements})
+                );
                 tree!((argument
                     name_node: (identifier #{lbl})?
-                    value: (call_expr callee: {constructor} argument: {elements})))
+                    value: {value}))
             }
         ),
         rule!(
@@ -862,6 +875,7 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
         // form is matched first.
         rule!(
             (optionalBindingCondition
+                bindingSpecifier: @@spec
                 pattern: (identifierPattern identifier: @name)
                 initializer: (initializerClause value: @val))
             =>
@@ -870,18 +884,20 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
                 pattern: (call_expr
                     callee: (member_access_expr base: (identifier "Optional") member_name_node: (identifier "some"))
                     argument: (argument value: (expr_pattern
-                        modifier: (modifier "let")
+                        modifier: (modifier #{spec})
                         expr: (identifier #{name})))))
         ),
         rule!(
-            (optionalBindingCondition pattern: (identifierPattern identifier: @name))
+            (optionalBindingCondition
+                bindingSpecifier: @@spec
+                pattern: (identifierPattern identifier: @name))
             =>
             (pattern_guard_expr
                 value: (identifier #{name})
                 pattern: (call_expr
                     callee: (member_access_expr base: (identifier "Optional") member_name_node: (identifier "some"))
                     argument: (argument value: (expr_pattern
-                        modifier: (modifier "let")
+                        modifier: (modifier #{spec})
                         expr: (identifier #{name})))))
         ),
         // A single condition in an `if`/`while`/`guard` condition list unwraps to
@@ -965,11 +981,19 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
         }),
         // try/try?/try! expr → unary_expr with operator "try", "try?" or "try!"
         rule!(
-            (tryExpr questionOrExclamationMark: _? @@m expression: @e)
+            (tryExpr
+                tryKeyword: @@keyword
+                questionOrExclamationMark: _? @@m
+                expression: @e)
             =>
             expr {
                 let op = format!("try{}", m.map(|m| ctx.source_text(m)).unwrap_or_default());
-                tree!((unary_expr operator: (prefix_operator #{op}) operand: {e}))
+                let operator = tree_spanning!(
+                    ctx,
+                    std::iter::once(keyword).chain(m),
+                    (prefix_operator #{op})
+                );
+                tree!((unary_expr operator: {operator} operand: {e}))
             }
         ),
         // Do-catch → try_expr
@@ -1003,17 +1027,29 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
         // Catch block without error binding
         rule!((catchClause body: @body) => (catch_clause body: {body})),
         // As expression (type cast) — as?, as!
-        rule!((asExpr expression: @val questionOrExclamationMark: _? @@mark type: @ty) => type_cast_expr {
+        rule!((asExpr expression: @val asKeyword: @@keyword questionOrExclamationMark: _? @@mark type: @ty) => type_cast_expr {
             let op = format!("as{}", mark.map(|m| ctx.source_text(m)).unwrap_or_default());
-            tree!((type_cast_expr expr: {val} operator: (infix_operator #{op}) type: {ty}))
+            let operator = tree_spanning!(
+                ctx,
+                std::iter::once(keyword).chain(mark),
+                (infix_operator #{op})
+            );
+            tree!((type_cast_expr expr: {val} operator: {operator} type: {ty}))
         }),
         // Check expression (`x is T`) → type_test_expr
-        rule!((isExpr expression: @val type: @ty) => (type_test_expr expr: {val} operator: (infix_operator "is") type: {ty})),
+        rule!((isExpr expression: @val isKeyword: @@keyword type: @ty) => (type_test_expr
+            expr: {val}
+            operator: {tree_at!(ctx, keyword, (infix_operator "is"))}
+            type: {ty})),
         // Await expression → unary_expr with operator "await"
-        rule!((awaitExpr expression: @val) => (unary_expr operator: (prefix_operator "await") operand: {val})),
+        rule!((awaitExpr awaitKeyword: @@keyword expression: @val) => (unary_expr
+            operator: {tree_at!(ctx, keyword, (prefix_operator "await"))}
+            operand: {val})),
         // Force-unwrap (`x!`) → postfix unary_expr, via swift-syntax's dedicated
         // `forceUnwrapExpr` node.
-        rule!((forceUnwrapExpr expression: @e) => (unary_expr operator: (postfix_operator "!") operand: {e})),
+        rule!((forceUnwrapExpr expression: @e exclamationMark: @@mark) => (unary_expr
+            operator: {tree_at!(ctx, mark, (postfix_operator "!"))}
+            operand: {e})),
         // ---- Imports ----
         // An import declaration. The dotted path (a list of
         // `importPathComponent`s) becomes a `name_node`/`member_access_expr`
@@ -1028,14 +1064,17 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
                 attributes: _* @attrs
                 modifiers: _* @mods
                 importKindSpecifier: _? @@kind
-                path: (importPathComponent name: @@parts)*)
+                path: (importPathComponent name: @@parts)*) @@decl
             =>
             import_declaration {
                 let last = *parts.last().ok_or("import has no path")?;
                 let pattern = match kind {
-                    None => tree!((named_pattern
-                        name_node: (identifier #{last})
-                        sub_pattern: (bulk_importing_pattern))),
+                    None => {
+                        let bulk = tree_at!(ctx, decl, (bulk_importing_pattern));
+                        tree!((named_pattern
+                            name_node: (identifier #{last})
+                            sub_pattern: {bulk}))
+                    }
                     Some(_) => tree!((identifier #{last})),
                 };
                 tree!((import_declaration

unified/extractor/tests/corpus/swift/desugar/import-with-deeply-nested-path-three-parts.output

@@ -38,4 +38,4 @@ top_level
           pattern:
             named_pattern
               name_node: identifier "URLSession"
-              sub_pattern: bulk_importing_pattern
+              sub_pattern: bulk_importing_pattern "import Foundation.Networking.URLSession"

unified/extractor/tests/corpus/swift/desugar/import-with-dotted-path-two-parts.output

@@ -32,4 +32,4 @@ top_level
           pattern:
             named_pattern
               name_node: identifier "Networking"
-              sub_pattern: bulk_importing_pattern
+              sub_pattern: bulk_importing_pattern "import Foundation.Networking"

unified/extractor/tests/corpus/swift/desugar/simple-import-with-single-name.output

@@ -26,4 +26,4 @@ top_level
           pattern:
             named_pattern
               name_node: identifier "Foundation"
-              sub_pattern: bulk_importing_pattern
+              sub_pattern: bulk_importing_pattern "import Foundation"

unified/extractor/tests/location_tests.rs

@@ -71,13 +71,63 @@ fn generic_type_children_have_local_ranges() {
     assert_has_span(&ast, source, "identifier", Some("Foo"), "Foo");
 }
 
+#[test]
+fn nested_calls_include_their_delimiters() {
+    let source = r#"sink(source("first"), source("second"))"#;
+    let ast = desugar(source);
+
+    assert_has_span(&ast, source, "call_expr", None, r#"source("first")"#);
+    assert_has_span(&ast, source, "call_expr", None, r#"source("second")"#);
+}
+
+#[test]
+fn enum_case_constructors_include_their_parameter_clause() {
+    let source = "enum Result<T> { case success(T) }";
+    let ast = desugar(source);
+
+    assert_has_span(&ast, source, "constructor_declaration", None, "success(T)");
+}
+
+#[test]
+fn synthesized_condition_and_switch_nodes_use_child_ranges() {
+    let source = "if a, b { c }\nswitch x { case a, b: c }";
+    let ast = desugar(source);
+
+    assert_has_span(&ast, source, "binary_expr", None, "a, b");
+    assert_has_empty_span(
+        &ast,
+        "infix_operator",
+        Some("&&"),
+        source.find(',').unwrap(),
+    );
+    assert_has_span(&ast, source, "or_pattern", None, "a, b");
+    assert_has_span(&ast, source, "block", None, "c");
+}
+
 #[test]
 fn declaration_and_operator_tokens_keep_precise_ranges() {
-    let source = "func f() { return x }";
+    let source = "func f() { return x }\nlet y = try? await value! as? T\nlet z = value is T";
     let ast = desugar(source);
 
     assert_has_span(&ast, source, "block", None, "{ return x }");
     assert_has_span(&ast, source, "return_expr", None, "return x");
+    assert_has_span(&ast, source, "prefix_operator", Some("try?"), "try?");
+    assert_has_span(&ast, source, "prefix_operator", Some("await"), "await");
+    assert_has_span(&ast, source, "postfix_operator", Some("!"), "!");
+    assert_has_span(&ast, source, "infix_operator", Some("as?"), "as?");
+    assert_has_span(&ast, source, "infix_operator", Some("is"), "is");
+}
+
+#[test]
+fn synthetic_optional_binding_nodes_anchor_to_binding_keyword() {
+    let source = "if let value = optional {}";
+    let ast = desugar(source);
+    let binding_start = source.find("let").unwrap();
+
+    assert_has_empty_span(&ast, "member_access_expr", None, binding_start);
+    assert_has_empty_span(&ast, "identifier", Some("Optional"), binding_start);
+    assert_has_empty_span(&ast, "identifier", Some("some"), binding_start);
+    assert_has_span(&ast, source, "modifier", Some("let"), "let");
 }
 
 #[test]