DsArrows.hs 44.9 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
21 22 23 24 25 26

-- 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.

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

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

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

data DsCmdEnv = DsCmdEnv {
53
        arr_id, compose_id, first_id, app_id, choice_id, loop_id :: CoreExpr
54 55
    }

56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72
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)) }
73

74 75 76
    find_meth prs std_name
      = assocDefault (mk_panic std_name) prs std_name
    mk_panic std_name = pprPanic "mkCmdEnv" (ptext (sLit "Not found:") <+> ppr std_name)
77 78 79 80 81 82 83

-- 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 ->
84
                CoreExpr -> CoreExpr -> CoreExpr
85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
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 ->
100
                CoreExpr -> CoreExpr -> CoreExpr
101 102 103 104 105 106 107 108 109
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]

110 111 112
-- 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 ->
113
                CoreExpr -> CoreExpr -> CoreExpr
114 115
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
116

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

121 122 123 124 125 126 127 128 129
-- 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)))

130 131
-- construct CoreExpr for \ (a :: a_ty, b :: b_ty) -> b
mkSndExpr :: Type -> Type -> DsM CoreExpr
132 133 134 135 136 137
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)))
138

Austin Seipp's avatar
Austin Seipp committed
139
{-
140 141
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
142
-}
143 144

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

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

ross's avatar
ross committed
152 153
coreCasePair :: Id -> Id -> Id -> CoreExpr -> CoreExpr
coreCasePair scrut_var var1 var2 body
154
  = Case (Var scrut_var) scrut_var (exprType body)
batterseapower's avatar
batterseapower committed
155
         [(DataAlt (tupleCon BoxedTuple 2), [var1, var2], body)]
156 157

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

mkCorePairExpr :: CoreExpr -> CoreExpr -> CoreExpr
mkCorePairExpr e1 e2 = mkCoreTup [e1, e2]
162 163 164

mkCoreUnitExpr :: CoreExpr
mkCoreUnitExpr = mkCoreTup []
165

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

171
        ((x1,...,xn), (s1,...(sk,())...))
172 173 174

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
175
-}
176

177 178 179 180 181 182 183 184 185 186
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)
187 188

----------------------------------------------
189
--              buildEnvStack
190
--
191
--      ((x1,...,xn),stk)
192

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

----------------------------------------------
198
--              matchEnvStack
199
--
200 201 202 203 204 205 206 207 208 209 210
--      \ ((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
211
matchEnvStack env_ids stack_id body = do
212 213
    uniqs <- newUniqueSupply
    tup_var <- newSysLocalDs (mkBigCoreVarTupTy env_ids)
214 215 216
    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))
217

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

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

----------------------------------------------
236
--              matchVarStack
237
--
238 239 240 241 242
--      case (x1, ...(xn, s)...) -> e
--      =>
--      case z0 of (x1,z1) ->
--      case zn-1 of (xn,s) ->
--      e
243 244 245 246 247 248
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)
249

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

Austin Seipp's avatar
Austin Seipp committed
254
-- Translation of arrow abstraction
255

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

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

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

285
        D; xs |-a c : stk --> t
286 287

to an expression e such that
288

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

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

297 298 299 300 301 302 303 304 305 306
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
                IdSet)          -- subset of local vars that occur free
307

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

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

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

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

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

353
    return (do_premap ids
354
              (envStackType env_ids stack_ty)
355 356 357 358
              (mkCorePairTy arrow_ty arg_ty)
              res_ty
              core_make_pair
              (do_app ids arg_ty res_ty),
359
            (exprFreeIds core_arrow `unionVarSet` exprFreeIds core_arg)
360
              `intersectVarSet` local_vars)
361

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

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

ross's avatar
ross committed
384
    -- match the environment and stack against the input
385
    core_map <- matchEnvStack env_ids stack_id core_body
386
    return (do_premap ids
387 388
                      (envStackType env_ids stack_ty)
                      (envStackType env_ids' stack_ty')
389 390 391
                      res_ty
                      core_map
                      core_cmd,
392 393
            free_vars `unionVarSet`
              (exprFreeIds core_arg `intersectVarSet` local_vars))
ross's avatar
ross committed
394

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

401
dsCmd ids local_vars stack_ty res_ty
Alan Zimmerman's avatar
Alan Zimmerman committed
402 403
        (HsCmdLam (MG { mg_alts = [L _ (Match _ pats _
                                       (GRHSs [L _ (GRHS [] body)] _ ))] }))
404
        env_ids = do
405 406
    let
        pat_vars = mkVarSet (collectPatsBinders pats)
407
        local_vars' = pat_vars `unionVarSet` local_vars
408
        (pat_tys, stack_ty') = splitTypeAt (length pats) stack_ty
409 410 411
    (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'
412 413 414 415 416

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

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

422
    fail_expr <- mkFailExpr LambdaExpr in_ty'
423 424 425 426
    -- 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
427
    -- match the old environment and stack against the input
428
    select_code <- matchEnvStack env_ids stack_id param_code
429
    return (do_premap ids in_ty in_ty' res_ty select_code core_body,
430
            free_vars `minusVarSet` pat_vars)
431

432 433
dsCmd ids local_vars stack_ty res_ty (HsCmdPar cmd) env_ids
  = dsLCmd ids local_vars stack_ty res_ty cmd env_ids
434

435 436 437 438 439
-- 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
440
--
441 442 443
--              ---> premap (\ ((xs),stk) ->
--                       if e then Left ((xs1),stk) else Right ((xs2),stk))
--                     (c1 ||| c2)
ross's avatar
ross committed
444

445
dsCmd ids local_vars stack_ty res_ty (HsCmdIf mb_fun cond then_cmd else_cmd)
446
        env_ids = do
447
    core_cond <- dsLExpr cond
448 449 450
    (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
451 452 453
    either_con <- dsLookupTyCon eitherTyConName
    left_con   <- dsLookupDataCon leftDataConName
    right_con  <- dsLookupDataCon rightDataConName
454

455 456
    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]
457

458 459 460
        in_ty = envStackType env_ids stack_ty
        then_ty = envStackType then_ids stack_ty
        else_ty = envStackType else_ids stack_ty
461
        sum_ty = mkTyConApp either_con [then_ty, else_ty]
462
        fvs_cond = exprFreeIds core_cond `intersectVarSet` local_vars
463

464 465
        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)
466

467
    core_if <- case mb_fun of
468
       Just fun -> do { core_fun <- dsExpr fun
469
                      ; matchEnvStack env_ids stack_id $
470
                        mkCoreApps core_fun [core_cond, core_left, core_right] }
471
       Nothing  -> matchEnvStack env_ids stack_id $
472 473
                   mkIfThenElse core_cond core_left core_right

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

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

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

is translated to

487 488 489 490 491
        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
492 493 494 495

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.
496
To build all this, we use triples describing segments of the list of
ross's avatar
ross committed
497
case bodies, containing the following fields:
498
 * a list of expressions of the form (Left|Right)* ((xs)*ts), to be put
ross's avatar
ross committed
499
   into the case replacing the commands
500
 * a sum type that is the common type of these expressions, and also the
ross's avatar
ross committed
501
   input type of the arrow
502
 * a CoreExpr for an arrow built by combining the translated command
ross's avatar
ross committed
503
   bodies with |||.
Austin Seipp's avatar
Austin Seipp committed
504
-}
ross's avatar
ross committed
505

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

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

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

523 524 525 526
    branches <- mapM make_branch leaves
    either_con <- dsLookupTyCon eitherTyConName
    left_con <- dsLookupDataCon leftDataConName
    right_con <- dsLookupDataCon rightDataConName
527
    let
528 529 530 531
        left_id  = HsVar (dataConWrapId left_con)
        right_id = HsVar (dataConWrapId right_con)
        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
532

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

536 537 538
        merge_branches (builds1, in_ty1, core_exp1)
                       (builds2, in_ty2, core_exp2)
          = (map (left_expr in_ty1 in_ty2) builds1 ++
539 540 541
                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)
