[ruby/prism] Fix locations derived from arguments.

(https://github.com/ruby/prism/pull/1897)

https://github.com/ruby/prism/commit/00b76ef254

4 files changed, 39 insertions, 45 deletions

diff --git a/prism/prism.c b/prism/prism.c
index 84a28b9345..004bf8ca24 100644
--- a/prism/prism.c
+++ b/prism/prism.c
@@ -797,6 +797,27 @@ typedef struct {
 } pm_arguments_t;
 
 /**
+ * Retrieve the end location of a `pm_arguments_t` object.
+ */
+static inline const uint8_t *
+pm_arguments_end(pm_arguments_t *arguments) {
+    if (arguments->block != NULL) {
+        const uint8_t *end = arguments->block->location.end;
+        if (arguments->closing_loc.start != NULL && arguments->closing_loc.end > end) {
+            end = arguments->closing_loc.end;
+        }
+        return end;
+    }
+    if (arguments->closing_loc.start != NULL) {
+        return arguments->closing_loc.end;
+    }
+    if (arguments->arguments != NULL) {
+        return arguments->arguments->base.location.end;
+    }
+    return arguments->closing_loc.end;
+}
+
+/**
  * Check that we're not about to attempt to attach a brace block to a call that
  * has arguments without parentheses.
  */
@@ -1611,11 +1632,7 @@ pm_call_node_aref_create(pm_parser_t *parser, pm_node_t *receiver, pm_arguments_
     pm_call_node_t *node = pm_call_node_create(parser);
 
     node->base.location.start = receiver->location.start;
-    if (arguments->block != NULL) {
-        node->base.location.end = arguments->block->location.end;
-    } else {
-        node->base.location.end = arguments->closing_loc.end;
-    }
+    node->base.location.end = pm_arguments_end(arguments);
 
     node->receiver = receiver;
     node->message_loc.start = arguments->opening_loc.start;
@@ -1664,15 +1681,11 @@ pm_call_node_call_create(pm_parser_t *parser, pm_node_t *receiver, pm_token_t *o
     pm_call_node_t *node = pm_call_node_create(parser);
 
     node->base.location.start = receiver->location.start;
-    if (arguments->block != NULL) {
-        node->base.location.end = arguments->block->location.end;
-    } else if (arguments->closing_loc.start != NULL) {
-        node->base.location.end = arguments->closing_loc.end;
-    } else if (arguments->arguments != NULL) {
-        node->base.location.end = arguments->arguments->base.location.end;
-    } else {
-        node->base.location.end = message->end;
+    const uint8_t *end = pm_arguments_end(arguments);
+    if (end == NULL) {
+        end = message->end;
     }
+    node->base.location.end = end;
 
     node->receiver = receiver;
     node->call_operator_loc = PM_OPTIONAL_LOCATION_TOKEN_VALUE(operator);
@@ -1699,15 +1712,7 @@ pm_call_node_fcall_create(pm_parser_t *parser, pm_token_t *message, pm_arguments
     pm_call_node_t *node = pm_call_node_create(parser);
 
     node->base.location.start = message->start;
-    if (arguments->block != NULL) {
-        node->base.location.end = arguments->block->location.end;
-    } else if (arguments->closing_loc.start != NULL) {
-        node->base.location.end = arguments->closing_loc.end;
-    } else if (arguments->arguments != NULL) {
-        node->base.location.end = arguments->arguments->base.location.end;
-    } else {
-        node->base.location.end = arguments->closing_loc.end;
-    }
+    node->base.location.end = pm_arguments_end(arguments);
 
     node->message_loc = PM_OPTIONAL_LOCATION_TOKEN_VALUE(message);
     node->opening_loc = arguments->opening_loc;
@@ -1755,11 +1760,7 @@ pm_call_node_shorthand_create(pm_parser_t *parser, pm_node_t *receiver, pm_token
     pm_call_node_t *node = pm_call_node_create(parser);
 
     node->base.location.start = receiver->location.start;
-    if (arguments->block != NULL) {
-        node->base.location.end = arguments->block->location.end;
-    } else {
-        node->base.location.end = arguments->closing_loc.end;
-    }
+    node->base.location.end = pm_arguments_end(arguments);
 
     node->receiver = receiver;
     node->call_operator_loc = PM_OPTIONAL_LOCATION_TOKEN_VALUE(operator);
@@ -5136,16 +5137,9 @@ pm_super_node_create(pm_parser_t *parser, const pm_token_t *keyword, pm_argument
     assert(keyword->type == PM_TOKEN_KEYWORD_SUPER);
     pm_super_node_t *node = PM_ALLOC_NODE(parser, pm_super_node_t);
 
-    const uint8_t *end;
-    if (arguments->block != NULL) {
-        end = arguments->block->location.end;
-    } else if (arguments->closing_loc.start != NULL) {
-        end = arguments->closing_loc.end;
-    } else if (arguments->arguments != NULL) {
-        end = arguments->arguments->base.location.end;
-    } else {
+    const uint8_t *end = pm_arguments_end(arguments);
+    if (end == NULL) {
         assert(false && "unreachable");
-        end = NULL;
     }
 
     *node = (pm_super_node_t) {
diff --git a/test/prism/snapshots/super.txt b/test/prism/snapshots/super.txt
index 96b74ce091..55d7526b37 100644
--- a/test/prism/snapshots/super.txt
+++ b/test/prism/snapshots/super.txt
@@ -1,7 +1,7 @@
-@ ProgramNode (location: (1,0)-(17,20))
+@ ProgramNode (location: (1,0)-(17,21))
 ├── locals: []
 └── statements:
-    @ StatementsNode (location: (1,0)-(17,20))
+    @ StatementsNode (location: (1,0)-(17,21))
     └── body: (length: 9)
         ├── @ ForwardingSuperNode (location: (1,0)-(1,5))
         │   └── block: ∅
@@ -51,7 +51,7 @@
         │       │   ├── closing_loc: ∅
         │       │   └── unescaped: "foo"
         │       └── operator_loc: (9,6)-(9,7) = "&"
-        ├── @ SuperNode (location: (11,0)-(11,11))
+        ├── @ SuperNode (location: (11,0)-(11,12))
         │   ├── keyword_loc: (11,0)-(11,5) = "super"
         │   ├── lparen_loc: (11,5)-(11,6) = "("
         │   ├── arguments: ∅
@@ -94,7 +94,7 @@
         │       ├── body: ∅
         │       ├── opening_loc: (15,15)-(15,16) = "{"
         │       └── closing_loc: (15,16)-(15,17) = "}"
-        └── @ SuperNode (location: (17,0)-(17,20))
+        └── @ SuperNode (location: (17,0)-(17,21))
             ├── keyword_loc: (17,0)-(17,5) = "super"
             ├── lparen_loc: (17,5)-(17,6) = "("
             ├── arguments:
diff --git a/test/prism/snapshots/unparser/corpus/literal/send.txt b/test/prism/snapshots/unparser/corpus/literal/send.txt
index 5a3e7ce6b3..68a83b210a 100644
--- a/test/prism/snapshots/unparser/corpus/literal/send.txt
+++ b/test/prism/snapshots/unparser/corpus/literal/send.txt
@@ -1004,7 +1004,7 @@
         │   ├── block: ∅
         │   ├── flags: ∅
         │   └── name: :foo
-        ├── @ CallNode (location: (58,0)-(58,12))
+        ├── @ CallNode (location: (58,0)-(58,13))
         │   ├── receiver:
         │   │   @ CallNode (location: (58,0)-(58,3))
         │   │   ├── receiver: ∅
@@ -1178,7 +1178,7 @@
         │   ├── block: ∅
         │   ├── flags: ∅
         │   └── name: :bar
-        ├── @ CallNode (location: (62,0)-(62,18))
+        ├── @ CallNode (location: (62,0)-(62,19))
         │   ├── receiver:
         │   │   @ CallNode (location: (62,0)-(62,3))
         │   │   ├── receiver: ∅
@@ -1364,7 +1364,7 @@
         │   ├── block: ∅
         │   ├── flags: ∅
         │   └── name: :bar
-        ├── @ CallNode (location: (66,0)-(66,26))
+        ├── @ CallNode (location: (66,0)-(66,27))
         │   ├── receiver:
         │   │   @ CallNode (location: (66,0)-(66,3))
         │   │   ├── receiver: ∅
diff --git a/test/prism/snapshots/unparser/corpus/literal/super.txt b/test/prism/snapshots/unparser/corpus/literal/super.txt
index 924c016c7e..3beb41916f 100644
--- a/test/prism/snapshots/unparser/corpus/literal/super.txt
+++ b/test/prism/snapshots/unparser/corpus/literal/super.txt
@@ -59,7 +59,7 @@
         │   │   └── flags: ∅
         │   ├── rparen_loc: (4,10)-(4,11) = ")"
         │   └── block: ∅
-        ├── @ SuperNode (location: (5,0)-(5,12))
+        ├── @ SuperNode (location: (5,0)-(5,13))
         │   ├── keyword_loc: (5,0)-(5,5) = "super"
         │   ├── lparen_loc: (5,5)-(5,6) = "("
         │   ├── arguments: ∅
@@ -78,7 +78,7 @@
         │       │   ├── flags: variable_call
         │       │   └── name: :block
         │       └── operator_loc: (5,6)-(5,7) = "&"
-        ├── @ SuperNode (location: (6,0)-(6,15))
+        ├── @ SuperNode (location: (6,0)-(6,16))
         │   ├── keyword_loc: (6,0)-(6,5) = "super"
         │   ├── lparen_loc: (6,5)-(6,6) = "("
         │   ├── arguments: