[ruby/prism] Add then keyword loc to when nodes

https://github.com/ruby/prism/commit/e1e613df16

lib/prism/translation/parser/compiler.rb

@@ -1607,7 +1607,11 @@ module Prism
           builder.when(
             token(node.keyword_loc),
             visit_all(node.conditions),
-            srange_find(node.conditions.last.location.end_offset, node.statements&.location&.start_offset || (node.conditions.last.location.end_offset + 1), [";", "then"]),
+            if node.then_keyword_loc
+              token(node.then_keyword_loc)
+            else
+              srange_find(node.conditions.last.location.end_offset, node.statements&.location&.start_offset || (node.conditions.last.location.end_offset + 1), [";"])
+            end,
             visit(node.statements)
           )
         end

prism/config.yml

@@ -2957,6 +2957,8 @@ nodes:
         type: location
       - name: conditions
         type: node[]
+      - name: then_keyword_loc
+        type: location?
       - name: statements
         type: node?
         kind: StatementsNode

prism/prism.c

@@ -6322,6 +6322,7 @@ pm_when_node_create(pm_parser_t *parser, const pm_token_t *keyword) {
         },
         .keyword_loc = PM_LOCATION_TOKEN_VALUE(keyword),
         .statements = NULL,
+        .then_keyword_loc = PM_OPTIONAL_LOCATION_NOT_PROVIDED_VALUE,
         .conditions = { 0 }
     };
 
@@ -6337,6 +6338,15 @@ pm_when_node_conditions_append(pm_when_node_t *node, pm_node_t *condition) {
     pm_node_list_append(&node->conditions, condition);
 }
 
+/**
+ * Set the location of the then keyword of a when node.
+ */
+static inline void
+pm_when_node_then_keyword_loc_set(pm_when_node_t *node, const pm_token_t *then_keyword) {
+    node->base.location.end = then_keyword->end;
+    node->then_keyword_loc = PM_LOCATION_TOKEN_VALUE(then_keyword);
+}
+
 /**
  * Set the statements list of a when node.
  */
