[ruby/yarp] Ensure correct location with if/elsif, allow FOCUS env var, fix newlines on __END__

https://github.com/ruby/yarp/commit/9da0bc4452

4 files changed, 49 insertions, 40 deletions

diff --git a/test/yarp/parse_test.rb b/test/yarp/parse_test.rb
index 33eb1475f5..c0f3ecf551 100644
--- a/test/yarp/parse_test.rb
+++ b/test/yarp/parse_test.rb
@@ -34,8 +34,12 @@ class ParseTest < Test::Unit::TestCase
   # running on Ruby 3.2+.
   check_ripper = RUBY_VERSION >= "3.2.0"
 
+  # The FOCUS environment variable allows you to specify one particular fixture
+  # to test, instead of all of them.
   base = File.join(__dir__, "fixtures")
-  Dir["**/*.txt", base: base].each do |relative|
+  relatives = ENV["FOCUS"] ? [ENV["FOCUS"]] : Dir["**/*.txt", base: base]
+
+  relatives.each do |relative|
     # These fail on TruffleRuby due to a difference in Symbol#inspect: :测试 vs :"测试"
     next if RUBY_ENGINE == "truffleruby" and %w[seattlerb/bug202.txt seattlerb/magic_encoding_comment.txt].include?(relative)
 
@@ -91,6 +95,13 @@ class ParseTest < Test::Unit::TestCase
       # Next, assert that the newlines are in the expected places.
       expected_newlines = [0]
       source.b.scan("\n") { expected_newlines << $~.offset(0)[0] + 1 }
+
+      # If there's a __END__, then we should trip out those newlines because we
+      # don't actually scan them during parsing (because we don't need to).
+      if found = result.comments.find { |comment| comment.type == :__END__ }
+        expected_newlines = expected_newlines[...found.location.start_line]
+      end
+
       assert_equal expected_newlines, YARP.const_get(:Debug).newlines(source)
 
       # This file has changed behavior in Ripper in Ruby 3.3, so we skip it if
diff --git a/test/yarp/snapshots/if.txt b/test/yarp/snapshots/if.txt
index c12691abfe..d151d6b547 100644
--- a/test/yarp/snapshots/if.txt
+++ b/test/yarp/snapshots/if.txt
@@ -203,7 +203,7 @@ ProgramNode(0...382)(
          (269...271)
        ),
        nil,
-       IfNode(274...289)(
+       IfNode(274...293)(
          (274...279),
          MatchPredicateNode(280...289)(
            CallNode(280...284)(
@@ -222,7 +222,7 @@ ProgramNode(0...382)(
          ),
          nil,
          nil,
-         nil
+         (290...293)
        ),
        (290...293)
      ),
diff --git a/test/yarp/snapshots/seattlerb/if_elsif.txt b/test/yarp/snapshots/seattlerb/if_elsif.txt
index 6f1edfb62e..51c89a9a32 100644
--- a/test/yarp/snapshots/seattlerb/if_elsif.txt
+++ b/test/yarp/snapshots/seattlerb/if_elsif.txt
@@ -5,7 +5,7 @@ ProgramNode(0...18)(
        (0...2),
        IntegerNode(3...4)(),
        nil,
-       IfNode(6...13)((6...11), IntegerNode(12...13)(), nil, nil, nil),
+       IfNode(6...18)((6...11), IntegerNode(12...13)(), nil, nil, (15...18)),
        (15...18)
      )]
   )
