Allow block arguments in arrow control operators

Arrow control operators have their own entries in the grammar, so they
did not cooperate with BlockArguments. This was just a minor oversight,
so this patch adjusts the grammar to add the desired behavior.

fixes #18050

7 files changed, 46 insertions, 8 deletions

diff --git a/compiler/GHC/Parser.y b/compiler/GHC/Parser.y
index afa8a0e1d8..34d46fd4db 100644
--- a/compiler/GHC/Parser.y
+++ b/compiler/GHC/Parser.y
@@ -2886,11 +2886,11 @@ aexp2   :: { ECP }
         | quasiquote          { ECP $ mkHsSplicePV $1 }
 
         -- arrow notation extension
-        | '(|' aexp2 cmdargs '|)'  {% runECP_P $2 >>= \ $2 ->
-                                      fmap ecpFromCmd $
-                                      ams (sLL $1 $> $ HsCmdArrForm noExtField $2 Prefix
-                                                           Nothing (reverse $3))
-                                          [mu AnnOpenB $1,mu AnnCloseB $4] }
+        | '(|' aexp cmdargs '|)'  {% runECP_P $2 >>= \ $2 ->
+                                     fmap ecpFromCmd $
+                                     ams (sLL $1 $> $ HsCmdArrForm noExtField $2 Prefix
+                                                          Nothing (reverse $3))
+                                         [mu AnnOpenB $1,mu AnnCloseB $4] }
 
 splice_exp :: { LHsExpr GhcPs }
         : splice_untyped { mapLoc (HsSpliceE noExtField) $1 }
@@ -2914,8 +2914,9 @@ cmdargs :: { [LHsCmdTop GhcPs] }
         | {- empty -}                   { [] }
 
 acmd    :: { LHsCmdTop GhcPs }
-        : aexp2                 {% runECP_P $1 >>= \ cmd ->
-                                    return (sL1 cmd $ HsCmdTop noExtField cmd) }
+        : aexp                  {% runECP_P $1 >>= \ cmd ->
+                                   runPV (checkCmdBlockArguments cmd) >>= \ _ ->
+                                   return (sL1 cmd $ HsCmdTop noExtField cmd) }
 
 cvtopbody :: { ([AddAnn],[LHsDecl GhcPs]) }
         :  '{'            cvtopdecls0 '}'      { ([mj AnnOpenC $1
diff --git a/compiler/GHC/Parser/PostProcess.hs b/compiler/GHC/Parser/PostProcess.hs
index 5a1817a1f6..fdc3085e3d 100644
--- a/compiler/GHC/Parser/PostProcess.hs
+++ b/compiler/GHC/Parser/PostProcess.hs
@@ -53,7 +53,7 @@ module GHC.Parser.PostProcess (
         -- Bunch of functions in the parser monad for
         -- checking and constructing values
         checkImportDecl,
-        checkExpBlockArguments,
+        checkExpBlockArguments, checkCmdBlockArguments,
         checkPrecP,           -- Int -> P Int
         checkContext,         -- HsType -> P HsContext
         checkPattern,         -- HsExp -> P HsPat
diff --git a/testsuite/tests/parser/should_compile/BlockArgumentsArrowCmds.hs b/testsuite/tests/parser/should_compile/BlockArgumentsArrowCmds.hs
new file mode 100644
index 0000000000..89ac332d31
--- /dev/null
+++ b/testsuite/tests/parser/should_compile/BlockArgumentsArrowCmds.hs
@@ -0,0 +1,22 @@
+{-# LANGUAGE Arrows, BlockArguments #-}
+
+module BlockArgumentsArrowCmds where
+
+import Control.Arrow
+
+cmdLam :: () -> ()
+cmdLam = proc () -> (| id \() -> () >- returnA |) ()
+
+cmdCase :: () -> ()
+cmdCase = proc () -> (| id case () of
+  () -> () >- returnA |)
+
+cmdIf :: () -> ()
+cmdIf = proc () -> (| id if True then () >- returnA else () >- returnA |)
+
+cmdLet :: () -> ()
+cmdLet = proc () -> (| id let x = () in x >- returnA |)
+
+cmdDo :: () -> ()
+cmdDo = proc () -> (| id do
+  () >- returnA |)
diff --git a/testsuite/tests/parser/should_compile/all.T b/testsuite/tests/parser/should_compile/all.T
index 85a7c3c172..fd69d32f0f 100644
--- a/testsuite/tests/parser/should_compile/all.T
+++ b/testsuite/tests/parser/should_compile/all.T
@@ -86,6 +86,7 @@ test('T3303', [], multimod_compile, ['T3303', '-v0'])
 test('T3741', normal, compile, [''])
 test('DoAndIfThenElse', normal, compile, [''])
 test('BlockArguments', normal, compile, [''])
+test('BlockArgumentsArrowCmds', normal, compile, [''])
 test('BlockArgumentsLambdaCase', normal, compile, [''])
 test('NoBlockArguments', normal, compile, [''])
 test('NondecreasingIndentation', normal, compile, [''])
diff --git a/testsuite/tests/parser/should_fail/NoBlockArgumentsFailArrowCmds.hs b/testsuite/tests/parser/should_fail/NoBlockArgumentsFailArrowCmds.hs
new file mode 100644
index 0000000000..866b0af4cb
--- /dev/null
+++ b/testsuite/tests/parser/should_fail/NoBlockArgumentsFailArrowCmds.hs
@@ -0,0 +1,7 @@
+{-# LANGUAGE Arrows #-}
+module NoBlockArgumentsFailArrowCmds where
+
+import Control.Arrow
+
+cmdLam :: () -> ()
+cmdLam = proc () -> (| id \() -> () >- returnA |) ()
diff --git a/testsuite/tests/parser/should_fail/NoBlockArgumentsFailArrowCmds.stderr b/testsuite/tests/parser/should_fail/NoBlockArgumentsFailArrowCmds.stderr
new file mode 100644
index 0000000000..2a99cda137
--- /dev/null
+++ b/testsuite/tests/parser/should_fail/NoBlockArgumentsFailArrowCmds.stderr
@@ -0,0 +1,6 @@
+
+NoBlockArgumentsFailArrowCmds.hs:7:27: error:
+    Unexpected lambda command in function application:
+        \ () -> () >- returnA
+    You could write it with parentheses
+    Or perhaps you meant to enable BlockArguments?
diff --git a/testsuite/tests/parser/should_fail/all.T b/testsuite/tests/parser/should_fail/all.T
index e0000f009e..c6d691bed3 100644
--- a/testsuite/tests/parser/should_fail/all.T
+++ b/testsuite/tests/parser/should_fail/all.T
@@ -77,6 +77,7 @@ test('NoPatternSynonyms', normal, compile_fail, [''])
 test('NoBlockArgumentsFail', normal, compile_fail, [''])
 test('NoBlockArgumentsFail2', normal, compile_fail, [''])
 test('NoBlockArgumentsFail3', normal, compile_fail, [''])
+test('NoBlockArgumentsFailArrowCmds', normal, compile_fail, [''])
 test('NondecreasingIndentationFail', normal, compile_fail, [''])
 test('readFailTraditionalRecords1', normal, compile_fail, [''])
 test('readFailTraditionalRecords2', normal, compile_fail, [''])