TcMatches.hs 48.5 KB
Newer Older
Austin Seipp's avatar
Austin Seipp committed
1 2 3 4
{-
(c) The University of Glasgow 2006
(c) The GRASP/AQUA Project, Glasgow University, 1992-1998

5 6

TcMatches: Typecheck some @Matches@
Austin Seipp's avatar
Austin Seipp committed
7
-}
8

quchen's avatar
quchen committed
9 10 11
{-# LANGUAGE CPP #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE MultiWayIf #-}
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
12
{-# LANGUAGE TupleSections #-}
13
{-# LANGUAGE FlexibleContexts #-}
14
{-# LANGUAGE TypeFamilies #-}
Ian Lynagh's avatar
Ian Lynagh committed
15

16
module TcMatches ( tcMatchesFun, tcGRHS, tcGRHSsPat, tcMatchesCase, tcMatchLambda,
17 18 19
                   TcMatchCtxt(..), TcStmtChecker, TcExprStmtChecker, TcCmdStmtChecker,
                   tcStmts, tcStmtsAndThen, tcDoStmts, tcBody,
                   tcDoStmt, tcGuardStmt
20
       ) where
21

eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
22 23
import {-# SOURCE #-}   TcExpr( tcSyntaxOp, tcInferSigmaNC, tcInferSigma
                              , tcCheckId, tcMonoExpr, tcMonoExprNC, tcPolyExpr )
sof's avatar
sof committed
24

Ben Gamari's avatar
Ben Gamari committed
25
import BasicTypes (LexicalFixity(..))
26
import HsSyn
27
import TcRnMonad
28 29 30 31 32 33 34 35 36 37
import TcEnv
import TcPat
import TcMType
import TcType
import TcBinds
import TcUnify
import Name
import TysWiredIn
import Id
import TyCon
38
import TysPrim
39
import TcEvidence
sof's avatar
sof committed
40
import Outputable
41
import Util
42
import SrcLoc
quchen's avatar
quchen committed
43 44
import DynFlags
import PrelNames (monadFailClassName)
45
import qualified GHC.LanguageExtensions as LangExt
46

47 48 49
-- Create chunkified tuple tybes for monad comprehensions
import MkCore

50
import Control.Monad
51
import Control.Arrow ( second )
52 53

#include "HsVersions.h"
54

Austin Seipp's avatar
Austin Seipp committed
55 56 57
{-
************************************************************************
*                                                                      *
58
\subsection{tcMatchesFun, tcMatchesCase}
Austin Seipp's avatar
Austin Seipp committed
59 60
*                                                                      *
************************************************************************
61

62 63 64 65 66
@tcMatchesFun@ typechecks a @[Match]@ list which occurs in a
@FunMonoBind@.  The second argument is the name of the function, which
is used in error messages.  It checks that all the equations have the
same number of arguments before using @tcMatches@ to do the work.

67 68 69
Note [Polymorphic expected type for tcMatchesFun]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
tcMatchesFun may be given a *sigma* (polymorphic) type
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
70
so it must be prepared to use tcSkolemise to skolemise it.
71
See Note [sig_tau may be polymorphic] in TcPat.
Austin Seipp's avatar
Austin Seipp committed
72
-}
73

74
tcMatchesFun :: Located Name
75
             -> MatchGroup GhcRn (LHsExpr GhcRn)
76
             -> ExpRhoType     -- Expected type of function
77
             -> TcM (HsWrapper, MatchGroup GhcTcId (LHsExpr GhcTcId))
78
                                -- Returns type of body
79
tcMatchesFun fn@(L _ fun_name) matches exp_ty
80 81 82 83 84 85
  = do  {  -- Check that they all have the same no of arguments
           -- Location is in the monad, set the caller so that
           -- any inter-equation error messages get some vaguely
           -- sensible location.        Note: we have to do this odd
           -- ann-grabbing, because we don't always have annotations in
           -- hand when we call tcMatchesFun...
86
          traceTc "tcMatchesFun" (ppr fun_name $$ ppr exp_ty)
87
        ; checkArgs fun_name matches
88

89
        ; (wrap_gen, (wrap_fun, group))
90
            <- tcSkolemiseET (FunSigCtxt fun_name True) exp_ty $ \ exp_rho ->
91
                  -- Note [Polymorphic expected type for tcMatchesFun]
92 93 94
               do { (matches', wrap_fun)
                       <- matchExpectedFunTys herald arity exp_rho $
                          \ pat_tys rhs_ty ->
Simon Peyton Jones's avatar
Simon Peyton Jones committed
95
                          tcMatches match_ctxt pat_tys rhs_ty matches
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
96
                  ; return (wrap_fun, matches') }
97
        ; return (wrap_gen <.> wrap_fun, group) }
ross's avatar
ross committed
98
  where
99
    arity = matchGroupArity matches
100 101
    herald = text "The equation(s) for"
             <+> quotes (ppr fun_name) <+> text "have"
Simon Peyton Jones's avatar
Simon Peyton Jones committed
102 103
    what = FunRhs { mc_fun = fn, mc_fixity = Prefix, mc_strictness = strictness }
    match_ctxt = MC { mc_what = what, mc_body = tcBody }
Ben Gamari's avatar
Ben Gamari committed
104 105
    strictness
      | [L _ match] <- unLoc $ mg_alts matches
Simon Peyton Jones's avatar
Simon Peyton Jones committed
106
      , FunRhs{ mc_strictness = SrcStrict } <- m_ctxt match
Ben Gamari's avatar
Ben Gamari committed
107 108 109
      = SrcStrict
      | otherwise
      = NoSrcStrict
110

Austin Seipp's avatar
Austin Seipp committed
111
{-
112 113
@tcMatchesCase@ doesn't do the argument-count check because the
parser guarantees that each equation has exactly one argument.
Austin Seipp's avatar
Austin Seipp committed
114
-}
115

116 117 118 119 120 121 122 123
tcMatchesCase :: (Outputable (body GhcRn)) =>
                TcMatchCtxt body                        -- Case context
             -> TcSigmaType                             -- Type of scrutinee
             -> MatchGroup GhcRn (Located (body GhcRn)) -- The case alternatives
             -> ExpRhoType                    -- Type of whole case expressions
             -> TcM (MatchGroup GhcTcId (Located (body GhcTcId)))
                -- Translated alternatives
                -- wrapper goes from MatchGroup's ty to expected ty
124

125
tcMatchesCase ctxt scrut_ty matches res_ty
Simon Peyton Jones's avatar
Simon Peyton Jones committed
126
  = tcMatches ctxt [mkCheckExpType scrut_ty] res_ty matches
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
127 128 129

tcMatchLambda :: SDoc -- see Note [Herald for matchExpectedFunTys] in TcUnify
              -> TcMatchCtxt HsExpr
130
              -> MatchGroup GhcRn (LHsExpr GhcRn)
131
              -> ExpRhoType   -- deeply skolemised
132
              -> TcM (MatchGroup GhcTcId (LHsExpr GhcTcId), HsWrapper)
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
133
tcMatchLambda herald match_ctxt match res_ty
Simon Peyton Jones's avatar
Simon Peyton Jones committed
134 135
  = matchExpectedFunTys herald n_pats res_ty $ \ pat_tys rhs_ty ->
    tcMatches match_ctxt pat_tys rhs_ty match
ross's avatar
ross committed
136
  where
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
137 138
    n_pats | isEmptyMatchGroup match = 1   -- must be lambda-case
           | otherwise               = matchGroupArity match
139

Austin Seipp's avatar
Austin Seipp committed
140
-- @tcGRHSsPat@ typechecks @[GRHSs]@ that occur in a @PatMonoBind@.
ross's avatar
ross committed
141

142 143
tcGRHSsPat :: GRHSs GhcRn (LHsExpr GhcRn) -> TcRhoType
           -> TcM (GRHSs GhcTcId (LHsExpr GhcTcId))
144
-- Used for pattern bindings
145
tcGRHSsPat grhss res_ty = tcGRHSs match_ctxt grhss (mkCheckExpType res_ty)
ross's avatar
ross committed
146 147
  where
    match_ctxt = MC { mc_what = PatBindRhs,
148
                      mc_body = tcBody }
149

Austin Seipp's avatar
Austin Seipp committed
150 151 152
{-
************************************************************************
*                                                                      *
153
\subsection{tcMatch}
Austin Seipp's avatar
Austin Seipp committed
154 155
*                                                                      *
************************************************************************
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
156

157 158
Note [Case branches must never infer a non-tau type]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
159 160 161 162 163 164 165 166 167 168
Consider

  case ... of
    ... -> \(x :: forall a. a -> a) -> x
    ... -> \y -> y

Should that type-check? The problem is that, if we check the second branch
first, then we'll get a type (b -> b) for the branches, which won't unify
with the polytype in the first branch. If we check the first branch first,
then everything is OK. This order-dependency is terrible. So we want only
169 170 171
proper tau-types in branches (unless a sigma-type is pushed down).
This is what expTypeToType ensures: it replaces an Infer with a fresh
tau-type.
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
172 173 174 175 176 177

An even trickier case looks like

  f x True  = x undefined
  f x False = x ()

178 179
Here, we see that the arguments must also be non-Infer. Thus, we must
use expTypeToType on the output of matchExpectedFunTys, not the input.
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
180 181 182 183 184 185

But we make a special case for a one-branch case. This is so that

  f = \(x :: forall a. a -> a) -> x

still gets assigned a polytype.
Austin Seipp's avatar
Austin Seipp committed
186
-}
187

188
-- | When the MatchGroup has multiple RHSs, convert an Infer ExpType in the
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
189
-- expected type into TauTvs.
190
-- See Note [Case branches must never infer a non-tau type]
Simon Peyton Jones's avatar
Simon Peyton Jones committed
191
tauifyMultipleMatches :: [LMatch id body]
192 193 194 195 196
                      -> [ExpType] -> TcM [ExpType]
tauifyMultipleMatches group exp_tys
  | isSingletonMatchGroup group = return exp_tys
  | otherwise                   = mapM tauifyExpType exp_tys
  -- NB: In the empty-match case, this ensures we fill in the ExpType
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
197

Simon Peyton Jones's avatar
Simon Peyton Jones committed
198
-- | Type-check a MatchGroup.
199
tcMatches :: (Outputable (body GhcRn)) => TcMatchCtxt body
200 201
          -> [ExpSigmaType]      -- Expected pattern types
          -> ExpRhoType          -- Expected result-type of the Match.
202 203
          -> MatchGroup GhcRn (Located (body GhcRn))
          -> TcM (MatchGroup GhcTcId (Located (body GhcTcId)))
204 205

data TcMatchCtxt body   -- c.f. TcStmtCtxt, also in this module
206 207
  = MC { mc_what :: HsMatchContext Name,  -- What kind of thing this is
         mc_body :: Located (body GhcRn)         -- Type checker for a body of
208
                                                -- an alternative
209
                 -> ExpRhoType
210
                 -> TcM (Located (body GhcTcId)) }
211

eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
212 213
tcMatches ctxt pat_tys rhs_ty (MG { mg_alts = L l matches
                                  , mg_origin = origin })
Simon Peyton Jones's avatar
Simon Peyton Jones committed
214
  = do { rhs_ty:pat_tys <- tauifyMultipleMatches matches (rhs_ty:pat_tys)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
215 216
            -- See Note [Case branches must never infer a non-tau type]

Simon Peyton Jones's avatar
Simon Peyton Jones committed
217
       ; matches' <- mapM (tcMatch ctxt pat_tys rhs_ty) matches
218 219
       ; pat_tys  <- mapM readExpType pat_tys
       ; rhs_ty   <- readExpType rhs_ty
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
220 221 222 223
       ; return (MG { mg_alts = L l matches'
                    , mg_arg_tys = pat_tys
                    , mg_res_ty = rhs_ty
                    , mg_origin = origin }) }
224 225

-------------
226
tcMatch :: (Outputable (body GhcRn)) => TcMatchCtxt body
227 228
        -> [ExpSigmaType]        -- Expected pattern types
        -> ExpRhoType            -- Expected result-type of the Match.
229 230
        -> LMatch GhcRn (Located (body GhcRn))
        -> TcM (LMatch GhcTcId (Located (body GhcTcId)))
231

232
tcMatch ctxt pat_tys rhs_ty match
233
  = wrapLocM (tc_match ctxt pat_tys rhs_ty) match
234
  where
235 236
    tc_match ctxt pat_tys rhs_ty
             match@(Match { m_pats = pats, m_type = maybe_rhs_sig, m_grhss = grhss })
237
      = add_match_ctxt match $
238
        do { (pats', grhss') <- tcPats (mc_what ctxt) pats pat_tys $
239
                                tc_grhss ctxt maybe_rhs_sig grhss rhs_ty
240 241
           ; return (Match { m_ctxt = mc_what ctxt, m_pats = pats'
                           , m_type = Nothing, m_grhss = grhss' }) }
242

243 244
    tc_grhss ctxt Nothing grhss rhs_ty
      = tcGRHSs ctxt grhss rhs_ty       -- No result signature
245

246
        -- Result type sigs are no longer supported
247
    tc_grhss _ (Just {}) _ _
248
      = panic "tc_ghrss"        -- Rejected by renamer
249

250
        -- For (\x -> e), tcExpr has already said "In the expression \x->e"
251
        -- so we don't want to add "In the lambda abstraction \x->e"
252
    add_match_ctxt match thing_inside
253 254
        = case mc_what ctxt of
            LambdaExpr -> thing_inside
255
            _          -> addErrCtxt (pprMatchInCtxt match) thing_inside
256

257
-------------
258 259
tcGRHSs :: TcMatchCtxt body -> GRHSs GhcRn (Located (body GhcRn)) -> ExpRhoType
        -> TcM (GRHSs GhcTcId (Located (body GhcTcId)))
260

261 262
-- Notice that we pass in the full res_ty, so that we get
-- good inference from simple things like
263
--      f = \(x::forall a.a->a) -> <stuff>
264 265
-- We used to force it to be a monotype when there was more than one guard
-- but we don't need to do that any more
ross's avatar
ross committed
266

267
tcGRHSs ctxt (GRHSs grhss (L l binds)) res_ty
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
268 269 270
  = do  { (binds', grhss')
            <- tcLocalBinds binds $
               mapM (wrapLocM (tcGRHS ctxt res_ty)) grhss
271

272
        ; return (GRHSs grhss' (L l binds')) }
273 274

-------------
275 276
tcGRHS :: TcMatchCtxt body -> ExpRhoType -> GRHS GhcRn (Located (body GhcRn))
       -> TcM (GRHS GhcTcId (Located (body GhcTcId)))
277

278
tcGRHS ctxt res_ty (GRHS guards rhs)
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
279 280 281
  = do  { (guards', rhs')
            <- tcStmtsAndThen stmt_ctxt tcGuardStmt guards res_ty $
               mc_body ctxt rhs
282
        ; return (GRHS guards' rhs') }
283
  where
284
    stmt_ctxt  = PatGuard (mc_what ctxt)
285

Austin Seipp's avatar
Austin Seipp committed
286 287 288
{-
************************************************************************
*                                                                      *
289
\subsection{@tcDoStmts@ typechecks a {\em list} of do statements}
Austin Seipp's avatar
Austin Seipp committed
290 291 292
*                                                                      *
************************************************************************
-}
293

294
tcDoStmts :: HsStmtContext Name
295
          -> Located [LStmt GhcRn (LHsExpr GhcRn)]
296
          -> ExpRhoType
297
          -> TcM (HsExpr GhcTcId)          -- Returns a HsDo
298
tcDoStmts ListComp (L l stmts) res_ty
299 300
  = do  { res_ty <- expTypeToType res_ty
        ; (co, elt_ty) <- matchExpectedListTy res_ty
301
        ; let list_ty = mkListTy elt_ty
302 303
        ; stmts' <- tcStmts ListComp (tcLcStmt listTyCon) stmts
                            (mkCheckExpType elt_ty)
304
        ; return $ mkHsWrapCo co (HsDo ListComp (L l stmts') list_ty) }
305

306
tcDoStmts PArrComp (L l stmts) res_ty
307 308
  = do  { res_ty <- expTypeToType res_ty
        ; (co, elt_ty) <- matchExpectedPArrTy res_ty
309
        ; let parr_ty = mkPArrTy elt_ty
310 311
        ; stmts' <- tcStmts PArrComp (tcLcStmt parrTyCon) stmts
                            (mkCheckExpType elt_ty)
312
        ; return $ mkHsWrapCo co (HsDo PArrComp (L l stmts') parr_ty) }
313

314
tcDoStmts DoExpr (L l stmts) res_ty
315
  = do  { stmts' <- tcStmts DoExpr tcDoStmt stmts res_ty
316
        ; res_ty <- readExpType res_ty
317
        ; return (HsDo DoExpr (L l stmts') res_ty) }
318

319
tcDoStmts MDoExpr (L l stmts) res_ty
320
  = do  { stmts' <- tcStmts MDoExpr tcDoStmt stmts res_ty
321
        ; res_ty <- readExpType res_ty
322
        ; return (HsDo MDoExpr (L l stmts') res_ty) }
323

324
tcDoStmts MonadComp (L l stmts) res_ty
325
  = do  { stmts' <- tcStmts MonadComp tcMcStmt stmts res_ty
326
        ; res_ty <- readExpType res_ty
327
        ; return (HsDo MonadComp (L l stmts') res_ty) }
328

329
tcDoStmts ctxt _ _ = pprPanic "tcDoStmts" (pprStmtContext ctxt)
330

331
tcBody :: LHsExpr GhcRn -> ExpRhoType -> TcM (LHsExpr GhcTcId)
332
tcBody body res_ty
333
  = do  { traceTc "tcBody" (ppr res_ty)
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
334
        ; tcMonoExpr body res_ty
335
        }
336

Austin Seipp's avatar
Austin Seipp committed
337 338 339
{-
************************************************************************
*                                                                      *
340
\subsection{tcStmts}
Austin Seipp's avatar
Austin Seipp committed
341 342 343
*                                                                      *
************************************************************************
-}
344

345 346
type TcExprStmtChecker = TcStmtChecker HsExpr ExpRhoType
type TcCmdStmtChecker  = TcStmtChecker HsCmd  TcRhoType
347

348
type TcStmtChecker body rho_type
349
  =  forall thing. HsStmtContext Name
350
                -> Stmt GhcRn (Located (body GhcRn))
351 352
                -> rho_type                 -- Result type for comprehension
                -> (rho_type -> TcM thing)  -- Checker for what follows the stmt
353
                -> TcM (Stmt GhcTcId (Located (body GhcTcId)), thing)
354

355
tcStmts :: (Outputable (body GhcRn)) => HsStmtContext Name
356
        -> TcStmtChecker body rho_type   -- NB: higher-rank type
357
        -> [LStmt GhcRn (Located (body GhcRn))]
358
        -> rho_type
359
        -> TcM [LStmt GhcTcId (Located (body GhcTcId))]
360
tcStmts ctxt stmt_chk stmts res_ty
361
  = do { (stmts', _) <- tcStmtsAndThen ctxt stmt_chk stmts res_ty $
362 363 364
                        const (return ())
       ; return stmts' }

365
tcStmtsAndThen :: (Outputable (body GhcRn)) => HsStmtContext Name
366
               -> TcStmtChecker body rho_type    -- NB: higher-rank type
367
               -> [LStmt GhcRn (Located (body GhcRn))]
368 369
               -> rho_type
               -> (rho_type -> TcM thing)
370
               -> TcM ([LStmt GhcTcId (Located (body GhcTcId))], thing)
371

372 373
-- Note the higher-rank type.  stmt_chk is applied at different
-- types in the equations for tcStmts
374

375
tcStmtsAndThen _ _ [] res_ty thing_inside
376 377
  = do  { thing <- thing_inside res_ty
        ; return ([], thing) }
378

379
-- LetStmts are handled uniformly, regardless of context
380 381
tcStmtsAndThen ctxt stmt_chk (L loc (LetStmt (L l binds)) : stmts)
                                                             res_ty thing_inside
382
  = do  { (binds', (stmts',thing)) <- tcLocalBinds binds $
Simon Marlow's avatar
Simon Marlow committed
383
              tcStmtsAndThen ctxt stmt_chk stmts res_ty thing_inside
384
        ; return (L loc (LetStmt (L l binds')) : stmts', thing) }
385

Simon Marlow's avatar
Simon Marlow committed
386 387 388
-- Don't set the error context for an ApplicativeStmt.  It ought to be
-- possible to do this with a popErrCtxt in the tcStmt case for
-- ApplicativeStmt, but it did someting strange and broke a test (ado002).
389
tcStmtsAndThen ctxt stmt_chk (L loc stmt : stmts) res_ty thing_inside
Simon Marlow's avatar
Simon Marlow committed
390 391 392 393 394 395 396 397 398
  | ApplicativeStmt{} <- stmt
  = do  { (stmt', (stmts', thing)) <-
             stmt_chk ctxt stmt res_ty $ \ res_ty' ->
               tcStmtsAndThen ctxt stmt_chk stmts res_ty'  $
                 thing_inside
        ; return (L loc stmt' : stmts', thing) }

  -- For the vanilla case, handle the location-setting part
  | otherwise
399 400
  = do  { (stmt', (stmts', thing)) <-
                setSrcSpan loc                              $
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
401
                addErrCtxt (pprStmtInCtxt ctxt stmt)        $
402
                stmt_chk ctxt stmt res_ty                   $ \ res_ty' ->
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
403
                popErrCtxt                                  $
404 405 406
                tcStmtsAndThen ctxt stmt_chk stmts res_ty'  $
                thing_inside
        ; return (L loc stmt' : stmts', thing) }
407

408
---------------------------------------------------
409
--              Pattern guards
410 411
---------------------------------------------------

412 413
tcGuardStmt :: TcExprStmtChecker
tcGuardStmt _ (BodyStmt guard _ _ _) res_ty thing_inside
414
  = do  { guard' <- tcMonoExpr guard (mkCheckExpType boolTy)
415 416
        ; thing  <- thing_inside res_ty
        ; return (BodyStmt guard' noSyntaxExpr noSyntaxExpr boolTy, thing) }
417

418
tcGuardStmt ctxt (BindStmt pat rhs _ _ _) res_ty thing_inside
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
419 420
  = do  { (rhs', rhs_ty) <- tcInferSigmaNC rhs
                                   -- Stmt has a context already
421
        ; (pat', thing)  <- tcPat_O (StmtCtxt ctxt) (lexprCtOrigin rhs)
422
                                    pat (mkCheckExpType rhs_ty) $
423
                            thing_inside res_ty
424
        ; return (mkTcBindStmt pat' rhs', thing) }
425

Ian Lynagh's avatar
Ian Lynagh committed
426
tcGuardStmt _ stmt _ _
427 428 429
  = pprPanic "tcGuardStmt: unexpected Stmt" (ppr stmt)


430
---------------------------------------------------
431 432
--           List comprehensions and PArrays
--               (no rebindable syntax)
433 434 435 436 437
---------------------------------------------------

-- Dealt with separately, rather than by tcMcStmt, because
--   a) PArr isn't (yet) an instance of Monad, so the generality seems overkill
--   b) We have special desugaring rules for list comprehensions,
438
--      which avoid creating intermediate lists.  They in turn
439 440 441 442
--      assume that the bind/return operations are the regular
--      polymorphic ones, and in particular don't have any
--      coercion matching stuff in them.  It's hard to avoid the
--      potential for non-trivial coercions in tcMcStmt
443

444 445
tcLcStmt :: TyCon       -- The list/Parray type constructor ([] or PArray)
         -> TcExprStmtChecker
446

Simon Marlow's avatar
Simon Marlow committed
447
tcLcStmt _ _ (LastStmt body noret _) elt_ty thing_inside
448
  = do { body' <- tcMonoExprNC body elt_ty
449
       ; thing <- thing_inside (panic "tcLcStmt: thing_inside")
Simon Marlow's avatar
Simon Marlow committed
450
       ; return (LastStmt body' noret noSyntaxExpr, thing) }
451

452
-- A generator, pat <- rhs
453
tcLcStmt m_tc ctxt (BindStmt pat rhs _ _ _) elt_ty thing_inside
454
 = do   { pat_ty <- newFlexiTyVarTy liftedTypeKind
455 456
        ; rhs'   <- tcMonoExpr rhs (mkCheckExpType $ mkTyConApp m_tc [pat_ty])
        ; (pat', thing)  <- tcPat (StmtCtxt ctxt) pat (mkCheckExpType pat_ty) $
457
                            thing_inside elt_ty
458
        ; return (mkTcBindStmt pat' rhs', thing) }
459 460

-- A boolean guard
461
tcLcStmt _ _ (BodyStmt rhs _ _ _) elt_ty thing_inside
462
  = do  { rhs'  <- tcMonoExpr rhs (mkCheckExpType boolTy)
463 464
        ; thing <- thing_inside elt_ty
        ; return (BodyStmt rhs' noSyntaxExpr noSyntaxExpr boolTy, thing) }
465

466
-- ParStmt: See notes with tcMcStmt
467
tcLcStmt m_tc ctxt (ParStmt bndr_stmts_s _ _ _) elt_ty thing_inside
468
  = do  { (pairs', thing) <- loop bndr_stmts_s
469
        ; return (ParStmt pairs' noExpr noSyntaxExpr unitTy, thing) }
470
  where
471 472
    -- loop :: [([LStmt GhcRn], [GhcRn])]
    --      -> TcM ([([LStmt GhcTcId], [GhcTcId])], thing)
473
    loop [] = do { thing <- thing_inside elt_ty
474
                 ; return ([], thing) }         -- matching in the branches
475

476
    loop (ParStmtBlock stmts names _ : pairs)
477
      = do { (stmts', (ids, pairs', thing))
478 479 480 481 482
                <- tcStmtsAndThen ctxt (tcLcStmt m_tc) stmts elt_ty $ \ _elt_ty' ->
                   do { ids <- tcLookupLocalIds names
                      ; (pairs', thing) <- loop pairs
                      ; return (ids, pairs', thing) }
           ; return ( ParStmtBlock stmts' ids noSyntaxExpr : pairs', thing ) }
483

484 485 486 487 488
tcLcStmt m_tc ctxt (TransStmt { trS_form = form, trS_stmts = stmts
                              , trS_bndrs =  bindersMap
                              , trS_by = by, trS_using = using }) elt_ty thing_inside
  = do { let (bndr_names, n_bndr_names) = unzip bindersMap
             unused_ty = pprPanic "tcLcStmt: inner ty" (ppr bindersMap)
489 490
             -- The inner 'stmts' lack a LastStmt, so the element type
             --  passed in to tcStmtsAndThen is never looked at
491
       ; (stmts', (bndr_ids, by'))
492
            <- tcStmtsAndThen (TransStmtCtxt ctxt) (tcLcStmt m_tc) stmts unused_ty $ \_ -> do
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
493
               { by' <- traverse tcInferSigma by
494 495 496 497 498 499 500 501 502 503 504 505
               ; bndr_ids <- tcLookupLocalIds bndr_names
               ; return (bndr_ids, by') }

       ; let m_app ty = mkTyConApp m_tc [ty]

       --------------- Typecheck the 'using' function -------------
       -- using :: ((a,b,c)->t) -> m (a,b,c) -> m (a,b,c)m      (ThenForm)
       --       :: ((a,b,c)->t) -> m (a,b,c) -> m (m (a,b,c)))  (GroupForm)

         -- n_app :: Type -> Type   -- Wraps a 'ty' into '[ty]' for GroupForm
       ; let n_app = case form of
                       ThenForm -> (\ty -> ty)
506
                       _        -> m_app
507 508 509 510

             by_arrow :: Type -> Type     -- Wraps 'ty' to '(a->t) -> ty' if the By is present
             by_arrow = case by' of
                          Nothing       -> \ty -> ty
511
                          Just (_,e_ty) -> \ty -> (alphaTy `mkFunTy` e_ty) `mkFunTy` ty
512 513 514

             tup_ty        = mkBigCoreVarTupTy bndr_ids
             poly_arg_ty   = m_app alphaTy
515
             poly_res_ty   = m_app (n_app alphaTy)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
516
             using_poly_ty = mkInvForAllTy alphaTyVar $
517
                             by_arrow $
518 519 520
                             poly_arg_ty `mkFunTy` poly_res_ty

       ; using' <- tcPolyExpr using using_poly_ty
521
       ; let final_using = fmap (mkHsWrap (WpTyApp tup_ty)) using'
522

523 524 525
             -- 'stmts' returns a result of type (m1_ty tuple_ty),
             -- typically something like [(Int,Bool,Int)]
             -- We don't know what tuple_ty is yet, so we use a variable
526
       ; let mk_n_bndr :: Name -> TcId -> TcId
527
             mk_n_bndr n_bndr_name bndr_id = mkLocalIdOrCoVar n_bndr_name (n_app (idType bndr_id))
528 529 530

             -- Ensure that every old binder of type `b` is linked up with its
             -- new binder which should have type `n b`
531
             -- See Note [GroupStmt binder map] in HsExpr
532 533 534
             n_bndr_ids  = zipWith mk_n_bndr n_bndr_names bndr_ids
             bindersMap' = bndr_ids `zip` n_bndr_ids

535
       -- Type check the thing in the environment with
536 537 538
       -- these new binders and return the result
       ; thing <- tcExtendIdEnv n_bndr_ids (thing_inside elt_ty)

539 540 541 542 543 544 545
       ; return (TransStmt { trS_stmts = stmts', trS_bndrs = bindersMap'
                           , trS_by = fmap fst by', trS_using = final_using
                           , trS_ret = noSyntaxExpr
                           , trS_bind = noSyntaxExpr
                           , trS_fmap = noExpr
                           , trS_bind_arg_ty = unitTy
                           , trS_form = form }, thing) }
546

Ian Lynagh's avatar
Ian Lynagh committed
547
tcLcStmt _ _ stmt _ _
548
  = pprPanic "tcLcStmt: unexpected Stmt" (ppr stmt)
549

550 551

---------------------------------------------------
552 553
--           Monad comprehensions
--        (supports rebindable syntax)
554
---------------------------------------------------
555

556
tcMcStmt :: TcExprStmtChecker
557

Simon Marlow's avatar
Simon Marlow committed
558
tcMcStmt _ (LastStmt body noret return_op) res_ty thing_inside
559 560 561 562
  = do  { (body', return_op')
            <- tcSyntaxOp MCompOrigin return_op [SynRho] res_ty $
               \ [a_ty] ->
               tcMonoExprNC body (mkCheckExpType a_ty)
563
        ; thing      <- thing_inside (panic "tcMcStmt: thing_inside")
Simon Marlow's avatar
Simon Marlow committed
564
        ; return (LastStmt body' noret return_op', thing) }
565

566 567 568 569 570
-- Generators for monad comprehensions ( pat <- rhs )
--
--   [ body | q <- gen ]  ->  gen :: m a
--                            q   ::   a
--
571

572
tcMcStmt ctxt (BindStmt pat rhs bind_op fail_op _) res_ty thing_inside
573
           -- (>>=) :: rhs_ty -> (pat_ty -> new_res_ty) -> res_ty
574 575 576 577 578 579 580 581 582
  = do  { ((rhs', pat', thing, new_res_ty), bind_op')
            <- tcSyntaxOp MCompOrigin bind_op
                          [SynRho, SynFun SynAny SynRho] res_ty $
               \ [rhs_ty, pat_ty, new_res_ty] ->
               do { rhs' <- tcMonoExprNC rhs (mkCheckExpType rhs_ty)
                  ; (pat', thing) <- tcPat (StmtCtxt ctxt) pat
                                           (mkCheckExpType pat_ty) $
                                     thing_inside (mkCheckExpType new_res_ty)
                  ; return (rhs', pat', thing, new_res_ty) }
583

584 585
        -- If (but only if) the pattern can fail, typecheck the 'fail' operator
        ; fail_op' <- tcMonadFailOp (MCompPatOrigin pat) pat' fail_op new_res_ty
quchen's avatar
quchen committed
586

587
        ; return (BindStmt pat' rhs' bind_op' fail_op' new_res_ty, thing) }
588 589 590 591 592

-- Boolean expressions.
--
--   [ body | stmts, expr ]  ->  expr :: m Bool
--
593
tcMcStmt _ (BodyStmt rhs then_op guard_op _) res_ty thing_inside
594
  = do  { -- Deal with rebindable syntax:
595 596 597
          --    guard_op :: test_ty -> rhs_ty
          --    then_op  :: rhs_ty -> new_res_ty -> res_ty
          -- Where test_ty is, for example, Bool
598 599 600 601 602 603 604 605 606 607
        ; ((thing, rhs', rhs_ty, guard_op'), then_op')
            <- tcSyntaxOp MCompOrigin then_op [SynRho, SynRho] res_ty $
               \ [rhs_ty, new_res_ty] ->
               do { (rhs', guard_op')
                      <- tcSyntaxOp MCompOrigin guard_op [SynAny]
                                    (mkCheckExpType rhs_ty) $
                         \ [test_ty] ->
                         tcMonoExpr rhs (mkCheckExpType test_ty)
                  ; thing <- thing_inside (mkCheckExpType new_res_ty)
                  ; return (thing, rhs', rhs_ty, guard_op') }
608
        ; return (BodyStmt rhs' then_op' guard_op' rhs_ty, thing) }
609 610 611 612 613 614 615 616

-- Grouping statements
--
--   [ body | stmts, then group by e using f ]
--     ->  e :: t
--         f :: forall a. (a -> t) -> m a -> m (m a)
--   [ body | stmts, then group using f ]
--     ->  f :: forall a. m a -> m (m a)
617 618 619

-- We type [ body | (stmts, group by e using f), ... ]
--     f <optional by> [ (a,b,c) | stmts ] >>= \(a,b,c) -> ...body....
620
--
621
-- We type the functions as follows:
622 623 624 625 626
--     f <optional by> :: m1 (a,b,c) -> m2 (a,b,c)              (ThenForm)
--                     :: m1 (a,b,c) -> m2 (n (a,b,c))          (GroupForm)
--     (>>=) :: m2 (a,b,c)     -> ((a,b,c)   -> res) -> res     (ThenForm)
--           :: m2 (n (a,b,c)) -> (n (a,b,c) -> res) -> res     (GroupForm)
--
627 628
tcMcStmt ctxt (TransStmt { trS_stmts = stmts, trS_bndrs = bindersMap
                         , trS_by = by, trS_using = using, trS_form = form
629
                         , trS_ret = return_op, trS_bind = bind_op
630
                         , trS_fmap = fmap_op }) res_ty thing_inside
631
  = do { let star_star_kind = liftedTypeKind `mkFunTy` liftedTypeKind
632 633 634
       ; m1_ty   <- newFlexiTyVarTy star_star_kind
       ; m2_ty   <- newFlexiTyVarTy star_star_kind
       ; tup_ty  <- newFlexiTyVarTy liftedTypeKind
635
       ; by_e_ty <- newFlexiTyVarTy liftedTypeKind  -- The type of the 'by' expression (if any)
636 637 638 639

         -- n_app :: Type -> Type   -- Wraps a 'ty' into '(n ty)' for GroupForm
       ; n_app <- case form of
                    ThenForm -> return (\ty -> ty)
640 641 642
                    _        -> do { n_ty <- newFlexiTyVarTy star_star_kind
                                   ; return (n_ty `mkAppTy`) }
       ; let by_arrow :: Type -> Type
643
             -- (by_arrow res) produces ((alpha->e_ty) -> res)     ('by' present)
644
             --                          or res                    ('by' absent)
645
             by_arrow = case by of
646 647
                          Nothing -> \res -> res
                          Just {} -> \res -> (alphaTy `mkFunTy` by_e_ty) `mkFunTy` res
648 649 650

             poly_arg_ty  = m1_ty `mkAppTy` alphaTy
             using_arg_ty = m1_ty `mkAppTy` tup_ty
651 652
             poly_res_ty  = m2_ty `mkAppTy` n_app alphaTy
             using_res_ty = m2_ty `mkAppTy` n_app tup_ty
Simon Peyton Jones's avatar
Simon Peyton Jones committed
653
             using_poly_ty = mkInvForAllTy alphaTyVar $
654
                             by_arrow $
655 656
                             poly_arg_ty `mkFunTy` poly_res_ty

657 658 659
             -- 'stmts' returns a result of type (m1_ty tuple_ty),
             -- typically something like [(Int,Bool,Int)]
             -- We don't know what tuple_ty is yet, so we use a variable
660 661
       ; let (bndr_names, n_bndr_names) = unzip bindersMap
       ; (stmts', (bndr_ids, by', return_op')) <-
662 663
            tcStmtsAndThen (TransStmtCtxt ctxt) tcMcStmt stmts
                           (mkCheckExpType using_arg_ty) $ \res_ty' -> do
664
                { by' <- case by of
665
                           Nothing -> return Nothing
666 667 668
                           Just e  -> do { e' <- tcMonoExpr e
                                                   (mkCheckExpType by_e_ty)
                                         ; return (Just e') }
669 670 671 672 673 674

                -- Find the Ids (and hence types) of all old binders
                ; bndr_ids <- tcLookupLocalIds bndr_names

                -- 'return' is only used for the binders, so we know its type.
                --   return :: (a,b,c,..) -> m (a,b,c,..)
675 676 677
                ; (_, return_op') <- tcSyntaxOp MCompOrigin return_op
                                       [synKnownType (mkBigCoreVarTupTy bndr_ids)]
                                       res_ty' $ \ _ -> return ()
678

679
                ; return (bndr_ids, by', return_op') }
680

681 682 683
       --------------- Typecheck the 'bind' function -------------
       -- (>>=) :: m2 (n (a,b,c)) -> ( n (a,b,c) -> new_res_ty ) -> res_ty
       ; new_res_ty <- newFlexiTyVarTy liftedTypeKind
684 685 686 687
       ; (_, bind_op')  <- tcSyntaxOp MCompOrigin bind_op
                             [ synKnownType using_res_ty
                             , synKnownType (n_app tup_ty `mkFunTy` new_res_ty) ]
                             res_ty $ \ _ -> return ()
688 689 690

       --------------- Typecheck the 'fmap' function -------------
       ; fmap_op' <- case form of
691
                       ThenForm -> return noExpr
692
                       _ -> fmap unLoc . tcPolyExpr (noLoc fmap_op) $
Simon Peyton Jones's avatar
Simon Peyton Jones committed
693 694
                            mkInvForAllTy alphaTyVar $
                            mkInvForAllTy betaTyVar  $
695 696 697 698 699 700 701 702
                            (alphaTy `mkFunTy` betaTy)
                            `mkFunTy` (n_app alphaTy)
                            `mkFunTy` (n_app betaTy)

       --------------- Typecheck the 'using' function -------------
       -- using :: ((a,b,c)->t) -> m1 (a,b,c) -> m2 (n (a,b,c))

       ; using' <- tcPolyExpr using using_poly_ty
703
       ; let final_using = fmap (mkHsWrap (WpTyApp tup_ty)) using'
704 705

       --------------- Bulding the bindersMap ----------------
706
       ; let mk_n_bndr :: Name -> TcId -> TcId
707
             mk_n_bndr n_bndr_name bndr_id = mkLocalIdOrCoVar n_bndr_name (n_app (idType bndr_id))
708 709 710

             -- Ensure that every old binder of type `b` is linked up with its
             -- new binder which should have type `n b`
711
             -- See Note [GroupStmt binder map] in HsExpr
712 713 714
             n_bndr_ids = zipWith mk_n_bndr n_bndr_names bndr_ids
             bindersMap' = bndr_ids `zip` n_bndr_ids

715
       -- Type check the thing in the environment with
716
       -- these new binders and return the result
717 718
       ; thing <- tcExtendIdEnv n_bndr_ids $
                  thing_inside (mkCheckExpType new_res_ty)
719

720 721
       ; return (TransStmt { trS_stmts = stmts', trS_bndrs = bindersMap'
                           , trS_by = by', trS_using = final_using
722
                           , trS_ret = return_op', trS_bind = bind_op'
723
                           , trS_bind_arg_ty = n_app tup_ty
724
                           , trS_fmap = fmap_op', trS_form = form }, thing) }