pretty-printer: Properly parenthesise LastStmt

After ApplicatveDo strips the last `return` during renaming, the pretty
printer has to restore it. However, if the return was followed by `$`,
the dollar was stripped too and not restored.

For example, the last stamement in:

```
  foo = do
    x <- ...
    ...
    return $ f x
```

would be printed as:

```
    return f x
```

This commit preserved the dolar, so it becomes:

```
    return $ f x
```

7 files changed, 40 insertions, 14 deletions

diff --git a/compiler/GHC/Hs/Expr.hs b/compiler/GHC/Hs/Expr.hs
index 0c14332d49..473868c4e8 100644
--- a/compiler/GHC/Hs/Expr.hs
+++ b/compiler/GHC/Hs/Expr.hs
@@ -1827,7 +1827,10 @@ data StmtLR idL idR body -- body should always be (LHs**** idR)
               -- Not used for GhciStmtCtxt, PatGuard, which scope over other stuff
           (XLastStmt idL idR body)
           body
-          Bool               -- True <=> return was stripped by ApplicativeDo
+          (Maybe Bool)  -- Whether return was stripped
+            -- Just True <=> return with a dollar was stripped by ApplicativeDo
+            -- Just False <=> return without a dollar was stripped by ApplicativeDo
+            -- Nothing <=> Nothing was stripped
           (SyntaxExpr idR)   -- The return operator
             -- The return operator is used only for MonadComp
             -- For ListComp we use the baked-in 'return'
@@ -2213,10 +2216,13 @@ pprStmt :: forall idL idR body . (OutputableBndrId idL,
                                   OutputableBndrId idR,
                                   Outputable body)
         => (StmtLR (GhcPass idL) (GhcPass idR) body) -> SDoc
-pprStmt (LastStmt _ expr ret_stripped _)
+pprStmt (LastStmt _ expr m_dollar_stripped _)
   = whenPprDebug (text "[last]") <+>
-       (if ret_stripped then text "return" else empty) <+>
-       ppr expr
+      (case m_dollar_stripped of
+        Just True -> text "return $"
+        Just False -> text "return"
+        Nothing -> empty) <+>
+      ppr expr
 pprStmt (BindStmt _ pat expr _ _) = hsep [ppr pat, larrow, ppr expr]
 pprStmt (LetStmt _ (L _ binds))   = hsep [text "let", pprBinds binds]
 pprStmt (BodyStmt _ expr _ _)     = ppr expr
@@ -2284,7 +2290,7 @@ pprStmt (ApplicativeStmt _ args mb_join)
      text "<-" <+>
      ppr (HsDo (panic "pprStmt") DoExpr (noLoc
                (stmts ++
-                   [noLoc (LastStmt noExtField (noLoc return) False noSyntaxExpr)])))
+                   [noLoc (LastStmt noExtField (noLoc return) Nothing noSyntaxExpr)])))
    pp_arg (_, XApplicativeArg x) = ppr x
 
 pprStmt (XStmtLR x) = ppr x
diff --git a/compiler/GHC/Hs/Utils.hs b/compiler/GHC/Hs/Utils.hs
index 36e0c6d2c4..ff1af93a6d 100644
--- a/compiler/GHC/Hs/Utils.hs
+++ b/compiler/GHC/Hs/Utils.hs
@@ -316,7 +316,7 @@ mkTransformByStmt  ss u b = emptyTransStmt { trS_form = ThenForm,  trS_stmts = s
 mkGroupUsingStmt   ss u   = emptyTransStmt { trS_form = GroupForm, trS_stmts = ss, trS_using = u }
 mkGroupByUsingStmt ss b u = emptyTransStmt { trS_form = GroupForm, trS_stmts = ss, trS_using = u, trS_by = Just b }
 
-mkLastStmt body = LastStmt noExtField body False noSyntaxExpr
+mkLastStmt body = LastStmt noExtField body Nothing noSyntaxExpr
 mkBodyStmt body
   = BodyStmt noExtField body noSyntaxExpr noSyntaxExpr
 mkBindStmt pat body
diff --git a/compiler/GHC/HsToCore/ListComp.hs b/compiler/GHC/HsToCore/ListComp.hs
index 6c58be3a47..281795c8ac 100644
--- a/compiler/GHC/HsToCore/ListComp.hs
+++ b/compiler/GHC/HsToCore/ListComp.hs
@@ -638,7 +638,7 @@ dsInnerMonadComp :: [ExprLStmt GhcTc]
                  -> DsM CoreExpr
 dsInnerMonadComp stmts bndrs ret_op
   = dsMcStmts (stmts ++
-                 [noLoc (LastStmt noExtField (mkBigLHsVarTupId bndrs) False ret_op)])
+                 [noLoc (LastStmt noExtField (mkBigLHsVarTupId bndrs) Nothing ret_op)])
 
 
 -- The `unzip` function for `GroupStmt` in a monad comprehensions
