DsArrows.hs 44.7 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

Simon Marlow's avatar
Simon Marlow committed
5 6

Desugaring arrow commands
Austin Seipp's avatar
Austin Seipp committed
7
-}
8

9
{-# LANGUAGE CPP #-}
Ian Lynagh's avatar
Ian Lynagh committed
10

11 12 13 14
module DsArrows ( dsProcExpr ) where

#include "HsVersions.h"

Simon Marlow's avatar
Simon Marlow committed
15 16
import Match
import DsUtils
17 18
import DsMonad

19
import HsSyn    hiding (collectPatBinders, collectPatsBinders, collectLStmtsBinders, collectLStmtBinders, collectStmtBinders )
Simon Marlow's avatar
Simon Marlow committed
20
import TcHsSyn
Simon Marlow's avatar
Simon Marlow committed
21
import qualified HsUtils
22 23 24 25 26 27

-- NB: The desugarer, which straddles the source and Core worlds, sometimes
--     needs to see source types (newtypes etc), and sometimes not
--     So WATCH OUT; check each use of split*Ty functions.
-- Sigh.  This is a pain.

28
import {-# SOURCE #-} DsExpr ( dsExpr, dsLExpr, dsLocalBinds, dsSyntaxExpr )
29

Simon Marlow's avatar
Simon Marlow committed
30
import TcType
31
import TcEvidence
32
import CoreSyn
Simon Marlow's avatar
Simon Marlow committed
33 34
import CoreFVs
import CoreUtils
35
import MkCore
36
import DsBinds (dsHsWrapper)
37

Simon Marlow's avatar
Simon Marlow committed
38
import Name
Ian Lynagh's avatar
Ian Lynagh committed
39
import Var
40
import Id
Simon Marlow's avatar
Simon Marlow committed
41 42 43 44
import DataCon
import TysWiredIn
import BasicTypes
import PrelNames
Ian Lynagh's avatar
Ian Lynagh committed
45
import Outputable
46
import Bag
Simon Marlow's avatar
Simon Marlow committed
47 48
import VarSet
import SrcLoc
49
import ListSetOps( assocDefault )
50
import Data.List
niteria's avatar
niteria committed
51
import Util
niteria's avatar
niteria committed
52
import UniqDFM
53 54

data DsCmdEnv = DsCmdEnv {
55
        arr_id, compose_id, first_id, app_id, choice_id, loop_id :: CoreExpr
56 57
    }

58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74
mkCmdEnv :: CmdSyntaxTable Id -> DsM ([CoreBind], DsCmdEnv)
-- See Note [CmdSyntaxTable] in HsExpr
mkCmdEnv tc_meths
  = do { (meth_binds, prs) <- mapAndUnzipM mk_bind tc_meths
       ; return (meth_binds, DsCmdEnv {
               arr_id     = Var (find_meth prs arrAName),
               compose_id = Var (find_meth prs composeAName),
               first_id   = Var (find_meth prs firstAName),
               app_id     = Var (find_meth prs appAName),
               choice_id  = Var (find_meth prs choiceAName),
               loop_id    = Var (find_meth prs loopAName)
             }) }
  where
    mk_bind (std_name, expr)
      = do { rhs <- dsExpr expr
           ; id <- newSysLocalDs (exprType rhs)
           ; return (NonRec id rhs, (std_name, id)) }
75

76 77
    find_meth prs std_name
      = assocDefault (mk_panic std_name) prs std_name
78
    mk_panic std_name = pprPanic "mkCmdEnv" (text "Not found:" <+> ppr std_name)
79 80 81 82 83 84 85

-- arr :: forall b c. (b -> c) -> a b c
do_arr :: DsCmdEnv -> Type -> Type -> CoreExpr -> CoreExpr
do_arr ids b_ty c_ty f = mkApps (arr_id ids) [Type b_ty, Type c_ty, f]

-- (>>>) :: forall b c d. a b c -> a c d -> a b d
do_compose :: DsCmdEnv -> Type -> Type -> Type ->
86
                CoreExpr -> CoreExpr -> CoreExpr
87 88 89 90 91 92 93 94 95 96 97 98 99 100 101
do_compose ids b_ty c_ty d_ty f g
  = mkApps (compose_id ids) [Type b_ty, Type c_ty, Type d_ty, f, g]

-- first :: forall b c d. a b c -> a (b,d) (c,d)
do_first :: DsCmdEnv -> Type -> Type -> Type -> CoreExpr -> CoreExpr
do_first ids b_ty c_ty d_ty f
  = mkApps (first_id ids) [Type b_ty, Type c_ty, Type d_ty, f]

-- app :: forall b c. a (a b c, b) c
do_app :: DsCmdEnv -> Type -> Type -> CoreExpr
do_app ids b_ty c_ty = mkApps (app_id ids) [Type b_ty, Type c_ty]

-- (|||) :: forall b d c. a b d -> a c d -> a (Either b c) d
-- note the swapping of d and c
do_choice :: DsCmdEnv -> Type -> Type -> Type ->
102
                CoreExpr -> CoreExpr -> CoreExpr
103 104 105 106 107 108 109 110 111
do_choice ids b_ty c_ty d_ty f g
  = mkApps (choice_id ids) [Type b_ty, Type d_ty, Type c_ty, f, g]

-- loop :: forall b d c. a (b,d) (c,d) -> a b c
-- note the swapping of d and c
do_loop :: DsCmdEnv -> Type -> Type -> Type -> CoreExpr -> CoreExpr
do_loop ids b_ty c_ty d_ty f
  = mkApps (loop_id ids) [Type b_ty, Type d_ty, Type c_ty, f]

112 113 114
-- premap :: forall b c d. (b -> c) -> a c d -> a b d
-- premap f g = arr f >>> g
do_premap :: DsCmdEnv -> Type -> Type -> Type ->
115
                CoreExpr -> CoreExpr -> CoreExpr
116 117
do_premap ids b_ty c_ty d_ty f g
   = do_compose ids b_ty c_ty d_ty (do_arr ids b_ty c_ty f) g
118

119
mkFailExpr :: HsMatchContext Id -> Type -> DsM CoreExpr
120 121 122
mkFailExpr ctxt ty
  = mkErrorAppDs pAT_ERROR_ID ty (matchContextErrString ctxt)

123 124 125 126 127 128 129 130 131
-- construct CoreExpr for \ (a :: a_ty, b :: b_ty) -> a
mkFstExpr :: Type -> Type -> DsM CoreExpr
mkFstExpr a_ty b_ty = do
    a_var <- newSysLocalDs a_ty
    b_var <- newSysLocalDs b_ty
    pair_var <- newSysLocalDs (mkCorePairTy a_ty b_ty)
    return (Lam pair_var
               (coreCasePair pair_var a_var b_var (Var a_var)))

132 133
-- construct CoreExpr for \ (a :: a_ty, b :: b_ty) -> b
mkSndExpr :: Type -> Type -> DsM CoreExpr
134 135 136 137 138 139
mkSndExpr a_ty b_ty = do
    a_var <- newSysLocalDs a_ty
    b_var <- newSysLocalDs b_ty
    pair_var <- newSysLocalDs (mkCorePairTy a_ty b_ty)
    return (Lam pair_var
               (coreCasePair pair_var a_var b_var (Var b_var)))
140

Austin Seipp's avatar
Austin Seipp committed
141
{-
142 143
Build case analysis of a tuple.  This cannot be done in the DsM monad,
because the list of variables is typically not yet defined.
Austin Seipp's avatar
Austin Seipp committed
144
-}
145 146

-- coreCaseTuple [u1..] v [x1..xn] body
147
--      = case v of v { (x1, .., xn) -> body }
148 149 150
-- But the matching may be nested if the tuple is very big

coreCaseTuple :: UniqSupply -> Id -> [Id] -> CoreExpr -> CoreExpr
ross's avatar
ross committed
151 152
coreCaseTuple uniqs scrut_var vars body
  = mkTupleCase uniqs vars body scrut_var (Var scrut_var)
153

ross's avatar
ross committed
154 155
coreCasePair :: Id -> Id -> Id -> CoreExpr -> CoreExpr
coreCasePair scrut_var var1 var2 body
156
  = Case (Var scrut_var) scrut_var (exprType body)
157
         [(DataAlt (tupleDataCon Boxed 2), [var1, var2], body)]
158 159

mkCorePairTy :: Type -> Type -> Type
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
160
mkCorePairTy t1 t2 = mkBoxedTupleTy [t1, t2]
161 162 163

mkCorePairExpr :: CoreExpr -> CoreExpr -> CoreExpr
mkCorePairExpr e1 e2 = mkCoreTup [e1, e2]
164 165 166

mkCoreUnitExpr :: CoreExpr
mkCoreUnitExpr = mkCoreTup []
167

Austin Seipp's avatar
Austin Seipp committed
168
{-
169
The input is divided into a local environment, which is a flat tuple
170 171
(unless it's too big), and a stack, which is a right-nested pair.
In general, the input has the form
172

173
        ((x1,...,xn), (s1,...(sk,())...))
174 175 176

where xi are the environment values, and si the ones on the stack,
with s1 being the "top", the first one to be matched with a lambda.
Austin Seipp's avatar
Austin Seipp committed
177
-}
178

179 180 181 182 183 184 185 186 187 188
envStackType :: [Id] -> Type -> Type
envStackType ids stack_ty = mkCorePairTy (mkBigCoreVarTupTy ids) stack_ty

-- splitTypeAt n (t1,... (tn,t)...) = ([t1, ..., tn], t)
splitTypeAt :: Int -> Type -> ([Type], Type)
splitTypeAt n ty
  | n == 0 = ([], ty)
  | otherwise = case tcTyConAppArgs ty of
      [t, ty'] -> let (ts, ty_r) = splitTypeAt (n-1) ty' in (t:ts, ty_r)
      _ -> pprPanic "splitTypeAt" (ppr ty)
189 190

----------------------------------------------
191
--              buildEnvStack
192
--
193
--      ((x1,...,xn),stk)
194

195 196 197
buildEnvStack :: [Id] -> Id -> CoreExpr
buildEnvStack env_ids stack_id
  = mkCorePairExpr (mkBigCoreVarTup env_ids) (Var stack_id)
198 199

----------------------------------------------
200
--              matchEnvStack
201
--
202 203 204 205 206 207 208 209 210 211 212
--      \ ((x1,...,xn),stk) -> body
--      =>
--      \ pair ->
--      case pair of (tup,stk) ->
--      case tup of (x1,...,xn) ->
--      body

matchEnvStack   :: [Id]         -- x1..xn
                -> Id           -- stk
                -> CoreExpr     -- e
                -> DsM CoreExpr
213
matchEnvStack env_ids stack_id body = do
214 215
    uniqs <- newUniqueSupply
    tup_var <- newSysLocalDs (mkBigCoreVarTupTy env_ids)
216 217 218
    let match_env = coreCaseTuple uniqs tup_var env_ids body
    pair_id <- newSysLocalDs (mkCorePairTy (idType tup_var) (idType stack_id))
    return (Lam pair_id (coreCasePair pair_id tup_var stack_id match_env))
219

220
----------------------------------------------
221
--              matchEnv
222
--
223 224 225 226 227 228 229 230 231
--      \ (x1,...,xn) -> body
--      =>
--      \ tup ->
--      case tup of (x1,...,xn) ->
--      body

matchEnv :: [Id]        -- x1..xn
         -> CoreExpr    -- e
         -> DsM CoreExpr
232 233 234 235
matchEnv env_ids body = do
    uniqs <- newUniqueSupply
    tup_id <- newSysLocalDs (mkBigCoreVarTupTy env_ids)
    return (Lam tup_id (coreCaseTuple uniqs tup_id env_ids body))
236 237

----------------------------------------------
238
--              matchVarStack
239
--
240 241 242 243 244
--      case (x1, ...(xn, s)...) -> e
--      =>
--      case z0 of (x1,z1) ->
--      case zn-1 of (xn,s) ->
--      e
245 246 247 248 249 250
matchVarStack :: [Id] -> Id -> CoreExpr -> DsM (Id, CoreExpr)
matchVarStack [] stack_id body = return (stack_id, body)
matchVarStack (param_id:param_ids) stack_id body = do
    (tail_id, tail_code) <- matchVarStack param_ids stack_id body
    pair_id <- newSysLocalDs (mkCorePairTy (idType param_id) (idType tail_id))
    return (pair_id, coreCasePair pair_id param_id tail_id tail_code)
251

252 253 254
mkHsEnvStackExpr :: [Id] -> Id -> LHsExpr Id
mkHsEnvStackExpr env_ids stack_id
  = mkLHsTupleExpr [mkLHsVarTuple env_ids, nlHsVar stack_id]
255

Austin Seipp's avatar
Austin Seipp committed
256
-- Translation of arrow abstraction
257

258
-- D; xs |-a c : () --> t'      ---> c'
259
-- --------------------------
260
-- D |- proc p -> c :: a t t'   ---> premap (\ p -> ((xs),())) c'
261
--
262
--              where (xs) is the tuple of variables bound by p
263 264

dsProcExpr
265 266 267
        :: LPat Id
        -> LHsCmdTop Id
        -> DsM CoreExpr
268
dsProcExpr pat (L _ (HsCmdTop cmd _unitTy cmd_ty ids)) = do
269
    (meth_binds, meth_ids) <- mkCmdEnv ids
270
    let locals = mkVarSet (collectPatBinders pat)
271
    (core_cmd, _free_vars, env_ids) <- dsfixCmd meth_ids locals unitTy cmd_ty cmd
272
    let env_ty = mkBigCoreVarTupTy env_ids
273 274 275
    let env_stk_ty = mkCorePairTy env_ty unitTy
    let env_stk_expr = mkCorePairExpr (mkBigCoreVarTup env_ids) mkCoreUnitExpr
    fail_expr <- mkFailExpr ProcExpr env_stk_ty
276
    var <- selectSimpleMatchVarL pat
277
    match_code <- matchSimply (Var var) ProcExpr pat env_stk_expr fail_expr
278
    let pat_ty = hsLPatType pat
279
    let proc_code = do_premap meth_ids pat_ty env_stk_ty cmd_ty
280 281
                    (Lam var match_code)
                    core_cmd
282
    return (mkLets meth_binds proc_code)
283

Austin Seipp's avatar
Austin Seipp committed
284
{-
285 286
Translation of a command judgement of the form

287
        D; xs |-a c : stk --> t
288 289

to an expression e such that
290

291
        D |- e :: a (xs, stk) t
Austin Seipp's avatar
Austin Seipp committed
292
-}
293

294
dsLCmd :: DsCmdEnv -> IdSet -> Type -> Type -> LHsCmd Id -> [Id]
niteria's avatar
niteria committed
295
       -> DsM (CoreExpr, DIdSet)
296 297
dsLCmd ids local_vars stk_ty res_ty cmd env_ids
  = dsCmd ids local_vars stk_ty res_ty (unLoc cmd) env_ids
298

299 300 301 302 303 304 305 306 307
dsCmd   :: DsCmdEnv             -- arrow combinators
        -> IdSet                -- set of local vars available to this command
        -> Type                 -- type of the stack (right-nested tuple)
        -> Type                 -- return type of the command
        -> HsCmd Id             -- command to desugar
        -> [Id]                 -- list of vars in the input to this command
                                -- This is typically fed back,
                                -- so don't pull on it too early
        -> DsM (CoreExpr,       -- desugared expression
niteria's avatar
niteria committed
308
                DIdSet)         -- subset of local vars that occur free
309

310 311 312 313
-- D |- fun :: a t1 t2
-- D, xs |- arg :: t1
-- -----------------------------
-- D; xs |-a fun -< arg : stk --> t2
ross's avatar
ross committed
314
--
315
--              ---> premap (\ ((xs), _stk) -> arg) fun
316

317
dsCmd ids local_vars stack_ty res_ty
318
        (HsCmdArrApp arrow arg arrow_ty HsFirstOrderApp _)
319
        env_ids = do
320 321
    let
        (a_arg_ty, _res_ty') = tcSplitAppTy arrow_ty
ross's avatar
ross committed
322
        (_a_ty, arg_ty) = tcSplitAppTy a_arg_ty
323 324
    core_arrow <- dsLExpr arrow
    core_arg   <- dsLExpr arg
325 326
    stack_id   <- newSysLocalDs stack_ty
    core_make_arg <- matchEnvStack env_ids stack_id core_arg
327
    return (do_premap ids
328
              (envStackType env_ids stack_ty)
329 330 331 332
              arg_ty
              res_ty
              core_make_arg
              core_arrow,
niteria's avatar
niteria committed
333
            exprFreeIdsDSet core_arg `udfmIntersectUFM` local_vars)
334

335 336 337 338
-- D, xs |- fun :: a t1 t2
-- D, xs |- arg :: t1
-- ------------------------------
-- D; xs |-a fun -<< arg : stk --> t2
ross's avatar
ross committed
339
--
340
--              ---> premap (\ ((xs), _stk) -> (fun, arg)) app
341

342
dsCmd ids local_vars stack_ty res_ty
343
        (HsCmdArrApp arrow arg arrow_ty HsHigherOrderApp _)
344
        env_ids = do
345 346
    let
        (a_arg_ty, _res_ty') = tcSplitAppTy arrow_ty
ross's avatar
ross committed
347
        (_a_ty, arg_ty) = tcSplitAppTy a_arg_ty
348

349 350
    core_arrow <- dsLExpr arrow
    core_arg   <- dsLExpr arg
351 352 353 354
    stack_id   <- newSysLocalDs stack_ty
    core_make_pair <- matchEnvStack env_ids stack_id
          (mkCorePairExpr core_arrow core_arg)

355
    return (do_premap ids
356
              (envStackType env_ids stack_ty)
357 358 359 360
              (mkCorePairTy arrow_ty arg_ty)
              res_ty
              core_make_pair
              (do_app ids arg_ty res_ty),
niteria's avatar
niteria committed
361 362
            (exprsFreeIdsDSet [core_arrow, core_arg])
              `udfmIntersectUFM` local_vars)
363

364 365 366 367
-- D; ys |-a cmd : (t,stk) --> t'
-- D, xs |-  exp :: t
-- ------------------------
-- D; xs |-a cmd exp : stk --> t'
ross's avatar
ross committed
368
--
369
--              ---> premap (\ ((xs),stk) -> ((ys),(e,stk))) cmd
ross's avatar
ross committed
370

371
dsCmd ids local_vars stack_ty res_ty (HsCmdApp cmd arg) env_ids = do
372
    core_arg <- dsLExpr arg
ross's avatar
ross committed
373
    let
374
        arg_ty = exprType core_arg
375
        stack_ty' = mkCorePairTy arg_ty stack_ty
376
    (core_cmd, free_vars, env_ids')
377 378
             <- dsfixCmd ids local_vars stack_ty' res_ty cmd
    stack_id <- newSysLocalDs stack_ty
379
    arg_id <- newSysLocalDs arg_ty
ross's avatar
ross committed
380 381
    -- push the argument expression onto the stack
    let
382
        stack' = mkCorePairExpr (Var arg_id) (Var stack_id)
383
        core_body = bindNonRec arg_id core_arg
384
                        (mkCorePairExpr (mkBigCoreVarTup env_ids') stack')
385

ross's avatar
ross committed
386
    -- match the environment and stack against the input
387
    core_map <- matchEnvStack env_ids stack_id core_body
388
    return (do_premap ids
389 390
                      (envStackType env_ids stack_ty)
                      (envStackType env_ids' stack_ty')
391 392 393
                      res_ty
                      core_map
                      core_cmd,
niteria's avatar
niteria committed
394 395
            free_vars `unionDVarSet`
              (exprFreeIdsDSet core_arg `udfmIntersectUFM` local_vars))
ross's avatar
ross committed
396

397 398 399
-- D; ys |-a cmd : stk t'
-- -----------------------------------------------
-- D; xs |-a \ p1 ... pk -> cmd : (t1,...(tk,stk)...) t'
400
--
401
--              ---> premap (\ ((xs), (p1, ... (pk,stk)...)) -> ((ys),stk)) cmd
402

403
dsCmd ids local_vars stack_ty res_ty
404 405
        (HsCmdLam (MG { mg_alts = L _ [L _ (Match _ pats _
                                           (GRHSs [L _ (GRHS [] body)] _ ))] }))
406
        env_ids = do
407
    let pat_vars = mkVarSet (collectPatsBinders pats)
408
    let
409
        local_vars' = pat_vars `unionVarSet` local_vars
410
        (pat_tys, stack_ty') = splitTypeAt (length pats) stack_ty
411 412 413
    (core_body, free_vars, env_ids') <- dsfixCmd ids local_vars' stack_ty' res_ty body
    param_ids <- mapM newSysLocalDs pat_tys
    stack_id' <- newSysLocalDs stack_ty'
414 415 416 417 418

    -- the expression is built from the inside out, so the actions
    -- are presented in reverse order

    let
419
        -- build a new environment, plus what's left of the stack
420 421 422
        core_expr = buildEnvStack env_ids' stack_id'
        in_ty = envStackType env_ids stack_ty
        in_ty' = envStackType env_ids' stack_ty'
423

424
    fail_expr <- mkFailExpr LambdaExpr in_ty'
425 426 427 428
    -- match the patterns against the parameters
    match_code <- matchSimplys (map Var param_ids) LambdaExpr pats core_expr fail_expr
    -- match the parameters against the top of the old stack
    (stack_id, param_code) <- matchVarStack param_ids stack_id' match_code
429
    -- match the old environment and stack against the input
430
    select_code <- matchEnvStack env_ids stack_id param_code
431
    return (do_premap ids in_ty in_ty' res_ty select_code core_body,
niteria's avatar
niteria committed
432
            free_vars `udfmMinusUFM` pat_vars)
433

434 435
dsCmd ids local_vars stack_ty res_ty (HsCmdPar cmd) env_ids
  = dsLCmd ids local_vars stack_ty res_ty cmd env_ids
436

437 438 439 440 441
-- D, xs |- e :: Bool
-- D; xs1 |-a c1 : stk --> t
-- D; xs2 |-a c2 : stk --> t
-- ----------------------------------------
-- D; xs |-a if e then c1 else c2 : stk --> t
ross's avatar
ross committed
442
--
443 444 445
--              ---> premap (\ ((xs),stk) ->
--                       if e then Left ((xs1),stk) else Right ((xs2),stk))
--                     (c1 ||| c2)
ross's avatar
ross committed
446

447
dsCmd ids local_vars stack_ty res_ty (HsCmdIf mb_fun cond then_cmd else_cmd)
448
        env_ids = do
449
    core_cond <- dsLExpr cond
450 451 452
    (core_then, fvs_then, then_ids) <- dsfixCmd ids local_vars stack_ty res_ty then_cmd
    (core_else, fvs_else, else_ids) <- dsfixCmd ids local_vars stack_ty res_ty else_cmd
    stack_id   <- newSysLocalDs stack_ty
453 454 455
    either_con <- dsLookupTyCon eitherTyConName
    left_con   <- dsLookupDataCon leftDataConName
    right_con  <- dsLookupDataCon rightDataConName
456

457 458
    let mk_left_expr ty1 ty2 e = mkCoreConApps left_con   [Type ty1, Type ty2, e]
        mk_right_expr ty1 ty2 e = mkCoreConApps right_con [Type ty1, Type ty2, e]
459

460 461 462
        in_ty = envStackType env_ids stack_ty
        then_ty = envStackType then_ids stack_ty
        else_ty = envStackType else_ids stack_ty
463
        sum_ty = mkTyConApp either_con [then_ty, else_ty]
niteria's avatar
niteria committed
464
        fvs_cond = exprFreeIdsDSet core_cond `udfmIntersectUFM` local_vars
465

466 467
        core_left  = mk_left_expr  then_ty else_ty (buildEnvStack then_ids stack_id)
        core_right = mk_right_expr then_ty else_ty (buildEnvStack else_ids stack_id)
468

469
    core_if <- case mb_fun of
470 471
       Just fun -> do { fun_apps <- dsSyntaxExpr fun [core_cond, core_left, core_right]
                      ; matchEnvStack env_ids stack_id fun_apps }
472
       Nothing  -> matchEnvStack env_ids stack_id $
473 474
                   mkIfThenElse core_cond core_left core_right

475
    return (do_premap ids in_ty sum_ty res_ty
476 477
                core_if
                (do_choice ids then_ty else_ty res_ty core_then core_else),
niteria's avatar
niteria committed
478
        fvs_cond `unionDVarSet` fvs_then `unionDVarSet` fvs_else)
ross's avatar
ross committed
479

Austin Seipp's avatar
Austin Seipp committed
480
{-
ross's avatar
ross committed
481 482 483
Case commands are treated in much the same way as if commands
(see above) except that there are more alternatives.  For example

484
        case e of { p1 -> c1; p2 -> c2; p3 -> c3 }
ross's avatar
ross committed
485 486 487

is translated to

488 489 490 491 492
        premap (\ ((xs)*ts) -> case e of
                p1 -> (Left (Left (xs1)*ts))
                p2 -> Left ((Right (xs2)*ts))
                p3 -> Right ((xs3)*ts))
        ((c1 ||| c2) ||| c3)
ross's avatar
ross committed
493 494 495 496

The idea is to extract the commands from the case, build a balanced tree
of choices, and replace the commands with expressions that build tagged
tuples, obtaining a case expression that can be desugared normally.
497
To build all this, we use triples describing segments of the list of
ross's avatar
ross committed
498
case bodies, containing the following fields:
499
 * a list of expressions of the form (Left|Right)* ((xs)*ts), to be put
ross's avatar
ross committed
500
   into the case replacing the commands
501
 * a sum type that is the common type of these expressions, and also the
ross's avatar
ross committed
502
   input type of the arrow
503
 * a CoreExpr for an arrow built by combining the translated command
ross's avatar
ross committed
504
   bodies with |||.
Austin Seipp's avatar
Austin Seipp committed
505
-}
ross's avatar
ross committed
506

507
dsCmd ids local_vars stack_ty res_ty
508 509
      (HsCmdCase exp (MG { mg_alts = L l matches, mg_arg_tys = arg_tys
                         , mg_origin = origin }))
510
      env_ids = do
511
    stack_id <- newSysLocalDs stack_ty
512 513 514 515 516 517

    -- Extract and desugar the leaf commands in the case, building tuple
    -- expressions that will (after tagging) replace these leaves

    let
        leaves = concatMap leavesMatch matches
518
        make_branch (leaf, bound_vars) = do
519
            (core_leaf, _fvs, leaf_ids) <-
520 521 522
                  dsfixCmd ids (bound_vars `unionVarSet` local_vars) stack_ty res_ty leaf
            return ([mkHsEnvStackExpr leaf_ids stack_id],
                    envStackType leaf_ids stack_ty,
523
                    core_leaf)
524

525 526 527 528
    branches <- mapM make_branch leaves
    either_con <- dsLookupTyCon eitherTyConName
    left_con <- dsLookupDataCon leftDataConName
    right_con <- dsLookupDataCon rightDataConName
529
    let
530 531
        left_id  = HsVar (noLoc (dataConWrapId left_con))
        right_id = HsVar (noLoc (dataConWrapId right_con))
532 533
        left_expr  ty1 ty2 e = noLoc $ HsApp (noLoc $ HsWrap (mkWpTyApps [ty1, ty2]) left_id ) e
        right_expr ty1 ty2 e = noLoc $ HsApp (noLoc $ HsWrap (mkWpTyApps [ty1, ty2]) right_id) e
534

535 536
        -- Prefix each tuple with a distinct series of Left's and Right's,
        -- in a balanced way, keeping track of the types.
537

538 539 540
        merge_branches (builds1, in_ty1, core_exp1)
                       (builds2, in_ty2, core_exp2)
          = (map (left_expr in_ty1 in_ty2) builds1 ++
541 542 543
                map (right_expr in_ty1 in_ty2) builds2,
             mkTyConApp either_con [in_ty1, in_ty2],
             do_choice ids in_ty1 in_ty2 res_ty core_exp1 core_exp2)
544
        (leaves', sum_ty, core_choices) = foldb merge_branches branches
545 546 547 548 549

        -- Replace the commands in the case with these tagged tuples,
        -- yielding a HsExpr Id we can feed to dsExpr.

        (_, matches') = mapAccumL (replaceLeavesMatch res_ty) leaves' matches
550
        in_ty = envStackType env_ids stack_ty
551

552 553
    core_body <- dsExpr (HsCase exp (MG { mg_alts = L l matches'
                                        , mg_arg_tys = arg_tys
554
                                        , mg_res_ty = sum_ty, mg_origin = origin }))
555 556
        -- Note that we replace the HsCase result type by sum_ty,
        -- which is the type of matches'
557

558
    core_matches <- matchEnvStack env_ids stack_id core_body
559
    return (do_premap ids in_ty sum_ty res_ty core_matches core_choices,
niteria's avatar
niteria committed
560
            exprFreeIdsDSet core_body `udfmIntersectUFM` local_vars)
561

562 563 564
-- D; ys |-a cmd : stk --> t
-- ----------------------------------
-- D; xs |-a let binds in cmd : stk --> t
565
--
566
--              ---> premap (\ ((xs),stk) -> let binds in ((ys),stk)) c
567

568
dsCmd ids local_vars stack_ty res_ty (HsCmdLet (L _ binds) body) env_ids = do
569
    let
570
        defined_vars = mkVarSet (collectLocalBinders binds)
571
        local_vars' = defined_vars `unionVarSet` local_vars
572

573 574
    (core_body, _free_vars, env_ids') <- dsfixCmd ids local_vars' stack_ty res_ty body
    stack_id <- newSysLocalDs stack_ty
575
    -- build a new environment, plus the stack, using the let bindings
576
    core_binds <- dsLocalBinds binds (buildEnvStack env_ids' stack_id)
577
    -- match the old environment and stack against the input
578
    core_map <- matchEnvStack env_ids stack_id core_binds
579
    return (do_premap ids
580 581
                        (envStackType env_ids stack_ty)
                        (envStackType env_ids' stack_ty)
582 583 584
                        res_ty
                        core_map
                        core_body,
niteria's avatar
niteria committed
585
        exprFreeIdsDSet core_binds `udfmIntersectUFM` local_vars)
586

587 588 589 590
-- D; xs |-a ss : t
-- ----------------------------------
-- D; xs |-a do { ss } : () --> t
--
591
--              ---> premap (\ (env,stk) -> env) c
592

593
dsCmd ids local_vars stack_ty res_ty (HsCmdDo (L _ stmts) _) env_ids = do
594 595 596 597
    (core_stmts, env_ids') <- dsCmdDo ids local_vars res_ty stmts env_ids
    let env_ty = mkBigCoreVarTupTy env_ids
    core_fst <- mkFstExpr env_ty stack_ty
    return (do_premap ids
598 599 600 601 602 603
                (mkCorePairTy env_ty stack_ty)
                env_ty
                res_ty
                core_fst
                core_stmts,
        env_ids')
604 605 606 607

-- D |- e :: forall e. a1 (e,stk1) t1 -> ... an (e,stkn) tn -> a (e,stk) t
-- D; xs |-a ci :: stki --> ti
-- -----------------------------------
608
-- D; xs |-a (|e c1 ... cn|) :: stk --> t       ---> e [t_xs] c1 ... cn
609

Alan Zimmerman's avatar
Alan Zimmerman committed
610
dsCmd _ids local_vars _stack_ty _res_ty (HsCmdArrForm op _ _ args) env_ids = do
611 612 613 614
    let env_ty = mkBigCoreVarTupTy env_ids
    core_op <- dsLExpr op
    (core_args, fv_sets) <- mapAndUnzipM (dsTrimCmdArg local_vars env_ids) args
    return (mkApps (App core_op (Type env_ty)) core_args,
niteria's avatar
niteria committed
615
            unionDVarSets fv_sets)
616

eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
617
dsCmd ids local_vars stack_ty res_ty (HsCmdWrap wrap cmd) env_ids = do
618
    (core_cmd, env_ids') <- dsCmd ids local_vars stack_ty res_ty cmd env_ids
Simon Peyton Jones's avatar
Simon Peyton Jones committed
619 620
    core_wrap <- dsHsWrapper wrap
    return (core_wrap core_cmd, env_ids')
621

Ian Lynagh's avatar
Ian Lynagh committed
622 623
dsCmd _ _ _ _ _ c = pprPanic "dsCmd" (ppr c)

624
-- D; ys |-a c : stk --> t      (ys <= xs)
625
-- ---------------------
626
-- D; xs |-a c : stk --> t      ---> premap (\ ((xs),stk) -> ((ys),stk)) c
627 628

dsTrimCmdArg
629 630 631 632
        :: IdSet                -- set of local vars available to this command
        -> [Id]                 -- list of vars in the input to this command
        -> LHsCmdTop Id         -- command argument to desugar
        -> DsM (CoreExpr,       -- desugared expression
niteria's avatar
niteria committed
633
                DIdSet)         -- subset of local vars that occur free
634
dsTrimCmdArg local_vars env_ids (L _ (HsCmdTop cmd stack_ty cmd_ty ids)) = do
635
    (meth_binds, meth_ids) <- mkCmdEnv ids
636 637 638
    (core_cmd, free_vars, env_ids') <- dsfixCmd meth_ids local_vars stack_ty cmd_ty cmd
    stack_id <- newSysLocalDs stack_ty
    trim_code <- matchEnvStack env_ids stack_id (buildEnvStack env_ids' stack_id)
639
    let
640 641
        in_ty = envStackType env_ids stack_ty
        in_ty' = envStackType env_ids' stack_ty
642
        arg_code = if env_ids' == env_ids then core_cmd else
643
                do_premap meth_ids in_ty in_ty' cmd_ty trim_code core_cmd
644
    return (mkLets meth_binds arg_code, free_vars)
645

646 647
-- Given D; xs |-a c : stk --> t, builds c with xs fed back.
-- Typically needs to be prefixed with arr (\(p, stk) -> ((xs),stk))
648 649

dsfixCmd
650 651 652 653 654 655
        :: DsCmdEnv             -- arrow combinators
        -> IdSet                -- set of local vars available to this command
        -> Type                 -- type of the stack (right-nested tuple)
        -> Type                 -- return type of the command
        -> LHsCmd Id            -- command to desugar
        -> DsM (CoreExpr,       -- desugared expression
niteria's avatar
niteria committed
656
                DIdSet,         -- subset of local vars that occur free
657
                [Id])           -- the same local vars as a list, fed back
658 659
dsfixCmd ids local_vars stk_ty cmd_ty cmd
  = trimInput (dsLCmd ids local_vars stk_ty cmd_ty cmd)
660 661 662 663 664

-- Feed back the list of local variables actually used a command,
-- for use as the input tuple of the generated arrow.

trimInput
niteria's avatar
niteria committed
665
        :: ([Id] -> DsM (CoreExpr, DIdSet))
666
        -> DsM (CoreExpr,       -- desugared expression
niteria's avatar
niteria committed
667
                DIdSet,         -- subset of local vars that occur free
668 669 670
                [Id])           -- same local vars as a list, fed back to
                                -- the inner function to form the tuple of
                                -- inputs to the arrow.
671 672 673
trimInput build_arrow
  = fixDs (\ ~(_,_,env_ids) -> do
        (core_cmd, free_vars) <- build_arrow env_ids
niteria's avatar
niteria committed
674
        return (core_cmd, free_vars, dVarSetElems free_vars))
675

Austin Seipp's avatar
Austin Seipp committed
676
{-
677 678
Translation of command judgements of the form

679
        D |-a do { ss } : t
Austin Seipp's avatar
Austin Seipp committed
680
-}
681

682 683 684 685 686 687 688 689
dsCmdDo :: DsCmdEnv             -- arrow combinators
        -> IdSet                -- set of local vars available to this statement
        -> Type                 -- return type of the statement
        -> [CmdLStmt Id]        -- statements to desugar
        -> [Id]                 -- list of vars in the input to this statement
                                -- This is typically fed back,
                                -- so don't pull on it too early
        -> DsM (CoreExpr,       -- desugared expression
niteria's avatar
niteria committed
690
                DIdSet)         -- subset of local vars that occur free
691

692
dsCmdDo _ _ _ [] _ = panic "dsCmdDo"
693

694 695 696 697
-- D; xs |-a c : () --> t
-- --------------------------
-- D; xs |-a do { c } : t
--
698
--              ---> premap (\ (xs) -> ((xs), ())) c
699

Simon Marlow's avatar
Simon Marlow committed
700
dsCmdDo ids local_vars res_ty [L _ (LastStmt body _ _)] env_ids = do
701 702 703 704 705 706
    (core_body, env_ids') <- dsLCmd ids local_vars unitTy res_ty body env_ids
    let env_ty = mkBigCoreVarTupTy env_ids
    env_var <- newSysLocalDs env_ty
    let core_map = Lam env_var (mkCorePairExpr (Var env_var) mkCoreUnitExpr)
    return (do_premap ids
                        env_ty
707
                        (mkCorePairTy env_ty unitTy)
708 709 710
                        res_ty
                        core_map
                        core_body,
711
        env_ids')
712

713
dsCmdDo ids local_vars res_ty (stmt:stmts) env_ids = do
714 715
    let bound_vars  = mkVarSet (collectLStmtBinders stmt)
    let local_vars' = bound_vars `unionVarSet` local_vars
716 717
    (core_stmts, _, env_ids') <- trimInput (dsCmdDo ids local_vars' res_ty stmts)
    (core_stmt, fv_stmt) <- dsCmdLStmt ids local_vars env_ids' stmt env_ids
718 719 720 721 722 723 724
    return (do_compose ids
                (mkBigCoreVarTupTy env_ids)
                (mkBigCoreVarTupTy env_ids')
                res_ty
                core_stmt
                core_stmts,
              fv_stmt)
725

Austin Seipp's avatar
Austin Seipp committed
726
{-
ross's avatar
ross committed
727 728 729
A statement maps one local environment to another, and is represented
as an arrow from one tuple type to another.  A statement sequence is
translated to a composition of such arrows.
Austin Seipp's avatar
Austin Seipp committed
730 731
-}

732
dsCmdLStmt :: DsCmdEnv -> IdSet -> [Id] -> CmdLStmt Id -> [Id]
niteria's avatar
niteria committed
733
           -> DsM (CoreExpr, DIdSet)
734 735
dsCmdLStmt ids local_vars out_ids cmd env_ids
  = dsCmdStmt ids local_vars out_ids (unLoc cmd) env_ids
ross's avatar
ross committed
736

737
dsCmdStmt
738 739 740 741 742 743 744 745
        :: DsCmdEnv             -- arrow combinators
        -> IdSet                -- set of local vars available to this statement
        -> [Id]                 -- list of vars in the output of this statement
        -> CmdStmt Id           -- statement to desugar
        -> [Id]                 -- list of vars in the input to this statement
                                -- This is typically fed back,
                                -- so don't pull on it too early
        -> DsM (CoreExpr,       -- desugared expression
niteria's avatar
niteria committed
746
                DIdSet)         -- subset of local vars that occur free
747

748 749 750 751
-- D; xs1 |-a c : () --> t
-- D; xs' |-a do { ss } : t'
-- ------------------------------
-- D; xs  |-a do { c; ss } : t'
752
--
753 754
--              ---> premap (\ ((xs)) -> (((xs1),()),(xs')))
--                      (first c >>> arr snd) >>> ss
755

756
dsCmdStmt ids local_vars out_ids (BodyStmt cmd _ _ c_ty) env_ids = do
757 758 759
    (core_cmd, fv_cmd, env_ids1) <- dsfixCmd ids local_vars unitTy c_ty cmd
    core_mux <- matchEnv env_ids
        (mkCorePairExpr
760 761
            (mkCorePairExpr (mkBigCoreVarTup env_ids1) mkCoreUnitExpr)
            (mkBigCoreVarTup out_ids))
762
    let
763 764 765 766 767
        in_ty = mkBigCoreVarTupTy env_ids
        in_ty1 = mkCorePairTy (mkBigCoreVarTupTy env_ids1) unitTy
        out_ty = mkBigCoreVarTupTy out_ids
        before_c_ty = mkCorePairTy in_ty1 out_ty
        after_c_ty = mkCorePairTy c_ty out_ty
768
    snd_fn <- mkSndExpr c_ty out_ty
769
    return (do_premap ids in_ty before_c_ty out_ty core_mux $
770 771 772
                do_compose ids before_c_ty after_c_ty out_ty
                        (do_first ids in_ty1 c_ty out_ty core_cmd) $
                do_arr ids after_c_ty out_ty snd_fn,
niteria's avatar
niteria committed
773
              extendDVarSetList fv_cmd out_ids)
774

775
-- D; xs1 |-a c : () --> t
776
-- D; xs' |-a do { ss } : t'            xs2 = xs' - defs(p)
777 778
-- -----------------------------------
-- D; xs  |-a do { p <- c; ss } : t'
779
--
780 781
--              ---> premap (\ (xs) -> (((xs1),()),(xs2)))
--                      (first c >>> arr (\ (p, (xs2)) -> (xs'))) >>> ss
782 783 784 785
--
-- It would be simpler and more consistent to do this using second,
-- but that's likely to be defined in terms of first.

786
dsCmdStmt ids local_vars out_ids (BindStmt pat cmd _ _ _) env_ids = do
787 788 789
    let pat_ty = hsLPatType pat
    (core_cmd, fv_cmd, env_ids1) <- dsfixCmd ids local_vars unitTy pat_ty cmd
    let pat_vars = mkVarSet (collectPatBinders pat)
790
    let
niteria's avatar
niteria committed
791
        env_ids2 = filterOut (`elemVarSet` pat_vars) out_ids
792
        env_ty2 = mkBigCoreVarTupTy env_ids2
simonpj's avatar