[ruby/yarp] Check predicate of while and until

https://github.com/ruby/yarp/commit/f57a6066b5

2 files changed, 22 insertions, 2 deletions

diff --git a/test/yarp/errors_test.rb b/test/yarp/errors_test.rb
index 3a7396ea14..2b3c8f5c89 100644
--- a/test/yarp/errors_test.rb
+++ b/test/yarp/errors_test.rb
@@ -1340,6 +1340,14 @@ module YARP
       ]
     end
 
+    def test_loop_conditional_is_closed
+      source = "while 0 0; foo; end; until 0 0; foo; end"
+      assert_errors expression(source), source, [
+        ["Expected a predicate expression for the `while` statement", 7..7],
+        ["Expected a predicate expression for the `until` statement", 28..28],
+      ]
+    end
+
     private
 
     def assert_errors(expected, source, errors, compare_ripper: RUBY_ENGINE == "ruby")
diff --git a/yarp/yarp.c b/yarp/yarp.c
index 3272d15d85..715708a4d3 100644
--- a/yarp/yarp.c
+++ b/yarp/yarp.c
@@ -8267,6 +8267,18 @@ expect2(yp_parser_t *parser, yp_token_type_t type1, yp_token_type_t type2, yp_di
     parser->previous.type = YP_TOKEN_MISSING;
 }
 
+// This function is the same as expect2, but it expects one of three token types.
+static void
+expect3(yp_parser_t *parser, yp_token_type_t type1, yp_token_type_t type2, yp_token_type_t type3, yp_diagnostic_id_t diag_id) {
+    if (accept3(parser, type1, type2, type3)) return;
+
+    const uint8_t *location = parser->previous.end;
+    yp_diagnostic_list_append(&parser->error_list, location, location, diag_id);
+
+    parser->previous.start = location;
+    parser->previous.type = YP_TOKEN_MISSING;
+}
+
 static yp_node_t *
 parse_expression(yp_parser_t *parser, yp_binding_power_t binding_power, yp_diagnostic_id_t diag_id);
 
@@ -12804,7 +12816,7 @@ parse_expression_prefix(yp_parser_t *parser, yp_binding_power_t binding_power) {
             yp_node_t *predicate = parse_expression(parser, YP_BINDING_POWER_COMPOSITION, YP_ERR_CONDITIONAL_UNTIL_PREDICATE);
             yp_do_loop_stack_pop(parser);
 
-            accept3(parser, YP_TOKEN_KEYWORD_DO_LOOP, YP_TOKEN_NEWLINE, YP_TOKEN_SEMICOLON);
+            expect3(parser, YP_TOKEN_KEYWORD_DO_LOOP, YP_TOKEN_NEWLINE, YP_TOKEN_SEMICOLON, YP_ERR_CONDITIONAL_UNTIL_PREDICATE);
             yp_statements_node_t *statements = NULL;
 
             if (!accept1(parser, YP_TOKEN_KEYWORD_END)) {
@@ -12825,7 +12837,7 @@ parse_expression_prefix(yp_parser_t *parser, yp_binding_power_t binding_power) {
             yp_node_t *predicate = parse_expression(parser, YP_BINDING_POWER_COMPOSITION, YP_ERR_CONDITIONAL_WHILE_PREDICATE);
             yp_do_loop_stack_pop(parser);
 
-            accept3(parser, YP_TOKEN_KEYWORD_DO_LOOP, YP_TOKEN_NEWLINE, YP_TOKEN_SEMICOLON);
+            expect3(parser, YP_TOKEN_KEYWORD_DO_LOOP, YP_TOKEN_NEWLINE, YP_TOKEN_SEMICOLON, YP_ERR_CONDITIONAL_WHILE_PREDICATE);
             yp_statements_node_t *statements = NULL;
 
             if (!accept1(parser, YP_TOKEN_KEYWORD_END)) {