Fix invalid block handling (sorbet#9861)

Author: jesse-shopify

parser/prism/Translator.cc

@@ -82,12 +82,14 @@ class Translator::DesugaredBlockArgument {
     // The literal block, if any, e.g. `{ ... }` or `do ... end`
     ast::ExpressionPtr literalBlockExpr;
 
-    // The expression passed as a block, if any, e.g. the `block` in `a.map(&block)`
+    // The expression passed as a block, if any, e.g. the `block` in `a.map(&block)`,
+    // or a `<fwd-block>` local var read, in the case of `foo(...)`.
     ast::ExpressionPtr blockPassExpr;
 
     // The location of the entire block pass, including the `&`.
     //     a.map(&:foo)
     //           ^^^^^
+    // or a zero-length loc for forwarding like `...`
     core::LocOffsets blockPassLoc;
 
 private:
@@ -3580,36 +3582,78 @@ Translator::translateNumberedParametersNode(pm_numbered_parameters_node *numbere
     return result;
 }
 
+// `parentLoc` is the location of the entire parent node, which we need for the `foo(...)` case
+// (as the location used for the `<fwd-block>` argument).
 Translator::DesugaredBlockArgument Translator::desugarBlock(pm_node_t *block, pm_arguments_node *otherArgs,
                                                             pm_location_t parentLoc) {
-    auto result = DesugaredBlockArgument::none();
+    auto hasFwdArgs = otherArgs != nullptr && PM_NODE_FLAG_P(otherArgs, PM_ARGUMENTS_NODE_FLAGS_CONTAINS_FORWARDING);
 
-    if (block) {
-        if (PM_NODE_TYPE_P(block, PM_BLOCK_NODE)) { // a literal block with `{ ... }` or `do ... end`
-            auto blockNode = down_cast<pm_block_node>(block);
+    // Check if there's a block argument in otherArgs (e.g., &block in `foo(&block) { }`)
+    pm_block_argument_node *blockArgInArgs = nullptr;
+    if (otherArgs != nullptr) {
+        auto args = absl::MakeSpan(otherArgs->arguments.nodes, otherArgs->arguments.size);
+        for (auto *arg : args) {
+            if (PM_NODE_TYPE_P(arg, PM_BLOCK_ARGUMENT_NODE)) {
+                blockArgInArgs = down_cast<pm_block_argument_node>(arg);
+                break;
+            }
+        }
+    }
 
-            result = DesugaredBlockArgument::literalBlock(desugarLiteralBlock(
-                blockNode->body, blockNode->parameters, blockNode->base.location, blockNode->opening_loc));
-        } else {
-            ENFORCE(PM_NODE_TYPE_P(block, PM_BLOCK_ARGUMENT_NODE)); // the `&b` in `a.map(&b)`
+    if (block == nullptr) {
+        // Desugar a call like `foo(...)` so it has a block argument like `foo(..., &<fwd-block>)`.
+        if (hasFwdArgs) {
+            // The local variable uses the full call location, but the Magic node uses zero-length at END
+            auto fullLoc = translateLoc(parentLoc);
+            auto magicLoc = fullLoc.copyEndWithZeroLength();
+            return DesugaredBlockArgument::blockPass(MK::Local(fullLoc, core::Names::fwdBlock()), magicLoc);
+        }
+        return DesugaredBlockArgument::none();
+    }
 
-            auto *bp = down_cast<pm_block_argument_node>(block);
+    if (PM_NODE_TYPE_P(block, PM_BLOCK_NODE)) { // a literal block with `{ ... }` or `do ... end`
+        auto blockNode = down_cast<pm_block_node>(block);
 
-            result = desugarBlockPassArgument(bp);
+        auto literalBlock = desugarLiteralBlock(blockNode->body, blockNode->parameters, blockNode->base.location,
+                                                blockNode->opening_loc);
+
+        // Handle combination of block pass argument AND a literal block.
+        // e.g., `foo(&block) { "literal" }` - both need to be kept.
+        if (blockArgInArgs != nullptr) {
+            auto blockPassResult = desugarBlockPassArgument(blockArgInArgs);
+            if (blockPassResult.hasBlockPass()) {
+                return DesugaredBlockArgument::both(move(literalBlock), move(blockPassResult.blockPassExpr),
+                                                    blockPassResult.blockPassLoc);
+            } else if (blockPassResult.hasLiteralBlock()) {
+                // Handle an error case like `a.map(&:foo) { "literal" }`
+                // The Symbol proc would have been desuraged to a literal block.'
+                // We keep both, moving the Symbol proc to the literal block position.
+                auto symbolProc = move(blockPassResult.literalBlockExpr);
+                auto symbolProcLoc = symbolProc.loc();
+                return DesugaredBlockArgument::both(move(literalBlock), move(symbolProc), symbolProcLoc);
+            } else {
+                unreachable("Expected either a block pass or a literal block, but got neither");
+            }
         }
-    }
 
-    auto hasFwdArgs = otherArgs != nullptr && PM_NODE_FLAG_P(otherArgs, PM_ARGUMENTS_NODE_FLAGS_CONTAINS_FORWARDING);
+        // Handle combination of forwarding args AND a literal block.
+        // e.g., `foo(...) { "literal" }` should pass both <fwd-block> AND the literal block.
+        if (hasFwdArgs) {
+            // The local variable uses the full call location, but the Magic node uses zero-length at END
+            auto fullLoc = translateLoc(parentLoc);
+            auto magicLoc = fullLoc.copyEndWithZeroLength();
+            return DesugaredBlockArgument::both(move(literalBlock), MK::Local(fullLoc, core::Names::fwdBlock()),
+                                                magicLoc);
+        }
 
-    if (hasFwdArgs) { // Desugar a call like `foo(...)` so it has a block argument like `foo(..., &b)`.
-        ENFORCE(!result.exists(), "The parser should have rejected a call with both a block pass "
-                                  "argument and forwarded args (e.g. `foo(&b, ...)`)");
+        return DesugaredBlockArgument::literalBlock(move(literalBlock));
 
-        result = DesugaredBlockArgument::blockPass(MK::Local(translateLoc(parentLoc), core::Names::fwdBlock()),
-                                                   core::LocOffsets::none());
-    }
+    } else {
+        ENFORCE(PM_NODE_TYPE_P(block, PM_BLOCK_ARGUMENT_NODE)); // the `&b` in `a.map(&b)`
+        auto *bp = down_cast<pm_block_argument_node>(block);
 
-    return result;
+        return desugarBlockPassArgument(bp);
+    }
 }
 
 ast::ExpressionPtr Translator::desugarLiteralBlock(pm_node *blockBodyNode, pm_node *blockParameters,
@@ -3761,13 +3805,15 @@ ast::ExpressionPtr Translator::desugarMethodCall(ast::ExpressionPtr receiver, co
     if (block.exists()) {
         // There's a block, so we need to calculate the location of the "send" node, excluding it.
         // Start with message location joined with receiver location
+        // (since receiverNode is null, we use the desugared receiver's location directly)
         auto initialLoc = receiver.loc().join(messageLoc);
         std::tie(sendLoc, blockLoc) = computeMethodCallLoc(initialLoc, receiverNode, prismArgs, closingLoc, block);
     } else {
         // There's no block, so the `sendLoc` and `sendWithBlockLoc` are the same, so we can just skip
         // the finicky logic in `computeMethodCallLoc()`.
         sendLoc = sendWithBlockLoc;
     }
+
     auto sendLoc0 = sendLoc.copyWithZeroLength();
 
     if (methodName == core::Names::squareBrackets() || methodName == core::Names::squareBracketsEq()) {
@@ -3847,6 +3893,8 @@ ast::ExpressionPtr Translator::desugarMethodCall(ast::ExpressionPtr receiver, co
                 continue; // Skip anonymous splats (like `f(*)`), which are handled separately in `PM_CALL_NODE`
             } else if (PM_NODE_TYPE_P(arg, PM_FORWARDING_ARGUMENTS_NODE)) {
                 continue; // Skip forwarded args (like `f(...)`), which are handled separately in `PM_CALL_NODE`
+            } else if (PM_NODE_TYPE_P(arg, PM_BLOCK_ARGUMENT_NODE)) {
+                continue; // Skip block args (like `f(&block)`), which are handled by `desugarBlock`
             }
 
             argExprs.emplace_back(desugar(arg));
@@ -3931,19 +3979,26 @@ ast::ExpressionPtr Translator::desugarMethodCall(ast::ExpressionPtr receiver, co
         magicSendArgs.emplace_back(move(kwargsExpr));
 
         if (block.hasBlockPass()) {
-            // Desugar a call with a splat, and any other expression as a block pass argument.
-            // E.g. `foo(*splat, &block)`
+            // Desugar a call with a splat and a block pass argument.
+            // E.g. `foo(*splat, &block)` or `foo(...) { "literal" }`
 
             auto blockPassLoc = hasFwdArgs ? sendLoc.copyEndWithZeroLength() : block.blockPassLoc;
 
             magicSendArgs.emplace_back(move(block.blockPassExpr));
             numPosArgs++;
 
+            if (block.hasLiteralBlock()) {
+                // Both block pass AND literal block: `foo(...) { "literal" }`
+                magicSendArgs.emplace_back(move(block.literalBlockExpr));
+                flags.hasBlock = true;
+            }
+
             return MK::Send(sendWithBlockLoc, MK::Magic(blockPassLoc), core::Names::callWithSplatAndBlockPass(),
                             messageLoc, numPosArgs, move(magicSendArgs), flags);
         }
 
         if (block.hasLiteralBlock()) {
+            // Just a literal block, no block pass
             magicSendArgs.emplace_back(move(block.literalBlockExpr));
             flags.hasBlock = true;
         }
@@ -3955,46 +4010,56 @@ ast::ExpressionPtr Translator::desugarMethodCall(ast::ExpressionPtr receiver, co
                         numPosArgs, move(magicSendArgs), flags);
     }
 
-    // Grab a copy of the argument count, before we concat in the kwargs key/value pairs. // huh?
-    int numPosArgs = prismArgs.size();
+    // Count args, excluding block arguments which are handled separately
+    int numPosArgs =
+        absl::c_count_if(prismArgs, [](auto *arg) { return !PM_NODE_TYPE_P(arg, PM_BLOCK_ARGUMENT_NODE); });
 
     if (block.hasBlockPass()) {
-        // FIXME: move this comment
-        // Special handling for non-Symbol block pass args, like `a.map(&block)`
-        // Symbol procs like `a.map(:to_s)` are rewritten into literal block arguments,
-        // and handled separately below.
+        // Desugar a call (without splat) with a block pass argument.
+        // E.g. `a.each(&block)` or `foo(&block) { "literal" }`
 
-        // Desugar a call without a splat, and any other expression as a block pass argument.
-        // E.g. `a.each(&block)`
-
-        auto blockPassLoc = block.blockPassLoc;
+        auto blockPassLoc = hasFwdArgs ? sendLoc.copyEndWithZeroLength() : block.blockPassLoc;
 
         ast::Send::ARGS_store magicSendArgs;
         magicSendArgs.reserve(3 + prismArgs.size());
         magicSendArgs.emplace_back(move(receiver));
         magicSendArgs.emplace_back(MK::Symbol(sendLoc0, methodName));
         magicSendArgs.emplace_back(move(block.blockPassExpr));
 
-        numPosArgs += 3;
+        // For block pass with literal block: pos_args = 3 (recv, sym, blockPass)
+        // For block pass without literal block: pos_args = 3 + numPosArgs (includes other args)
+        int magicNumPosArgs = block.hasLiteralBlock() ? 3 : numPosArgs + 3;
+
+        if (block.hasLiteralBlock()) {
+            // This supports the invalid case of having both a block pass AND a literal block
+            magicSendArgs.emplace_back(move(block.literalBlockExpr));
+            flags.hasBlock = true;
+        }
 
         for (auto *arg : prismArgs) {
+            if (PM_NODE_TYPE_P(arg, PM_BLOCK_ARGUMENT_NODE)) {
+                continue; // Skip block args, handled above
+            }
             magicSendArgs.emplace_back(desugar(arg));
         }
 
         if (kwargsHashNode) {
             flattenKwargs(kwargsHashNode, magicSendArgs);
-            ast::desugar::DuplicateHashKeyCheck::checkSendArgs(ctx, numPosArgs, magicSendArgs);
+            ast::desugar::DuplicateHashKeyCheck::checkSendArgs(ctx, magicNumPosArgs, magicSendArgs);
         }
 
         return MK::Send(sendWithBlockLoc, MK::Magic(blockPassLoc), core::Names::callWithBlockPass(), messageLoc,
-                        numPosArgs, move(magicSendArgs), flags);
+                        magicNumPosArgs, move(magicSendArgs), flags);
     }
 
     ast::Send::ARGS_store sendArgs{};
     // TODO: reserve size for kwargs Hash keys and values, if needed.
     // TODO: reserve size for the block, if needed.
     sendArgs.reserve(prismArgs.size());
     for (auto *arg : prismArgs) {
+        if (PM_NODE_TYPE_P(arg, PM_BLOCK_ARGUMENT_NODE)) {
+            continue; // Skip block args, handled by `desugarBlock`
+        }
         sendArgs.emplace_back(desugar(arg));
     }
 

test/BUILD

@@ -290,7 +290,6 @@ pipeline_tests(
             # Tests having to do with tree differences in invalid Ruby code; will address later
             "testdata/parser/error_recovery/assign.rb",
             "testdata/parser/error_recovery/begin_1.rb",
-            "testdata/parser/error_recovery/block_arg_and_block.rb",
             "testdata/parser/error_recovery/block_do_1.rb",
             "testdata/parser/error_recovery/case_1.rb",
             "testdata/parser/error_recovery/case_2.rb",
@@ -327,7 +326,6 @@ pipeline_tests(
             "testdata/parser/error_recovery/eof_7.rb",
             "testdata/parser/error_recovery/eof_9.rb",
             "testdata/parser/error_recovery/forward_args.rb",
-            "testdata/parser/error_recovery/forward_args_with_block.rb",
             "testdata/parser/error_recovery/if_do_1.rb",
             "testdata/parser/error_recovery/if_do_2.rb",
             "testdata/parser/error_recovery/if_indent_1.rb",
@@ -516,6 +514,9 @@ pipeline_tests(
             "testdata/parser/error_recovery/eof_8.rb",
             "testdata/parser/error_recovery/other_missing_end.rb",
             "testdata/rbi/proc.rb",
+
+            # Prism preserves more of the structure than the legacy parser
+            "testdata/parser/error_recovery/block_forwarding_invalid_def.rb",
         ],
     ),
     "PrismPosTests",

test/prism_regression/call_forwarding_permutations.rb

@@ -0,0 +1,102 @@
+# typed: false
+
+# This file tests combinations of block forwarding and literal blocks.
+#
+# Summary of 12 combinations:
+#   Valid:
+#     1. None                 - def f; bar; end
+#     2. &                    - def f(&b); bar(&b); end
+#     3. & (anonymous)        - def f(&); bar(&); end
+#     4. ...                  - def f(...); bar(...); end
+#     5. { }                  - def f; bar { }; end
+#     6. ... + { }            - def f(...); bar(...) { }; end
+#     7. & + { }              - def f(&b); bar(&b) { }; end
+#     8. & (anonymous) + { }  - def f(&); bar(&) { }; end
+#   Error cases (definition - cannot be tested, see below):
+#     9. & + ...              - def f(&b, ...); end
+#    10. & (anonymous) + ...  - def f(&, ...); end
+#    11. & + ... + { }        - def f(&b, ...); bar(...) { }; end
+#    12. & (anonymous) + ... + { } - def f(&, ...); bar(...) { }; end
+
+# Cases 9-12 are excluded from this test.
+#
+# These cases have `&` and `...` together in the definition, which is a parser
+# error. The error recovery is fundamentally different between parsers:
+#   - Whitequark: Aggressive recovery destroys the entire file from error point
+#   - Prism: Graceful recovery keeps method definitions with error flags
+#
+# Since this tests parser error recovery (not desugaring), these cases cannot
+# be meaningfully compared between parsers.
+# Tests for these cases are in test/testdata/parser/error_recovery/block_forwarding_invalid_def.rb.
+
+# ==============================================================================
+# VALID CASES - These should produce identical output
+# ==============================================================================
+
+# Case 1: None
+def case1_none
+  bar
+end
+
+# Case 2: Just & - Explicit block parameter
+def case2_block_param(&block)
+  bar(&block)
+end
+
+# Case 3: Anonymous block parameter
+def case3_anonymous_block_param(&)
+  bar(&)
+end
+
+# Case 4: Just ... - Forwarding only
+def case4_forwarding(...)
+  bar(...)
+end
+
+# Case 5: Just { } - Block literal only
+def case5_block_literal
+  bar { "block body" }
+end
+
+# CASE 6: ... + { } at call site
+# Forwarding includes <fwd-block>, AND there's a literal block.
+# Both should be kept in the output.
+def case6_forwarding_and_literal(...)
+  foo(...) { "literal block" }
+end
+
+# CASE 7: & + { } at call site
+# Block pass argument AND a literal block.
+# Both should be kept in the output.
+def case7_block_pass_and_literal(&block)
+  foo(&block) { "literal block" }
+end
+
+# Case 8: Anonymous block param with literal block
+def case8_anonymous_block_pass_and_literal(&)
+  foo(&) { "literal block" }
+end
+
+# ==============================================================================
+# ERROR CASES - These cannot be tested, see above
+# ==============================================================================
+
+# Case 9: & + ... in definition (named block param)
+# def case9_block_param_and_forwarding(&block, ...)
+#   bar(...)
+# end
+
+# Case 10: & + ... in definition (anonymous block param)
+# def case10_anonymous_block_param_and_forwarding(&, ...)
+#   bar(...)
+# end
+
+# Case 11: & + ... + { } (all three, named block param)
+# def case11_all_three(&block, ...)
+#   foo(...) { "literal block" }
+# end
+
+# Case 12: & + ... + { } (all three, anonymous block param)
+# def case12_all_three_anonymous(&, ...)
+#   foo(...) { "literal block" }
+# end