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 )
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 50
import ListSetOps( assocDefault )
import FastString
51
import Data.List
52 53

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

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

75 76 77
    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)
78 79 80 81 82 83 84

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

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

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

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

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

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

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

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

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

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

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

mkCoreUnitExpr :: CoreExpr
mkCoreUnitExpr = mkCoreTup []
166

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

to an expression e such that
289

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

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

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
                IdSet)          -- subset of local vars that occur free
308

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

468
    core_if <- case mb_fun of
469
       Just fun -> do { core_fun <- dsExpr fun
470
                      ; matchEnvStack env_ids stack_id $
471
                        mkCoreApps core_fun [core_cond, core_left, core_right] }
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 478
                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
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,
560
            exprFreeIds core_body  `intersectVarSet` 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,
585
        exprFreeIds core_binds `intersectVarSet` 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 610

dsCmd _ids local_vars _stack_ty _res_ty (HsCmdArrForm op _ args) env_ids = do
611 612 613 614 615
    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)
616

617 618
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
619
    wrapped_cmd <- dsHsWrapper (mkWpCastN coercion) core_cmd
620
    return (wrapped_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 633
        :: 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
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 656 657
        :: 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
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
665 666 667 668 669 670
        :: ([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.
671 672 673
trimInput build_arrow
  = fixDs (\ ~(_,_,env_ids) -> do
        (core_cmd, free_vars) <- build_arrow env_ids
674
        return (core_cmd, free_vars, varSetElems 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 690
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
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
    let
715
        bound_vars = mkVarSet (collectLStmtBinders stmt)
716 717 718
        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
719 720 721 722 723 724 725
    return (do_compose ids
                (mkBigCoreVarTupTy env_ids)
                (mkBigCoreVarTupTy env_ids')
                res_ty
                core_stmt
                core_stmts,
              fv_stmt)
726

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

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

738
dsCmdStmt
739 740 741 742 743 744 745 746 747
        :: 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
748

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

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

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

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

    -- multiplexing function
796
    --          \ (xs) -> (((xs1),()),(xs2))
797

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

    -- projection function
804
    --          \ (p, (xs2)) -> (zs)
805

806 807
    env_id <- newSysLocalDs env_ty2
    uniqs <- newUniqueSupply