diff --git a/compiler/GHC/Rename/Expr.hs b/compiler/GHC/Rename/Expr.hs
index baadf419ad..eca5b42e3e 100644
--- a/compiler/GHC/Rename/Expr.hs
+++ b/compiler/GHC/Rename/Expr.hs
@@ -1956,19 +1956,20 @@ needJoin :: MonadNames
          -> (Bool, [ExprLStmt GhcRn])
 needJoin _monad_names [] = (False, [])  -- we're in an ApplicativeArg
 needJoin monad_names  [L loc (LastStmt _ e _ t)]
- | Just arg <- isReturnApp monad_names e =
-       (False, [L loc (LastStmt noExtField arg True t)])
+ | Just (arg, wasDollar) <- isReturnApp monad_names e =
+       (False, [L loc (LastStmt noExtField arg (Just wasDollar) t)])
 needJoin _monad_names stmts = (True, stmts)
 
--- | @Just e@, if the expression is @return e@ or @return $ e@,
--- otherwise @Nothing@
+-- | @(Just e, False)@, if the expression is @return e@
+--   @(Just e, True)@ if the expression is @return $ e@,
+--   otherwise @Nothing@.
 isReturnApp :: MonadNames
             -> LHsExpr GhcRn
-            -> Maybe (LHsExpr GhcRn)
+            -> Maybe (LHsExpr GhcRn, Bool)
 isReturnApp monad_names (L _ (HsPar _ expr)) = isReturnApp monad_names expr
 isReturnApp monad_names (L _ e) = case e of
-  OpApp _ l op r | is_return l, is_dollar op -> Just r
-  HsApp _ f arg  | is_return f               -> Just arg
+  OpApp _ l op r | is_return l, is_dollar op -> Just (r, True)
+  HsApp _ f arg  | is_return f               -> Just (arg, False)
   _otherwise -> Nothing
  where
   is_var f (L _ (HsPar _ e)) = is_var f e
diff --git a/testsuite/tests/ado/ado009.hs b/testsuite/tests/ado/ado009.hs
new file mode 100644
index 0000000000..876a4fba90
--- /dev/null
+++ b/testsuite/tests/ado/ado009.hs
@@ -0,0 +1,10 @@
+{-# LANGUAGE ApplicativeDo #-}
+{-# OPTIONS_GHC -ddump-rn -dsuppress-uniques #-}
+
+module Test where
+
+-- Make sure the $ stripped from the last stmt is printed
+q :: IO ()
+q = do
+  a <- return ()
+  return $ (\_ -> ()) a
diff --git a/testsuite/tests/ado/ado009.stderr b/testsuite/tests/ado/ado009.stderr
new file mode 100644
index 0000000000..19a5defd4c
--- /dev/null
+++ b/testsuite/tests/ado/ado009.stderr
@@ -0,0 +1,8 @@
+
+==================== Renamer ====================
+Test.q :: IO ()
+Test.q
+  = do a <- return ()
+       return $ (\ _ -> ()) a
+
+
diff --git a/testsuite/tests/ado/all.T b/testsuite/tests/ado/all.T
index 634aae2314..11a9f4d6c8 100644
--- a/testsuite/tests/ado/all.T
+++ b/testsuite/tests/ado/all.T
@@ -6,6 +6,7 @@ test('ado005', normal, compile_fail, [''])
 test('ado006', normal, compile, [''])
 test('ado007', normal, compile, [''])
 test('ado008', normal, compile, [''])
+test('ado009', normal, compile, [''])
 test('T11607', normal, compile_and_run, [''])
 test('ado-optimal', normal, compile_and_run, [''])
 test('T12490', normal, compile, [''])