Allow `let` just before pure/return in ApplicativeDo

The following is currently rejected:

```haskell
-- F is an Applicative but not a Monad
x :: F (Int, Int)
x = do
  a <- pure 0
  let b = 1
  pure (a, b)
```

This has bitten me multiple times. This MR contains a simple fix:
only allow a "let only" segment to be merged with the next (and not
the previous) segment. As a result, when the last one or more
statements before pure/return are `LetStmt`s, there will be one
more segment containing only those `LetStmt`s.

Note that if the `let` statement mentions a name bound previously, then
the program is still rejected, for example

```haskell
x = do
  a <- pure 0
  let b = a + 1
  pure (a, b)
```

or the example in #18559. To support this would require a more
complex approach, but this is IME much less common than the
previous case.

7 files changed, 50 insertions, 19 deletions

diff --git a/compiler/GHC/Rename/Expr.hs b/compiler/GHC/Rename/Expr.hs
index 148d401f91..055afd7e84 100644
--- a/compiler/GHC/Rename/Expr.hs
+++ b/compiler/GHC/Rename/Expr.hs
@@ -2107,7 +2107,9 @@ stmtTreeToStmts monad_names ctxt (StmtTreeApplicative trees) tail tail_fvs = do
         if | L _ ApplicativeStmt{} <- last stmts' ->
              return (unLoc tup, emptyNameSet)
            | otherwise -> do