@@ -15564,9 +15574,12 @@ parse_expression_prefix(pm_parser_t *parser, pm_binding_power_t binding_power, b
                     } while (accept1(parser, PM_TOKEN_COMMA));
 
                     if (accept2(parser, PM_TOKEN_NEWLINE, PM_TOKEN_SEMICOLON)) {
-                        accept1(parser, PM_TOKEN_KEYWORD_THEN);
+                        if (accept1(parser, PM_TOKEN_KEYWORD_THEN)) {
+                            pm_when_node_then_keyword_loc_set(when_node, &parser->previous);
+                        }
                     } else {
                         expect1(parser, PM_TOKEN_KEYWORD_THEN, PM_ERR_EXPECT_WHEN_DELIMITER);
+                        pm_when_node_then_keyword_loc_set(when_node, &parser->previous);
                     }
 
                     if (!match3(parser, PM_TOKEN_KEYWORD_WHEN, PM_TOKEN_KEYWORD_ELSE, PM_TOKEN_KEYWORD_END)) {

test/prism/snapshots/case.txt

@@ -21,6 +21,7 @@
         │   │       │       ├── value_loc: (2,6)-(2,8) = "hi"
         │   │       │       ├── closing_loc: ∅
         │   │       │       └── unescaped: "hi"
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements: ∅
         │   ├── consequent: ∅
         │   ├── case_keyword_loc: (1,0)-(1,4) = "case"
@@ -33,6 +34,7 @@
         │   │   │   ├── keyword_loc: (5,11)-(5,15) = "when"
         │   │   │   ├── conditions: (length: 1)
         │   │   │   │   └── @ TrueNode (location: (5,16)-(5,20))
+        │   │   │   ├── then_keyword_loc: ∅
         │   │   │   └── statements:
         │   │   │       @ StatementsNode (location: (5,22)-(5,30))
         │   │   │       └── body: (length: 1)
@@ -59,6 +61,7 @@
         │   │       ├── keyword_loc: (5,32)-(5,36) = "when"
         │   │       ├── conditions: (length: 1)
         │   │       │   └── @ FalseNode (location: (5,37)-(5,42))
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements:
         │   │           @ StatementsNode (location: (5,44)-(5,53))
         │   │           └── body: (length: 1)
@@ -103,6 +106,7 @@
         │   │       │           ├── arguments: ∅
         │   │       │           ├── closing_loc: ∅
         │   │       │           └── block: ∅
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements: ∅
         │   ├── consequent: ∅
         │   ├── case_keyword_loc: (7,0)-(7,4) = "case"
@@ -125,6 +129,7 @@
         │   │       │       ├── value_loc: (10,6)-(10,8) = "hi"
         │   │       │       ├── closing_loc: ∅
         │   │       │       └── unescaped: "hi"
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements: ∅
         │   ├── consequent:
         │   │   @ ElseNode (location: (11,0)-(13,3))
@@ -161,6 +166,7 @@
         │   │       │   │   └── name: :FooBar
         │   │       │   └── @ ConstantReadNode (location: (15,24)-(15,31))
         │   │       │       └── name: :BazBonk
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements: ∅
         │   ├── consequent: ∅
         │   ├── case_keyword_loc: (15,0)-(15,4) = "case"
@@ -204,6 +210,7 @@
         │   │       │       │           └── block: ∅
         │   │       │       ├── closing_loc: ∅
         │   │       │       └── block: ∅
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements: ∅
         │   ├── consequent: ∅
         │   ├── case_keyword_loc: (17,0)-(17,4) = "case"
@@ -224,6 +231,7 @@
         │   │       │       ├── arguments: ∅
         │   │       │       ├── closing_loc: ∅
         │   │       │       └── block: ∅
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements: ∅
         │   ├── consequent:
         │   │   @ ElseNode (location: (23,0)-(25,3))
@@ -254,6 +262,7 @@
         │   │       │       ├── value_loc: (28,9)-(28,10) = "b"
         │   │       │       ├── closing_loc: ∅
         │   │       │       └── unescaped: "b"
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements: ∅
         │   ├── consequent:
         │   │   @ ElseNode (location: (29,5)-(30,6))
@@ -271,6 +280,7 @@
         │   │       │   └── @ IntegerNode (location: (32,19)-(32,20))
         │   │       │       ├── flags: decimal
         │   │       │       └── value: 1
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements: ∅
         │   ├── consequent: ∅
         │   ├── case_keyword_loc: (32,0)-(32,4) = "case"
@@ -294,6 +304,7 @@
         │   │       │   └── @ IntegerNode (location: (35,5)-(35,6))
         │   │       │       ├── flags: decimal
         │   │       │       └── value: 3
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements: ∅
         │   ├── consequent: ∅
         │   ├── case_keyword_loc: (34,0)-(34,4) = "case"
@@ -317,6 +328,7 @@
         │   │       │   └── @ IntegerNode (location: (38,18)-(38,19))
         │   │       │       ├── flags: decimal
         │   │       │       └── value: 3
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements: ∅
         │   ├── consequent: ∅
         │   ├── case_keyword_loc: (38,0)-(38,4) = "case"

test/prism/snapshots/seattlerb/bug_case_when_regexp.txt

@@ -12,7 +12,7 @@
             │   ├── closing_loc: ∅
             │   └── unescaped: "x"
             ├── conditions: (length: 1)
-            │   └── @ WhenNode (location: (1,9)-(1,17))
+            │   └── @ WhenNode (location: (1,9)-(1,22))
             │       ├── keyword_loc: (1,9)-(1,13) = "when"
             │       ├── conditions: (length: 1)
             │       │   └── @ RegularExpressionNode (location: (1,14)-(1,17))
@@ -21,6 +21,7 @@
             │       │       ├── content_loc: (1,15)-(1,16) = "x"
             │       │       ├── closing_loc: (1,16)-(1,17) = "/"
             │       │       └── unescaped: "x"
+            │       ├── then_keyword_loc: (1,18)-(1,22) = "then"
             │       └── statements: ∅
             ├── consequent: ∅
             ├── case_keyword_loc: (1,0)-(1,4) = "case"

test/prism/snapshots/seattlerb/bug_cond_pct.txt

@@ -15,6 +15,7 @@
             │       │       ├── content_loc: (1,14)-(1,22) = "blahblah"
             │       │       ├── closing_loc: (1,22)-(1,23) = "%"
             │       │       └── unescaped: "blahblah"
+            │       ├── then_keyword_loc: ∅
             │       └── statements: ∅
             ├── consequent: ∅
             ├── case_keyword_loc: (1,0)-(1,4) = "case"

test/prism/snapshots/seattlerb/when_splat.txt

@@ -16,7 +16,7 @@
             │   ├── closing_loc: ∅
             │   └── block: ∅
             ├── conditions: (length: 1)
-            │   └── @ WhenNode (location: (1,8)-(1,15))
+            │   └── @ WhenNode (location: (1,8)-(1,20))
             │       ├── keyword_loc: (1,8)-(1,12) = "when"
             │       ├── conditions: (length: 1)
             │       │   └── @ SplatNode (location: (1,13)-(1,15))
@@ -32,6 +32,7 @@
             │       │           ├── arguments: ∅
             │       │           ├── closing_loc: ∅
             │       │           └── block: ∅
+            │       ├── then_keyword_loc: (1,16)-(1,20) = "then"
             │       └── statements: ∅
             ├── consequent: ∅
             ├── case_keyword_loc: (1,0)-(1,4) = "case"

test/prism/snapshots/unparser/corpus/literal/case.txt

@@ -19,6 +19,7 @@
         │   │   │   │       ├── arguments: ∅
         │   │   │   │       ├── closing_loc: ∅
         │   │   │   │       └── block: ∅
+        │   │   │   ├── then_keyword_loc: ∅
         │   │   │   └── statements:
         │   │   │       @ StatementsNode (location: (3,2)-(3,5))
         │   │   │       └── body: (length: 1)
@@ -45,6 +46,7 @@
         │   │       │       ├── arguments: ∅
         │   │       │       ├── closing_loc: ∅
         │   │       │       └── block: ∅
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements:
         │   │           @ StatementsNode (location: (5,2)-(5,5))
         │   │           └── body: (length: 1)
@@ -87,6 +89,7 @@
         │   │   │   │       ├── arguments: ∅
         │   │   │   │       ├── closing_loc: ∅
         │   │   │   │       └── block: ∅
+        │   │   │   ├── then_keyword_loc: ∅
         │   │   │   └── statements: ∅
         │   │   └── @ WhenNode (location: (9,0)-(10,5))
         │   │       ├── keyword_loc: (9,0)-(9,4) = "when"
@@ -101,6 +104,7 @@
         │   │       │       ├── arguments: ∅
         │   │       │       ├── closing_loc: ∅
         │   │       │       └── block: ∅
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements:
         │   │           @ StatementsNode (location: (10,2)-(10,5))
         │   │           └── body: (length: 1)
@@ -143,6 +147,7 @@
         │   │   │   │       ├── arguments: ∅
         │   │   │   │       ├── closing_loc: ∅
         │   │   │   │       └── block: ∅
+        │   │   │   ├── then_keyword_loc: ∅
         │   │   │   └── statements:
         │   │   │       @ StatementsNode (location: (14,2)-(14,5))
         │   │   │       └── body: (length: 1)
@@ -169,6 +174,7 @@
         │   │       │       ├── arguments: ∅
         │   │       │       ├── closing_loc: ∅
         │   │       │       └── block: ∅
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements:
         │   │           @ StatementsNode (location: (16,2)-(16,5))
         │   │           └── body: (length: 1)
@@ -221,6 +227,7 @@
         │   │       │       ├── arguments: ∅
         │   │       │       ├── closing_loc: ∅
         │   │       │       └── block: ∅
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements:
         │   │           @ StatementsNode (location: (20,2)-(20,8))
         │   │           └── body: (length: 1)
@@ -262,6 +269,7 @@
         │   │       │           ├── arguments: ∅
         │   │       │           ├── closing_loc: ∅
         │   │       │           └── block: ∅
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements:
         │   │           @ StatementsNode (location: (24,2)-(24,8))
         │   │           └── body: (length: 1)
@@ -300,6 +308,7 @@
         │   │       │       ├── arguments: ∅
         │   │       │       ├── closing_loc: ∅
         │   │       │       └── block: ∅
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements:
         │   │           @ StatementsNode (location: (28,2)-(28,5))
         │   │           └── body: (length: 1)
@@ -380,6 +389,7 @@
         │   │       │           │           └── block: ∅
         │   │       │           ├── closing_loc: ∅
         │   │       │           └── block: ∅
+        │   │       ├── then_keyword_loc: ∅
         │   │       └── statements: ∅
         │   ├── consequent: ∅
         │   ├── case_keyword_loc: (32,0)-(32,4) = "case"
@@ -429,6 +439,7 @@
             │       │           │           └── value: 1
             │       │           ├── closing_loc: ∅
             │       │           └── block: ∅
+            │       ├── then_keyword_loc: ∅
             │       └── statements: ∅
             ├── consequent: ∅
             ├── case_keyword_loc: (35,0)-(35,4) = "case"

test/prism/snapshots/unparser/corpus/literal/send.txt

@@ -191,6 +191,7 @@
         │   │   │       │       ├── arguments: ∅
         │   │   │       │       ├── closing_loc: ∅
         │   │   │       │       └── block: ∅
+        │   │   │       ├── then_keyword_loc: ∅
         │   │   │       └── statements: ∅
         │   │   ├── consequent: ∅
         │   │   ├── case_keyword_loc: (16,0)-(16,4) = "case"
@@ -231,6 +232,7 @@
         │   │   │       │       ├── arguments: ∅
         │   │   │       │       ├── closing_loc: ∅
         │   │   │       │       └── block: ∅
+        │   │   │       ├── then_keyword_loc: ∅
         │   │   │       └── statements: ∅
         │   │   ├── consequent: ∅
         │   │   ├── case_keyword_loc: (20,0)-(20,4) = "case"

test/prism/snapshots/whitequark/case_cond.txt

@@ -19,6 +19,7 @@
             │       │       ├── arguments: ∅
             │       │       ├── closing_loc: ∅
             │       │       └── block: ∅
+            │       ├── then_keyword_loc: ∅
             │       └── statements:
             │           @ StatementsNode (location: (1,16)-(1,21))
             │           └── body: (length: 1)

test/prism/snapshots/whitequark/case_cond_else.txt

@@ -19,6 +19,7 @@
             │       │       ├── arguments: ∅
             │       │       ├── closing_loc: ∅
             │       │       └── block: ∅
+            │       ├── then_keyword_loc: ∅
             │       └── statements:
             │           @ StatementsNode (location: (1,16)-(1,21))
             │           └── body: (length: 1)

test/prism/snapshots/whitequark/case_expr.txt

@@ -25,6 +25,7 @@
             │       │       ├── content_loc: (1,16)-(1,19) = "bar"
             │       │       ├── closing_loc: (1,19)-(1,20) = "'"
             │       │       └── unescaped: "bar"
+            │       ├── then_keyword_loc: ∅
             │       └── statements:
             │           @ StatementsNode (location: (1,22)-(1,25))
             │           └── body: (length: 1)

test/prism/snapshots/whitequark/case_expr_else.txt

@@ -25,6 +25,7 @@
             │       │       ├── content_loc: (1,16)-(1,19) = "bar"
             │       │       ├── closing_loc: (1,19)-(1,20) = "'"
             │       │       └── unescaped: "bar"
+            │       ├── then_keyword_loc: ∅
             │       └── statements:
             │           @ StatementsNode (location: (1,22)-(1,25))
             │           └── body: (length: 1)

test/prism/snapshots/whitequark/when_multi.txt

@@ -31,6 +31,7 @@
             │       │       ├── content_loc: (1,23)-(1,26) = "baz"
             │       │       ├── closing_loc: (1,26)-(1,27) = "'"
             │       │       └── unescaped: "baz"
+            │       ├── then_keyword_loc: ∅
             │       └── statements:
             │           @ StatementsNode (location: (1,29)-(1,32))
             │           └── body: (length: 1)

test/prism/snapshots/whitequark/when_splat.txt

@@ -35,6 +35,7 @@
             │   │   │           ├── arguments: ∅
             │   │   │           ├── closing_loc: ∅
             │   │   │           └── block: ∅
+            │   │   ├── then_keyword_loc: ∅
             │   │   └── statements:
             │   │       @ StatementsNode (location: (1,24)-(1,27))
             │   │       └── body: (length: 1)
@@ -64,6 +65,7 @@
             │       │           ├── arguments: ∅
             │       │           ├── closing_loc: ∅
             │       │           └── block: ∅
+            │       ├── then_keyword_loc: ∅
             │       └── statements: ∅
             ├── consequent: ∅
             ├── case_keyword_loc: (1,0)-(1,4) = "case"

test/prism/snapshots/whitequark/when_then.txt

@@ -25,6 +25,7 @@
             │       │       ├── content_loc: (1,16)-(1,19) = "bar"
             │       │       ├── closing_loc: (1,19)-(1,20) = "'"
             │       │       └── unescaped: "bar"
+            │       ├── then_keyword_loc: (1,21)-(1,25) = "then"
             │       └── statements:
             │           @ StatementsNode (location: (1,26)-(1,29))
             │           └── body: (length: 1)