542
        (leaves', sum_ty, core_choices) = foldb merge_branches branches
543 544 545 546 547

        -- 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
548
        in_ty = envStackType env_ids stack_ty
549

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

555
    core_matches <- matchEnvStack env_ids stack_id core_body
556
    return (do_premap ids in_ty sum_ty res_ty core_matches core_choices,
557
            exprFreeIds core_body  `intersectVarSet` local_vars)
558

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

565
dsCmd ids local_vars stack_ty res_ty (HsCmdLet binds body) env_ids = do
566
    let
567
        defined_vars = mkVarSet (collectLocalBinders binds)
568
        local_vars' = defined_vars `unionVarSet` local_vars
569

570 571
    (core_body, _free_vars, env_ids') <- dsfixCmd ids local_vars' stack_ty res_ty body
    stack_id <- newSysLocalDs stack_ty
572
    -- build a new environment, plus the stack, using the let bindings
573
    core_binds <- dsLocalBinds binds (buildEnvStack env_ids' stack_id)
574
    -- match the old environment and stack against the input
575
    core_map <- matchEnvStack env_ids stack_id core_binds
576
    return (do_premap ids
577 578
                        (envStackType env_ids stack_ty)
                        (envStackType env_ids' stack_ty)
579 580 581
                        res_ty
                        core_map
                        core_body,
582
        exprFreeIds core_binds `intersectVarSet` local_vars)
583

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

590 591 592 593 594
dsCmd ids local_vars stack_ty res_ty (HsCmdDo stmts _) env_ids = do
    (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
595 596 597 598 599 600
                (mkCorePairTy env_ty stack_ty)
                env_ty
                res_ty
                core_fst
                core_stmts,
        env_ids')
601 602 603 604

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

dsCmd _ids local_vars _stack_ty _res_ty (HsCmdArrForm op _ args) env_ids = do
608 609 610 611 612
    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,
            unionVarSets fv_sets)
613

614 615
dsCmd ids local_vars stack_ty res_ty (HsCmdCast coercion cmd) env_ids = do
    (core_cmd, env_ids') <- dsCmd ids local_vars stack_ty res_ty cmd env_ids
Joachim Breitner's avatar
Joachim Breitner committed
616
    wrapped_cmd <- dsHsWrapper (mkWpCast coercion) core_cmd
617
    return (wrapped_cmd, env_ids')
618

Ian Lynagh's avatar
Ian Lynagh committed
619 620
dsCmd _ _ _ _ _ c = pprPanic "dsCmd" (ppr c)

621
-- D; ys |-a c : stk --> t      (ys <= xs)
622
-- ---------------------
623
-- D; xs |-a c : stk --> t      ---> premap (\ ((xs),stk) -> ((ys),stk)) c
624 625

dsTrimCmdArg
626 627 628 629 630
        :: 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
                IdSet)          -- subset of local vars that occur free
631
dsTrimCmdArg local_vars env_ids (L _ (HsCmdTop cmd stack_ty cmd_ty ids)) = do
632
    (meth_binds, meth_ids) <- mkCmdEnv ids
633 634 635
    (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)
636
    let
637 638
        in_ty = envStackType env_ids stack_ty
        in_ty' = envStackType env_ids' stack_ty
639
        arg_code = if env_ids' == env_ids then core_cmd else
640
                do_premap meth_ids in_ty in_ty' cmd_ty trim_code core_cmd
641
    return (mkLets meth_binds arg_code, free_vars)
642

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

dsfixCmd
647 648 649 650 651 652 653 654
        :: 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
                IdSet,          -- subset of local vars that occur free
                [Id])           -- the same local vars as a list, fed back
655 656
dsfixCmd ids local_vars stk_ty cmd_ty cmd
  = trimInput (dsLCmd ids local_vars stk_ty cmd_ty cmd)
657 658 659 660 661

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

trimInput
662 663 664 665 666 667
        :: ([Id] -> DsM (CoreExpr, IdSet))
        -> DsM (CoreExpr,       -- desugared expression
                IdSet,          -- subset of local vars that occur free
                [Id])           -- same local vars as a list, fed back to
                                -- the inner function to form the tuple of
                                -- inputs to the arrow.
668 669 670
trimInput build_arrow
  = fixDs (\ ~(_,_,env_ids) -> do
        (core_cmd, free_vars) <- build_arrow env_ids
671
        return (core_cmd, free_vars, varSetElems free_vars))
672

Austin Seipp's avatar
Austin Seipp committed
673
{-
674 675
Translation of command judgements of the form

676
        D |-a do { ss } : t
Austin Seipp's avatar
Austin Seipp committed
677
-}
678

679 680 681 682 683 684 685 686 687
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
                IdSet)          -- subset of local vars that occur free
688

689
dsCmdDo _ _ _ [] _ = panic "dsCmdDo"
690

691 692 693 694
-- D; xs |-a c : () --> t
-- --------------------------
-- D; xs |-a do { c } : t
--
695
--              ---> premap (\ (xs) -> ((xs), ())) c
696 697 698 699 700 701 702 703

dsCmdDo ids local_vars res_ty [L _ (LastStmt body _)] env_ids = do
    (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
704
                        (mkCorePairTy env_ty unitTy)
705 706 707
                        res_ty
                        core_map
                        core_body,
708
        env_ids')
709

710
dsCmdDo ids local_vars res_ty (stmt:stmts) env_ids = do
711
    let
712
        bound_vars = mkVarSet (collectLStmtBinders stmt)
713 714 715
        local_vars' = bound_vars `unionVarSet` local_vars
    (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
716 717 718 719 720 721 722
    return (do_compose ids
                (mkBigCoreVarTupTy env_ids)
                (mkBigCoreVarTupTy env_ids')
                res_ty
                core_stmt
                core_stmts,
              fv_stmt)
723

Austin Seipp's avatar
Austin Seipp committed
724
{-
ross's avatar
ross committed
725 726 727
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
728 729
-}

730
dsCmdLStmt :: DsCmdEnv -> IdSet -> [Id] -> CmdLStmt Id -> [Id]
Ian Lynagh's avatar
Ian Lynagh committed
731
           -> DsM (CoreExpr, IdSet)
732 733
dsCmdLStmt ids local_vars out_ids cmd env_ids
  = dsCmdStmt ids local_vars out_ids (unLoc cmd) env_ids
ross's avatar
ross committed
734

735
dsCmdStmt
736 737 738 739 740 741 742 743 744
        :: 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
                IdSet)          -- subset of local vars that occur free
745

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

754
dsCmdStmt ids local_vars out_ids (BodyStmt cmd _ _ c_ty) env_ids = do
755 756 757
    (core_cmd, fv_cmd, env_ids1) <- dsfixCmd ids local_vars unitTy c_ty cmd
    core_mux <- matchEnv env_ids
        (mkCorePairExpr
758 759
            (mkCorePairExpr (mkBigCoreVarTup env_ids1) mkCoreUnitExpr)
            (mkBigCoreVarTup out_ids))
760
    let
761 762 763 764 765
        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
766
    snd_fn <- mkSndExpr c_ty out_ty
767
    return (do_premap ids in_ty before_c_ty out_ty core_mux $
768 769 770 771
                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,
              extendVarSetList fv_cmd out_ids)
772

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

784
dsCmdStmt ids local_vars out_ids (BindStmt pat cmd _ _) env_ids = do
785
    (core_cmd, fv_cmd, env_ids1) <- dsfixCmd ids local_vars unitTy (hsLPatType pat) cmd
786
    let
787 788 789 790
        pat_ty = hsLPatType pat
        pat_vars = mkVarSet (collectPatBinders pat)
        env_ids2 = varSetElems (mkVarSet out_ids `minusVarSet` pat_vars)
        env_ty2 = mkBigCoreVarTupTy env_ids2
791 792

    -- multiplexing function
793
    --          \ (xs) -> (((xs1),()),(xs2))
794

795 796
    core_mux <- matchEnv env_ids
        (mkCorePairExpr
797 798
            (mkCorePairExpr (mkBigCoreVarTup env_ids1) mkCoreUnitExpr)
            (mkBigCoreVarTup env_ids2))
799 800

    -- projection function
801
    --          \ (p, (xs2)) -> (zs)
802

803 804
    env_id <- newSysLocalDs env_ty2
    uniqs <- newUniqueSupply
805
    let
806 807 808 809
        after_c_ty = mkCorePairTy pat_ty env_ty2
        out_ty = mkBigCoreVarTupTy out_ids
        body_expr = coreCaseTuple uniqs env_id env_ids2 (mkBigCoreVarTup out_ids)

810 811 812 813
    fail_expr <- mkFailExpr (StmtCtxt DoExpr) out_ty
    pat_id    <- selectSimpleMatchVarL pat
    match_code <- matchSimply (Var pat_id) (StmtCtxt DoExpr) pat body_expr fail_expr
    pair_id   <- newSysLocalDs after_c_ty
814
    let
815
        proj_expr = Lam pair_id (coreCasePair pair_id pat_id env_id match_code)
816

ross's avatar
ross committed
817
    -- put it all together
818
    let
819 820 821 822
        in_ty = mkBigCoreVarTupTy env_ids
        in_ty1 = mkCorePairTy (mkBigCoreVarTupTy env_ids1) unitTy
        in_ty2 = mkBigCoreVarTupTy env_ids2
        before_c_ty = mkCorePairTy in_ty1 in_ty2
823
    return (do_premap ids in_ty before_c_ty out_ty core_mux $
824 825 826 827
                do_compose ids before_c_ty after_c_ty out_ty
                        (do_first ids in_ty1 pat_ty in_ty2 core_cmd) $
                do_arr ids after_c_ty out_ty proj_expr,
              fv_cmd `unionVarSet` (mkVarSet out_ids `minusVarSet` pat_vars))
828

829 830 831
-- D; xs' |-a do { ss } : t
-- --------------------------------------
-- D; xs  |-a do { let binds; ss } : t
832
--
833
--              ---> arr (\ (xs) -> let binds in (xs')) >>> ss
834

835
dsCmdStmt ids local_vars out_ids (LetStmt binds) env_ids = do
836
    -- build a new environment using the let bindings
837
    core_binds <- dsLocalBinds binds (mkBigCoreVarTup out_ids)
838
    -- match the old environment against the input
839
    core_map <- matchEnv env_ids core_binds
840
    return (do_arr ids
841 842 843 844
                        (mkBigCoreVarTupTy env_ids)
                        (mkBigCoreVarTupTy out_