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, dsSyntaxExpr )
29

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

Simon Marlow's avatar
Simon Marlow committed
38
import Name
Ian Lynagh's avatar
Ian Lynagh committed
39
import Var
40
import Id
Simon Marlow's avatar
Simon Marlow committed
41 42 43 44
import DataCon
import TysWiredIn
import BasicTypes
import PrelNames
Ian Lynagh's avatar
Ian Lynagh committed
45
import Outputable
46
import Bag
Simon Marlow's avatar
Simon Marlow committed
47 48
import VarSet
import SrcLoc
49
import ListSetOps( assocDefault )
50
import Data.List
niteria's avatar
niteria committed
51
import Util
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
    find_meth prs std_name
      = assocDefault (mk_panic std_name) prs std_name
77
    mk_panic std_name = pprPanic "mkCmdEnv" (text "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
    let 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
    let pat_vars = mkVarSet (collectPatsBinders pats)
407
    let
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 470
       Just fun -> do { fun_apps <- dsSyntaxExpr fun [core_cond, core_left, core_right]
                      ; matchEnvStack env_ids stack_id fun_apps }
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
niteria's avatar
niteria committed
790
        env_ids2 = filterOut (`elemVarSet` pat_vars) out_ids
791
        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