-             (ret, _) <- lookupQualifiedDoExpr (HsDoStmt ctxt) returnMName
+             -- Need 'pureAName' and not 'returnMName' here, so that it requires
+             -- 'Applicative' and not 'Monad' whenever possible (until #20540 is fixed).
+             (ret, _) <- lookupQualifiedDoExpr (HsDoStmt ctxt) pureAName
              let expr = HsApp noComments (noLocA ret) tup
              return (expr, emptyFVs)
      return ( ApplicativeArgMany
@@ -2125,19 +2127,19 @@ stmtTreeToStmts monad_names ctxt (StmtTreeApplicative trees) tail tail_fvs = do
 segments
   :: [(ExprLStmt GhcRn, FreeVars)]
   -> [[(ExprLStmt GhcRn, FreeVars)]]
-segments stmts = map fst $ merge $ reverse $ map reverse $ walk (reverse stmts)
+segments stmts = merge $ reverse $ map reverse $ walk (reverse stmts)
   where
     allvars = mkNameSet (concatMap (collectStmtBinders CollNoDictBinders . unLoc . fst) stmts)
 
     -- We would rather not have a segment that just has LetStmts in
-    -- it, so combine those with an adjacent segment where possible.
+    -- it, so combine those with the next segment where possible.
+    -- We don't merge it with the previous segment because the merged segment
+    -- would require 'Monad' while it may otherwise only require 'Applicative'.
     merge [] = []
     merge (seg : segs)
        = case rest of
-          [] -> [(seg,all_lets)]
-          ((s,s_lets):ss) | all_lets || s_lets
-               -> (seg ++ s, all_lets && s_lets) : ss
-          _otherwise -> (seg,all_lets) : rest
+          s:ss | all_lets -> (seg ++ s) : ss
+          _otherwise -> seg : rest
       where
         rest = merge segs
         all_lets = all (isLetStmt . fst) seg
diff --git a/docs/users_guide/exts/applicative_do.rst b/docs/users_guide/exts/applicative_do.rst
index 460ae3d162..a46d74d957 100644
--- a/docs/users_guide/exts/applicative_do.rst
+++ b/docs/users_guide/exts/applicative_do.rst
@@ -71,9 +71,11 @@ is as follows. If the do-expression has the following form: ::
 
 where none of the variables defined by ``p1...pn`` are mentioned in ``E1...En``,
 and ``p1...pn`` are all variables or lazy patterns,
-then the expression will only require ``Applicative``. Otherwise, the expression
+then the expression will only require ``Applicative``. The do expression may also
+contain ``let`` statements anywhere, provided that the right-hand-sides of the ``let``
+bindings do not mention any of ``p1...pn``. Otherwise, the expression
 will require ``Monad``. The block may return a pure expression ``E`` depending
-upon the results ``p1...pn`` with either ``return`` or ``pure``.
+upon the results ``p1...pn`` and the ``let`` bindings, with either ``return`` or ``pure``.
 
 Note: the final statement must match one of these patterns exactly:
 
@@ -187,6 +189,3 @@ terms of ``Monad`` is to use the ``Monad`` operations directly, e.g. ::
     instance Applicative MyType where
         pure = return
         (<*>) = ap
-
-
-
diff --git a/testsuite/tests/ado/ado001.hs b/testsuite/tests/ado/ado001.hs
index 6abce177e0..a72d0ca3a7 100644
--- a/testsuite/tests/ado/ado001.hs
+++ b/testsuite/tests/ado/ado001.hs
@@ -2,6 +2,7 @@
 module Main where
 
 import Control.Applicative
+import Data.Maybe
 import Text.PrettyPrint as PP
 
 (a:b:c:d:e:f:g:h:_) = map (\c -> doc [c]) ['a'..]
@@ -85,7 +86,7 @@ test8 = do
   return () `const` (x2,x3,x4)
 
 -- test that Lets don't get in the way
--- ((a | (b; c)) | d) | e
+-- (((a | (b; c)) | d) | e) | pure
 test9 :: M ()
 test9 = do
   x1 <- a
@@ -109,7 +110,7 @@ test10 = do
   x5 <- e
   return (const () (x3,x4,x5))
 
--- (a | b)
+-- (a | b) | pure
 -- This demonstrated a bug in RnExpr.segments (#11612)
 test11 :: M ()
 test11 = do
@@ -160,10 +161,10 @@ instance Functor M where
 instance Applicative M where
   pure a = M $ \_ -> (Nothing, a)
   M f <*> M a = M $ \p ->
-    let (Just d1, f') = f True
-        (Just d2, a') = a True
+    let (d1, f') = f True
+        (d2, a') = a True
     in
-        (Just (maybeParen p (d1 <+> char '|' <+> d2)), f' a')
+        (Just (maybeParen p (fromMaybe (text "pure") d1 <+> char '|' <+> fromMaybe (text "pure") d2)), f' a')
 
 instance Monad M where
   return = pure
diff --git a/testsuite/tests/ado/ado001.stdout b/testsuite/tests/ado/ado001.stdout
index 365860f55e..f4f576e084 100644
--- a/testsuite/tests/ado/ado001.stdout
+++ b/testsuite/tests/ado/ado001.stdout
@@ -6,7 +6,7 @@ a; g
 a; (b | (c; (d; (e | (f; g)))))
 (a | b); (c | d)
 a; ((b | c) | d)
-((a | (b; c)) | d) | e
+(((a | (b; c)) | d) | e) | pure
 ((a | b); (c | d)) | e
-a | b
+(a | b) | pure
 a | (b; c)
diff --git a/testsuite/tests/ado/ado011.hs b/testsuite/tests/ado/ado011.hs
new file mode 100644
index 0000000000..5a164b4787
--- /dev/null
+++ b/testsuite/tests/ado/ado011.hs
@@ -0,0 +1,26 @@
+{-# LANGUAGE ApplicativeDo, DerivingStrategies, GeneralizedNewtypeDeriving #-}
+
+module Main where
+
+import Data.Functor.Identity
+
+newtype F a = F (Identity a)
+  deriving newtype (Functor, Applicative, Show)
+
+x :: F (Int, Int)
+x = do
+  a <- pure 0
+  let b = 1
+  pure (a, b)
+
+y :: F (Int, Int, Int)
+y = do
+   a <- pure 0
+   let b = 1
+   let c = b + 1
+   pure (a, b, c)
+
+main :: IO ()
+main = do
+  print x
+  print y
diff --git a/testsuite/tests/ado/ado011.stdout b/testsuite/tests/ado/ado011.stdout
new file mode 100644
index 0000000000..6cb5d29c6c
--- /dev/null
+++ b/testsuite/tests/ado/ado011.stdout
@@ -0,0 +1,2 @@
+Identity (0,1)
+Identity (0,1,2)
diff --git a/testsuite/tests/ado/all.T b/testsuite/tests/ado/all.T
index 7369f9e986..7bd416a570 100644
--- a/testsuite/tests/ado/all.T
+++ b/testsuite/tests/ado/all.T
@@ -8,6 +8,7 @@ test('ado007', normal, compile, [''])
 test('ado008', normal, compile, [''])
 test('ado009', normal, compile, [''])
 test('ado010', normal, compile, [''])
+test('ado011', normal, compile_and_run, [''])
 test('T11607', normal, compile_and_run, [''])
 test('ado-optimal', normal, compile_and_run, [''])
 test('T12490', normal, compile, [''])