diff --git a/yarp/yarp.c b/yarp/yarp.c
index 8cc1130952..419d3ddada 100644
--- a/yarp/yarp.c
+++ b/yarp/yarp.c
@@ -6337,7 +6337,7 @@ parser_lex(yp_parser_t *parser) {
                         ((parser->current.end - parser->current.start) == 7) &&
                         current_token_starts_line(parser) &&
                         (strncmp(parser->current.start, "__END__", 7) == 0) &&
-                        (*parser->current.end == '\n' || (*parser->current.end == '\r' && parser->current.end[1] == '\n'))
+                        (parser->current.end == parser->end || *parser->current.end == '\n' || (*parser->current.end == '\r' && parser->current.end[1] == '\n'))
                     ) {
                         parser->current.end = parser->end;
                         parser->current.type = YP_TOKEN___END__;
@@ -8674,7 +8674,7 @@ parse_conditional(yp_parser_t *parser, yp_context_t context) {
     }
 
     yp_token_t end_keyword = not_provided(parser);
-    yp_node_t *parent;
+    yp_node_t *parent = NULL;
 
     switch (context) {
         case YP_CONTEXT_IF:
@@ -8684,7 +8684,6 @@ parse_conditional(yp_parser_t *parser, yp_context_t context) {
             parent = (yp_node_t *) yp_unless_node_create(parser, &keyword, predicate, statements);
             break;
         default:
-            parent = NULL;
             assert(false && "unreachable");
             break;
     }
@@ -8730,50 +8729,49 @@ parse_conditional(yp_parser_t *parser, yp_context_t context) {
         switch (context) {
             case YP_CONTEXT_IF:
                 ((yp_if_node_t *) current)->consequent = (yp_node_t *) else_node;
-                // Recurse down if nodes setting the appropriate end location in
-                // all cases.
-                yp_node_t *recursing_node = parent;
-                bool recursing = true;
-
-                while (recursing) {
-                    switch (YP_NODE_TYPE(recursing_node)) {
-                        case YP_NODE_IF_NODE:
-                            yp_if_node_end_keyword_loc_set((yp_if_node_t *) recursing_node, &parser->previous);
-                            recursing_node = ((yp_if_node_t *) recursing_node)->consequent;
-                            break;
-                        case YP_NODE_ELSE_NODE:
-                            yp_else_node_end_keyword_loc_set((yp_else_node_t *) recursing_node, &parser->previous);
-                            recursing = false;
-                            break;
-                        default: {
-                            recursing = false;
-                            break;
-                        }
-                    }
-                }
                 break;
             case YP_CONTEXT_UNLESS:
                 ((yp_unless_node_t *) parent)->consequent = else_node;
-                yp_unless_node_end_keyword_loc_set((yp_unless_node_t *) parent, &parser->previous);
                 break;
             default:
                 assert(false && "unreachable");
                 break;
         }
     } else {
-        expect(parser, YP_TOKEN_KEYWORD_END, "Expected `end` to close `if` statement.");
+        expect(parser, YP_TOKEN_KEYWORD_END, "Expected `end` to close conditional statement.");
+    }
 
-        switch (context) {
-            case YP_CONTEXT_IF:
-                yp_if_node_end_keyword_loc_set((yp_if_node_t *) parent, &parser->previous);
-                break;
-            case YP_CONTEXT_UNLESS:
-                yp_unless_node_end_keyword_loc_set((yp_unless_node_t *) parent, &parser->previous);
-                break;
-            default:
-                assert(false && "unreachable");
-                break;
+    // Set the appropriate end location for all of the nodes in the subtree.
+    switch (context) {
+        case YP_CONTEXT_IF: {
+            yp_node_t *current = parent;
+            bool recursing = true;
+
+            while (recursing) {
+                switch (YP_NODE_TYPE(current)) {
+                    case YP_NODE_IF_NODE:
+                        yp_if_node_end_keyword_loc_set((yp_if_node_t *) current, &parser->previous);
+                        current = ((yp_if_node_t *) current)->consequent;
+                        recursing = current != NULL;
+                        break;
+                    case YP_NODE_ELSE_NODE:
+                        yp_else_node_end_keyword_loc_set((yp_else_node_t *) current, &parser->previous);
+                        recursing = false;
+                        break;
+                    default: {
+                        recursing = false;
+                        break;
+                    }
+                }
+            }
+            break;
         }
+        case YP_CONTEXT_UNLESS:
+            yp_unless_node_end_keyword_loc_set((yp_unless_node_t *) parent, &parser->previous);
+            break;
+        default:
+            assert(false && "unreachable");
+            break;
     }
 
     return parent;