DsArrows.hs 44.6 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 )
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
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
    find_meth prs std_name
      = assocDefault (mk_panic std_name) prs std_name
76
    mk_panic std_name = pprPanic "mkCmdEnv" (text "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)
155
         [(DataAlt (tupleDataCon Boxed 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
    let 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
402 403
        (HsCmdLam (MG { mg_alts = L _ [L _ (Match _ pats _
                                           (GRHSs [L _ (GRHS [] body)] _ ))] }))
404
        env_ids = do
405
    let pat_vars = mkVarSet (collectPatsBinders pats)
406
    let
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 508
      (HsCmdCase exp (MG { mg_alts = L l matches, mg_arg_tys = arg_tys
                         , mg_origin = origin }))
509
      env_ids = do
510
    stack_id <- newSysLocalDs stack_ty
511 512 513 514 515 516

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

dsfixCmd
649 650 651 652 653 654 655 656
        :: 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
657 658
dsfixCmd ids local_vars stk_ty cmd_ty cmd
  = trimInput (dsLCmd ids local_vars stk_ty cmd_ty cmd)
659 660 661 662 663

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

trimInput
664 665 666 667 668 669
        :: ([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.
670 671 672
trimInput build_arrow
  = fixDs (\ ~(_,_,env_ids) -> do
        (core_cmd, free_vars) <- build_arrow env_ids
673
        return (core_cmd, free_vars, varSetElems free_vars))
674

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

678
        D |-a do { ss } : t
Austin Seipp's avatar
Austin Seipp committed
679
-}
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
                IdSet)          -- subset of local vars that occur free
690

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

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

Simon Marlow's avatar
Simon Marlow committed
699
dsCmdDo ids local_vars res_ty [L _ (LastStmt body _ _)] env_ids = do
700 701 702 703 704 705
    (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
706
                        (mkCorePairTy env_ty unitTy)
707 708 709
                        res_ty
                        core_map
                        core_body,
710
        env_ids')
711

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

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

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

736
dsCmdStmt
737 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
                IdSet)          -- subset of local vars that occur free
746

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

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

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

785
dsCmdStmt ids local_vars out_ids (BindStmt pat cmd _ _) env_ids = do
786 787 788
    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)
789
    let
790 791
        env_ids2 = varSetElems (mkVarSet out_ids `minusVarSet` pat_vars)
        env_ty2 = mkBigCoreVarTupTy env_ids2
792 793

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

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

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