DsArrows.lhs 42.9 KB
Newer Older
1
%
Simon Marlow's avatar
Simon Marlow committed
2
% (c) The University of Glasgow 2006
3 4
% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
%
Simon Marlow's avatar
Simon Marlow committed
5 6

Desugaring arrow commands
7 8

\begin{code}
Ian Lynagh's avatar
Ian Lynagh committed
9 10 11 12
{-# OPTIONS -fno-warn-tabs #-}
-- The above warning supression flag is a temporary kludge.
-- While working on this module you are encouraged to remove it and
-- detab the module (please do the detabbing in a separate patch). See
13
--     http://ghc.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#TabsvsSpaces
Ian Lynagh's avatar
Ian Lynagh committed
14 15
-- for details

16 17 18 19
module DsArrows ( dsProcExpr ) where

#include "HsVersions.h"

Simon Marlow's avatar
Simon Marlow committed
20 21
import Match
import DsUtils
22 23
import DsMonad

24
import HsSyn	hiding (collectPatBinders, collectPatsBinders, collectLStmtsBinders, collectLStmtBinders, collectStmtBinders )
Simon Marlow's avatar
Simon Marlow committed
25
import TcHsSyn
26 27 28 29 30 31

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

32
import {-# SOURCE #-} DsExpr ( dsExpr, dsLExpr, dsLocalBinds )
33

Simon Marlow's avatar
Simon Marlow committed
34
import TcType
35
import TcEvidence
36
import CoreSyn
Simon Marlow's avatar
Simon Marlow committed
37 38
import CoreFVs
import CoreUtils
39
import MkCore
40
import DsBinds (dsHsWrapper)
41

Simon Marlow's avatar
Simon Marlow committed
42
import Name
Ian Lynagh's avatar
Ian Lynagh committed
43
import Var
44
import Id
Simon Marlow's avatar
Simon Marlow committed
45 46 47 48
import DataCon
import TysWiredIn
import BasicTypes
import PrelNames
Ian Lynagh's avatar
Ian Lynagh committed
49
import Outputable
50
import Bag
Simon Marlow's avatar
Simon Marlow committed
51 52
import VarSet
import SrcLoc
53 54
import ListSetOps( assocDefault )
import FastString
55
import Data.List
56 57 58 59 60 61 62
\end{code}

\begin{code}
data DsCmdEnv = DsCmdEnv {
	arr_id, compose_id, first_id, app_id, choice_id, loop_id :: CoreExpr
    }

63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83
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)) }
 
    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)
84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116

-- 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 ->
		CoreExpr -> CoreExpr -> CoreExpr
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 ->
		CoreExpr -> CoreExpr -> CoreExpr
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]

117 118 119
-- 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 ->
120
		CoreExpr -> CoreExpr -> CoreExpr
121 122
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
123

124
mkFailExpr :: HsMatchContext Id -> Type -> DsM CoreExpr
125 126 127
mkFailExpr ctxt ty
  = mkErrorAppDs pAT_ERROR_ID ty (matchContextErrString ctxt)

128 129 130 131 132 133 134 135 136
-- 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)))

137 138
-- construct CoreExpr for \ (a :: a_ty, b :: b_ty) -> b
mkSndExpr :: Type -> Type -> DsM CoreExpr
139 140 141 142 143 144
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)))
145 146 147 148 149 150 151 152 153 154 155
\end{code}

Build case analysis of a tuple.  This cannot be done in the DsM monad,
because the list of variables is typically not yet defined.

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

coreCaseTuple :: UniqSupply -> Id -> [Id] -> CoreExpr -> CoreExpr
ross's avatar
ross committed
156 157
coreCaseTuple uniqs scrut_var vars body
  = mkTupleCase uniqs vars body scrut_var (Var scrut_var)
158

ross's avatar
ross committed
159 160
coreCasePair :: Id -> Id -> Id -> CoreExpr -> CoreExpr
coreCasePair scrut_var var1 var2 body
161
  = Case (Var scrut_var) scrut_var (exprType body)
batterseapower's avatar
batterseapower committed
162
         [(DataAlt (tupleCon BoxedTuple 2), [var1, var2], body)]
163 164 165 166
\end{code}

\begin{code}
mkCorePairTy :: Type -> Type -> Type
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
167
mkCorePairTy t1 t2 = mkBoxedTupleTy [t1, t2]
168 169 170

mkCorePairExpr :: CoreExpr -> CoreExpr -> CoreExpr
mkCorePairExpr e1 e2 = mkCoreTup [e1, e2]
171 172 173

mkCoreUnitExpr :: CoreExpr
mkCoreUnitExpr = mkCoreTup []
174 175 176
\end{code}

The input is divided into a local environment, which is a flat tuple
177 178
(unless it's too big), and a stack, which is a right-nested pair.
In general, the input has the form
179

180
	((x1,...,xn), (s1,...(sk,())...))
181 182 183 184 185

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.

\begin{code}
186 187 188 189 190 191 192 193 194 195
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)
196 197 198 199

----------------------------------------------
--		buildEnvStack
--
200
--	((x1,...,xn),stk)
201

202 203 204
buildEnvStack :: [Id] -> Id -> CoreExpr
buildEnvStack env_ids stack_id
  = mkCorePairExpr (mkBigCoreVarTup env_ids) (Var stack_id)
205 206 207 208

----------------------------------------------
-- 		matchEnvStack
--
209
--	\ ((x1,...,xn),stk) -> body
210
--	=>
211 212 213 214
--	\ pair ->
--	case pair of (tup,stk) ->
--	case tup of (x1,...,xn) ->
--	body
215

ross's avatar
ross committed
216
matchEnvStack	:: [Id] 	-- x1..xn
217
		-> Id	 	-- stk
218 219
		-> CoreExpr 	-- e
		-> DsM CoreExpr
220
matchEnvStack env_ids stack_id body = do
221 222
    uniqs <- newUniqueSupply
    tup_var <- newSysLocalDs (mkBigCoreVarTupTy env_ids)
223 224 225
    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))
226

227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242
----------------------------------------------
-- 		matchEnv
--
--	\ (x1,...,xn) -> body
--	=>
--	\ tup ->
--	case tup of (x1,...,xn) ->
--	body

matchEnv :: [Id] 	-- x1..xn
	 -> CoreExpr 	-- e
	 -> DsM CoreExpr
matchEnv env_ids body = do
    uniqs <- newUniqueSupply
    tup_id <- newSysLocalDs (mkBigCoreVarTupTy env_ids)
    return (Lam tup_id (coreCaseTuple uniqs tup_id env_ids body))
243 244

----------------------------------------------
245
--		matchVarStack
246
--
247
--	case (x1, ...(xn, s)...) -> e
248
--	=>
249 250
--	case z0 of (x1,z1) ->
--	case zn-1 of (xn,s) ->
251
--	e
252 253 254 255 256 257
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)
258 259 260
\end{code}

\begin{code}
261 262 263
mkHsEnvStackExpr :: [Id] -> Id -> LHsExpr Id
mkHsEnvStackExpr env_ids stack_id
  = mkLHsTupleExpr [mkLHsVarTuple env_ids, nlHsVar stack_id]
264 265 266 267 268 269
\end{code}

Translation of arrow abstraction

\begin{code}

270 271 272
-- D; xs |-a c : () --> t'  	---> c'
-- --------------------------
-- D |- proc p -> c :: a t t'	---> premap (\ p -> ((xs),())) c'
273 274 275 276
--
--		where (xs) is the tuple of variables bound by p

dsProcExpr
277 278
	:: LPat Id
	-> LHsCmdTop Id
279
	-> DsM CoreExpr
280
dsProcExpr pat (L _ (HsCmdTop cmd _unitTy cmd_ty ids)) = do
281
    (meth_binds, meth_ids) <- mkCmdEnv ids
282
    let locals = mkVarSet (collectPatBinders pat)
283
    (core_cmd, _free_vars, env_ids) <- dsfixCmd meth_ids locals unitTy cmd_ty cmd
284
    let env_ty = mkBigCoreVarTupTy env_ids
285 286 287
    let env_stk_ty = mkCorePairTy env_ty unitTy
    let env_stk_expr = mkCorePairExpr (mkBigCoreVarTup env_ids) mkCoreUnitExpr
    fail_expr <- mkFailExpr ProcExpr env_stk_ty
288
    var <- selectSimpleMatchVarL pat
289
    match_code <- matchSimply (Var var) ProcExpr pat env_stk_expr fail_expr
290
    let pat_ty = hsLPatType pat
291
        proc_code = do_premap meth_ids pat_ty env_stk_ty cmd_ty
292 293
                    (Lam var match_code)
                    core_cmd
294
    return (mkLets meth_binds proc_code)
295 296
\end{code}

297 298 299 300 301
Translation of a command judgement of the form

	D; xs |-a c : stk --> t

to an expression e such that
302

303
	D |- e :: a (xs, stk) t
304 305

\begin{code}
306
dsLCmd :: DsCmdEnv -> IdSet -> Type -> Type -> LHsCmd Id -> [Id]
Ian Lynagh's avatar
Ian Lynagh committed
307
       -> DsM (CoreExpr, IdSet)
308 309
dsLCmd ids local_vars stk_ty res_ty cmd env_ids
  = dsCmd ids local_vars stk_ty res_ty (unLoc cmd) env_ids
310

311
dsCmd   :: DsCmdEnv		-- arrow combinators
312
	-> IdSet		-- set of local vars available to this command
313
	-> Type			-- type of the stack (right-nested tuple)
314
	-> Type			-- return type of the command
315
	-> HsCmd Id		-- command to desugar
316 317 318
	-> [Id]			-- list of vars in the input to this command
				-- This is typically fed back,
				-- so don't pull on it too early
319
	-> DsM (CoreExpr,	-- desugared expression
320
		IdSet)		-- subset of local vars that occur free
321

322 323 324 325
-- D |- fun :: a t1 t2
-- D, xs |- arg :: t1
-- -----------------------------
-- D; xs |-a fun -< arg : stk --> t2
ross's avatar
ross committed
326
--
327
--		---> premap (\ ((xs), _stk) -> arg) fun
328

329
dsCmd ids local_vars stack_ty res_ty
330
        (HsCmdArrApp arrow arg arrow_ty HsFirstOrderApp _)
331
        env_ids = do
332 333
    let
        (a_arg_ty, _res_ty') = tcSplitAppTy arrow_ty
ross's avatar
ross committed
334
        (_a_ty, arg_ty) = tcSplitAppTy a_arg_ty
335 336
    core_arrow <- dsLExpr arrow
    core_arg   <- dsLExpr arg
337 338
    stack_id   <- newSysLocalDs stack_ty
    core_make_arg <- matchEnvStack env_ids stack_id core_arg
339
    return (do_premap ids
340
              (envStackType env_ids stack_ty)
341 342 343 344
              arg_ty
              res_ty
              core_make_arg
              core_arrow,
345
            exprFreeIds core_arg `intersectVarSet` local_vars)
346

347 348 349 350
-- D, xs |- fun :: a t1 t2
-- D, xs |- arg :: t1
-- ------------------------------
-- D; xs |-a fun -<< arg : stk --> t2
ross's avatar
ross committed
351
--
352
--		---> premap (\ ((xs), _stk) -> (fun, arg)) app
353

354
dsCmd ids local_vars stack_ty res_ty
355
        (HsCmdArrApp arrow arg arrow_ty HsHigherOrderApp _)
356
        env_ids = do
357 358
    let
        (a_arg_ty, _res_ty') = tcSplitAppTy arrow_ty
ross's avatar
ross committed
359
        (_a_ty, arg_ty) = tcSplitAppTy a_arg_ty
360 361 362
    
    core_arrow <- dsLExpr arrow
    core_arg   <- dsLExpr arg
363 364 365 366
    stack_id   <- newSysLocalDs stack_ty
    core_make_pair <- matchEnvStack env_ids stack_id
          (mkCorePairExpr core_arrow core_arg)

367
    return (do_premap ids
368
              (envStackType env_ids stack_ty)
369 370 371 372
              (mkCorePairTy arrow_ty arg_ty)
              res_ty
              core_make_pair
              (do_app ids arg_ty res_ty),
373
            (exprFreeIds core_arrow `unionVarSet` exprFreeIds core_arg)
374
              `intersectVarSet` local_vars)
375

376 377 378 379
-- D; ys |-a cmd : (t,stk) --> t'
-- D, xs |-  exp :: t
-- ------------------------
-- D; xs |-a cmd exp : stk --> t'
ross's avatar
ross committed
380
--
381
--		---> premap (\ ((xs),stk) -> ((ys),(e,stk))) cmd
ross's avatar
ross committed
382

383
dsCmd ids local_vars stack_ty res_ty (HsCmdApp cmd arg) env_ids = do
384
    core_arg <- dsLExpr arg
ross's avatar
ross committed
385
    let
386
        arg_ty = exprType core_arg
387
        stack_ty' = mkCorePairTy arg_ty stack_ty
388
    (core_cmd, free_vars, env_ids')
389 390
             <- dsfixCmd ids local_vars stack_ty' res_ty cmd
    stack_id <- newSysLocalDs stack_ty
391
    arg_id <- newSysLocalDs arg_ty
ross's avatar
ross committed
392 393
    -- push the argument expression onto the stack
    let
394
	stack' = mkCorePairExpr (Var arg_id) (Var stack_id)
395
        core_body = bindNonRec arg_id core_arg
396 397
			(mkCorePairExpr (mkBigCoreVarTup env_ids') stack')

ross's avatar
ross committed
398
    -- match the environment and stack against the input
399
    core_map <- matchEnvStack env_ids stack_id core_body
400
    return (do_premap ids
401 402
                      (envStackType env_ids stack_ty)
                      (envStackType env_ids' stack_ty')
403 404 405
                      res_ty
                      core_map
                      core_cmd,
406 407
            free_vars `unionVarSet`
              (exprFreeIds core_arg `intersectVarSet` local_vars))
ross's avatar
ross committed
408

409 410 411
-- D; ys |-a cmd : stk t'
-- -----------------------------------------------
-- D; xs |-a \ p1 ... pk -> cmd : (t1,...(tk,stk)...) t'
412
--
413
--		---> premap (\ ((xs), (p1, ... (pk,stk)...)) -> ((ys),stk)) cmd
414

415
dsCmd ids local_vars stack_ty res_ty
416
        (HsCmdLam (MG { mg_alts = [L _ (Match pats _ (GRHSs [L _ (GRHS [] body)] _ ))] }))
417
        env_ids = do
418 419
    let
        pat_vars = mkVarSet (collectPatsBinders pats)
420
        local_vars' = pat_vars `unionVarSet` local_vars
421 422 423 424
	(pat_tys, stack_ty') = splitTypeAt (length pats) stack_ty
    (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'
425 426 427 428 429

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

    let
430
        -- build a new environment, plus what's left of the stack
431 432 433
        core_expr = buildEnvStack env_ids' stack_id'
        in_ty = envStackType env_ids stack_ty
        in_ty' = envStackType env_ids' stack_ty'
434 435
    
    fail_expr <- mkFailExpr LambdaExpr in_ty'
436 437 438 439
    -- 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
440
    -- match the old environment and stack against the input
441
    select_code <- matchEnvStack env_ids stack_id param_code
442
    return (do_premap ids in_ty in_ty' res_ty select_code core_body,
443
            free_vars `minusVarSet` pat_vars)
444

445 446
dsCmd ids local_vars stack_ty res_ty (HsCmdPar cmd) env_ids
  = dsLCmd ids local_vars stack_ty res_ty cmd env_ids
447

448 449 450 451 452
-- 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
453
--
454 455
--		---> premap (\ ((xs),stk) ->
--			 if e then Left ((xs1),stk) else Right ((xs2),stk))
456
--		       (c1 ||| c2)
ross's avatar
ross committed
457

458
dsCmd ids local_vars stack_ty res_ty (HsCmdIf mb_fun cond then_cmd else_cmd)
459
        env_ids = do
460
    core_cond <- dsLExpr cond
461 462 463
    (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
464 465 466
    either_con <- dsLookupTyCon eitherTyConName
    left_con   <- dsLookupDataCon leftDataConName
    right_con  <- dsLookupDataCon rightDataConName
467 468 469

    let mk_left_expr ty1 ty2 e = mkConApp left_con [Type ty1, Type ty2, e]
        mk_right_expr ty1 ty2 e = mkConApp right_con [Type ty1, Type ty2, e]
470

471 472 473
        in_ty = envStackType env_ids stack_ty
        then_ty = envStackType then_ids stack_ty
        else_ty = envStackType else_ids stack_ty
474
        sum_ty = mkTyConApp either_con [then_ty, else_ty]
475
        fvs_cond = exprFreeIds core_cond `intersectVarSet` local_vars
476
        
477 478
        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)
479 480 481

    core_if <- case mb_fun of 
       Just fun -> do { core_fun <- dsExpr fun
482
                      ; matchEnvStack env_ids stack_id $
483
                        mkCoreApps core_fun [core_cond, core_left, core_right] }
484
       Nothing  -> matchEnvStack env_ids stack_id $
485 486
                   mkIfThenElse core_cond core_left core_right

487
    return (do_premap ids in_ty sum_ty res_ty
488 489 490
                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
491 492 493 494 495 496 497 498 499
\end{code}

Case commands are treated in much the same way as if commands
(see above) except that there are more alternatives.  For example

	case e of { p1 -> c1; p2 -> c2; p3 -> c3 }

is translated to

500
	premap (\ ((xs)*ts) -> case e of
ross's avatar
ross committed
501 502
		p1 -> (Left (Left (xs1)*ts))
		p2 -> Left ((Right (xs2)*ts))
503 504
		p3 -> Right ((xs3)*ts))
	((c1 ||| c2) ||| c3)
ross's avatar
ross committed
505 506 507 508

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.
509
To build all this, we use triples describing segments of the list of
ross's avatar
ross committed
510
case bodies, containing the following fields:
511
 * a list of expressions of the form (Left|Right)* ((xs)*ts), to be put
ross's avatar
ross committed
512
   into the case replacing the commands
513
 * a sum type that is the common type of these expressions, and also the
ross's avatar
ross committed
514
   input type of the arrow
515
 * a CoreExpr for an arrow built by combining the translated command
ross's avatar
ross committed
516 517 518
   bodies with |||.

\begin{code}
519
dsCmd ids local_vars stack_ty res_ty 
520 521
      (HsCmdCase exp (MG { mg_alts = matches, mg_arg_tys = arg_tys }))
      env_ids = do
522
    stack_id <- newSysLocalDs stack_ty
523 524 525 526 527 528

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

    let
        leaves = concatMap leavesMatch matches
529
        make_branch (leaf, bound_vars) = do
530
            (core_leaf, _fvs, leaf_ids) <-
531 532 533
                  dsfixCmd ids (bound_vars `unionVarSet` local_vars) stack_ty res_ty leaf
            return ([mkHsEnvStackExpr leaf_ids stack_id],
                    envStackType leaf_ids stack_ty,
534 535 536 537 538 539
                    core_leaf)
    
    branches <- mapM make_branch leaves
    either_con <- dsLookupTyCon eitherTyConName
    left_con <- dsLookupDataCon leftDataConName
    right_con <- dsLookupDataCon rightDataConName
540
    let
541 542 543 544
        left_id  = HsVar (dataConWrapId left_con)
        right_id = HsVar (dataConWrapId right_con)
        left_expr  ty1 ty2 e = noLoc $ HsApp (noLoc $ HsWrap (mkWpTyApps [ty1, ty2]) left_id ) e
        right_expr ty1 ty2 e = noLoc $ HsApp (noLoc $ HsWrap (mkWpTyApps [ty1, ty2]) right_id) e
545

546 547
        -- Prefix each tuple with a distinct series of Left's and Right's,
        -- in a balanced way, keeping track of the types.
548

549 550 551
        merge_branches (builds1, in_ty1, core_exp1)
                       (builds2, in_ty2, core_exp2)
          = (map (left_expr in_ty1 in_ty2) builds1 ++
552 553 554
                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)
555
        (leaves', sum_ty, core_choices) = foldb merge_branches branches
556 557 558 559 560

        -- 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
561
        in_ty = envStackType env_ids stack_ty
562

563 564
    core_body <- dsExpr (HsCase exp (MG { mg_alts = matches', mg_arg_tys = arg_tys
                                        , mg_res_ty = sum_ty }))
565 566
        -- Note that we replace the HsCase result type by sum_ty,
        -- which is the type of matches'
567

568
    core_matches <- matchEnvStack env_ids stack_id core_body
569
    return (do_premap ids in_ty sum_ty res_ty core_matches core_choices,
570
            exprFreeIds core_body  `intersectVarSet` local_vars)
571

572 573 574
-- D; ys |-a cmd : stk --> t
-- ----------------------------------
-- D; xs |-a let binds in cmd : stk --> t
575
--
576
--		---> premap (\ ((xs),stk) -> let binds in ((ys),stk)) c
577

578
dsCmd ids local_vars stack_ty res_ty (HsCmdLet binds body) env_ids = do
579
    let
580
        defined_vars = mkVarSet (collectLocalBinders binds)
581
        local_vars' = defined_vars `unionVarSet` local_vars
582
    
583 584
    (core_body, _free_vars, env_ids') <- dsfixCmd ids local_vars' stack_ty res_ty body
    stack_id <- newSysLocalDs stack_ty
585
    -- build a new environment, plus the stack, using the let bindings
586
    core_binds <- dsLocalBinds binds (buildEnvStack env_ids' stack_id)
587
    -- match the old environment and stack against the input
588
    core_map <- matchEnvStack env_ids stack_id core_binds
589
    return (do_premap ids
590 591
                        (envStackType env_ids stack_ty)
                        (envStackType env_ids' stack_ty)
592 593 594
                        res_ty
                        core_map
                        core_body,
595
        exprFreeIds core_binds `intersectVarSet` local_vars)
596

597 598 599 600 601
-- D; xs |-a ss : t
-- ----------------------------------
-- D; xs |-a do { ss } : () --> t
--
--		---> premap (\ (env,stk) -> env) c
602

603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620
dsCmd ids local_vars stack_ty res_ty (HsCmdDo stmts _) env_ids = do
    (core_stmts, env_ids') <- dsCmdDo ids local_vars res_ty stmts env_ids
    let env_ty = mkBigCoreVarTupTy env_ids
    core_fst <- mkFstExpr env_ty stack_ty
    return (do_premap ids
		(mkCorePairTy env_ty stack_ty)
		env_ty
		res_ty
		core_fst
		core_stmts,
	env_ids')

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

dsCmd _ids local_vars _stack_ty _res_ty (HsCmdArrForm op _ args) env_ids = do
621 622 623 624 625
    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)
626

627 628 629 630
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
    wrapped_cmd <- dsHsWrapper (WpCast coercion) core_cmd
    return (wrapped_cmd, env_ids')
631

Ian Lynagh's avatar
Ian Lynagh committed
632 633
dsCmd _ _ _ _ _ c = pprPanic "dsCmd" (ppr c)

634 635 636
-- D; ys |-a c : stk --> t	(ys <= xs)
-- ---------------------
-- D; xs |-a c : stk --> t	---> premap (\ ((xs),stk) -> ((ys),stk)) c
637 638 639 640

dsTrimCmdArg
	:: IdSet		-- set of local vars available to this command
	-> [Id]			-- list of vars in the input to this command
641
	-> LHsCmdTop Id		-- command argument to desugar
642
	-> DsM (CoreExpr,	-- desugared expression
643
		IdSet)		-- subset of local vars that occur free
644
dsTrimCmdArg local_vars env_ids (L _ (HsCmdTop cmd stack_ty cmd_ty ids)) = do
645
    (meth_binds, meth_ids) <- mkCmdEnv ids
646 647 648
    (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)
649
    let
650 651
        in_ty = envStackType env_ids stack_ty
        in_ty' = envStackType env_ids' stack_ty
652
        arg_code = if env_ids' == env_ids then core_cmd else
653
                do_premap meth_ids in_ty in_ty' cmd_ty trim_code core_cmd
654
    return (mkLets meth_binds arg_code, free_vars)
655

656 657
-- Given D; xs |-a c : stk --> t, builds c with xs fed back.
-- Typically needs to be prefixed with arr (\(p, stk) -> ((xs),stk))
658 659 660 661

dsfixCmd
	:: DsCmdEnv		-- arrow combinators
	-> IdSet		-- set of local vars available to this command
662
	-> Type			-- type of the stack (right-nested tuple)
663
	-> Type			-- return type of the command
664
	-> LHsCmd Id		-- command to desugar
665
	-> DsM (CoreExpr,	-- desugared expression
666 667
		IdSet,		-- subset of local vars that occur free
		[Id])		-- the same local vars as a list, fed back
668 669
dsfixCmd ids local_vars stk_ty cmd_ty cmd
  = trimInput (dsLCmd ids local_vars stk_ty cmd_ty cmd)
670 671 672 673 674 675 676 677 678 679 680 681 682 683

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

trimInput
	:: ([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.
trimInput build_arrow
  = fixDs (\ ~(_,_,env_ids) -> do
        (core_cmd, free_vars) <- build_arrow env_ids
684
        return (core_cmd, free_vars, varSetElems free_vars))
685 686 687 688 689

\end{code}

Translation of command judgements of the form

690
	D |-a do { ss } : t
691 692 693 694 695

\begin{code}

dsCmdDo :: DsCmdEnv		-- arrow combinators
	-> IdSet		-- set of local vars available to this statement
696
	-> Type			-- return type of the statement
697
	-> [CmdLStmt Id]        -- statements to desugar
698 699 700 701
	-> [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
702
		IdSet)		-- subset of local vars that occur free
703

704
dsCmdDo _ _ _ [] _ = panic "dsCmdDo"
705

706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723
-- D; xs |-a c : () --> t
-- --------------------------
-- D; xs |-a do { c } : t
--
--		---> premap (\ (xs) -> ((xs), ())) c

dsCmdDo ids local_vars res_ty [L _ (LastStmt body _)] env_ids = do
    (core_body, env_ids') <- dsLCmd ids local_vars unitTy res_ty body env_ids
    let env_ty = mkBigCoreVarTupTy env_ids
    env_var <- newSysLocalDs env_ty
    let core_map = Lam env_var (mkCorePairExpr (Var env_var) mkCoreUnitExpr)
    return (do_premap ids
                        env_ty
			(mkCorePairTy env_ty unitTy)
                        res_ty
                        core_map
                        core_body,
	env_ids')
724

725
dsCmdDo ids local_vars res_ty (stmt:stmts) env_ids = do
726
    let
727
        bound_vars = mkVarSet (collectLStmtBinders stmt)
728 729 730
        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
731 732 733 734 735 736 737
    return (do_compose ids
                (mkBigCoreVarTupTy env_ids)
                (mkBigCoreVarTupTy env_ids')
                res_ty
                core_stmt
                core_stmts,
              fv_stmt)
738

ross's avatar
ross committed
739 740 741 742 743
\end{code}
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.
\begin{code}
744
dsCmdLStmt :: DsCmdEnv -> IdSet -> [Id] -> CmdLStmt Id -> [Id]
Ian Lynagh's avatar
Ian Lynagh committed
745
           -> DsM (CoreExpr, IdSet)
746 747
dsCmdLStmt ids local_vars out_ids cmd env_ids
  = dsCmdStmt ids local_vars out_ids (unLoc cmd) env_ids
ross's avatar
ross committed
748

749 750 751
dsCmdStmt
	:: DsCmdEnv		-- arrow combinators
	-> IdSet		-- set of local vars available to this statement
752
	-> [Id]			-- list of vars in the output of this statement
753
	-> CmdStmt Id           -- statement to desugar
754 755 756 757
	-> [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
758
		IdSet)		-- subset of local vars that occur free
759

760 761 762 763
-- D; xs1 |-a c : () --> t
-- D; xs' |-a do { ss } : t'
-- ------------------------------
-- D; xs  |-a do { c; ss } : t'
764
--
765
--		---> premap (\ ((xs)) -> (((xs1),()),(xs')))
766
--			(first c >>> arr snd) >>> ss
767

768
dsCmdStmt ids local_vars out_ids (BodyStmt cmd _ _ c_ty) env_ids = do
769 770 771 772 773
    (core_cmd, fv_cmd, env_ids1) <- dsfixCmd ids local_vars unitTy c_ty cmd
    core_mux <- matchEnv env_ids
        (mkCorePairExpr
	    (mkCorePairExpr (mkBigCoreVarTup env_ids1) mkCoreUnitExpr)
	    (mkBigCoreVarTup out_ids))
774
    let
775
	in_ty = mkBigCoreVarTupTy env_ids
776
	in_ty1 = mkCorePairTy (mkBigCoreVarTupTy env_ids1) unitTy
777
	out_ty = mkBigCoreVarTupTy out_ids
778 779
	before_c_ty = mkCorePairTy in_ty1 out_ty
	after_c_ty = mkCorePairTy c_ty out_ty
780
    snd_fn <- mkSndExpr c_ty out_ty
781
    return (do_premap ids in_ty before_c_ty out_ty core_mux $
782 783 784
		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,
785
	      extendVarSetList fv_cmd out_ids)
786

787 788 789 790
-- D; xs1 |-a c : () --> t
-- D; xs' |-a do { ss } : t'		xs2 = xs' - defs(p)
-- -----------------------------------
-- D; xs  |-a do { p <- c; ss } : t'
791
--
792
--		---> premap (\ (xs) -> (((xs1),()),(xs2)))
793
--			(first c >>> arr (\ (p, (xs2)) -> (xs'))) >>> ss
794 795 796 797
--
-- It would be simpler and more consistent to do this using second,
-- but that's likely to be defined in terms of first.

798
dsCmdStmt ids local_vars out_ids (BindStmt pat cmd _ _) env_ids = do
799
    (core_cmd, fv_cmd, env_ids1) <- dsfixCmd ids local_vars unitTy (hsLPatType pat) cmd
800
    let
801
	pat_ty = hsLPatType pat
802 803
	pat_vars = mkVarSet (collectPatBinders pat)
	env_ids2 = varSetElems (mkVarSet out_ids `minusVarSet` pat_vars)
804
	env_ty2 = mkBigCoreVarTupTy env_ids2
805 806

    -- multiplexing function
807
    --		\ (xs) -> (((xs1),()),(xs2))
808

809 810 811 812
    core_mux <- matchEnv env_ids
        (mkCorePairExpr
	    (mkCorePairExpr (mkBigCoreVarTup env_ids1) mkCoreUnitExpr)
	    (mkBigCoreVarTup env_ids2))
813 814 815 816

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

817 818
    env_id <- newSysLocalDs env_ty2
    uniqs <- newUniqueSupply
819
    let
ross's avatar
ross committed
820
	after_c_ty = mkCorePairTy pat_ty env_ty2
821 822
	out_ty = mkBigCoreVarTupTy out_ids
	body_expr = coreCaseTuple uniqs env_id env_ids2 (mkBigCoreVarTup out_ids)
823 824 825 826 827
    
    fail_expr <- mkFailExpr (StmtCtxt DoExpr) out_ty
    pat_id    <- selectSimpleMatchVarL pat
    match_code <- matchSimply (Var pat_id) (StmtCtxt DoExpr) pat body_expr fail_expr
    pair_id   <- newSysLocalDs after_c_ty
828
    let
ross's avatar
ross committed
829
	proj_expr = Lam pair_id (coreCasePair pair_id pat_id env_id match_code)
830

ross's avatar
ross committed
831
    -- put it all together
832
    let
833
	in_ty = mkBigCoreVarTupTy env_ids
834
	in_ty1 = mkCorePairTy (mkBigCoreVarTupTy env_ids1) unitTy
835
	in_ty2 = mkBigCoreVarTupTy env_ids2
836
	before_c_ty = mkCorePairTy in_ty1 in_ty2
837
    return (do_premap ids in_ty before_c_ty out_ty core_mux $
838 839 840 841 842
		do_compose ids before_c_ty after_c_ty out_ty
			(do_first ids in_ty1 pat_ty in_ty2 core_cmd) $
		do_arr ids after_c_ty out_ty proj_expr,
	      fv_cmd `unionVarSet` (mkVarSet out_ids `minusVarSet` pat_vars))

843 844 845
-- D; xs' |-a do { ss } : t
-- --------------------------------------
-- D; xs  |-a do { let binds; ss } : t
846 847 848
--
--		---> arr (\ (xs) -> let binds in (xs')) >>> ss

849
dsCmdStmt ids local_vars out_ids (LetStmt binds) env_ids = do
850
    -- build a new environment using the let bindings
851
    core_binds <- dsLocalBinds binds (mkBigCoreVarTup out_ids)
852
    -- match the old environment against the input
853
    core_map <- matchEnv env_ids core_binds
854
    return (do_arr ids
855 856
			(mkBigCoreVarTupTy env_ids)
			(mkBigCoreVarTupTy out_ids)
857
			core_map,
858
	    exprFreeIds core_binds `intersectVarSet` local_vars)
859

860 861 862 863
-- D; ys  |-a do { ss; returnA -< ((xs1), (ys2)) } : ...
-- D; xs' |-a do { ss' } : t
-- ------------------------------------
-- D; xs  |-a do { rec ss; ss' } : t
864 865 866 867 868 869 870 871 872 873
--
--			xs1 = xs' /\ defs(ss)
--			xs2 = xs' - defs(ss)
--			ys1 = ys - defs(ss)
--			ys2 = ys /\ defs(ss)
--
--		---> arr (\(xs) -> ((ys1),(xs2))) >>>
--			first (loop (arr (\((ys1),~(ys2)) -> (ys)) >>> ss)) >>>
--			arr (\((xs1),(xs2)) -> (xs')) >>> ss'

874 875 876 877 878
dsCmdStmt ids local_vars out_ids
        (RecStmt { recS_stmts = stmts
                 , recS_later_ids = later_ids, recS_rec_ids = rec_ids
                 , recS_later_rets = later_rets, recS_rec_rets = rec_rets })
        env_ids = do
879
    let
880 881 882
        env2_id_set = mkVarSet out_ids `minusVarSet` mkVarSet later_ids
        env2_ids = varSetElems env2_id_set
        env2_ty = mkBigCoreVarTupTy env2_ids
ross's avatar
ross committed
883 884 885

    -- post_loop_fn = \((later_ids),(env2_ids)) -> (out_ids)

886 887
    uniqs <- newUniqueSupply
    env2_id <- newSysLocalDs env2_ty
ross's avatar
ross committed
888
    let
889 890 891 892
        later_ty = mkBigCoreVarTupTy later_ids
        post_pair_ty = mkCorePairTy later_ty env2_ty
        post_loop_body = coreCaseTuple uniqs env2_id env2_ids (mkBigCoreVarTup out_ids)

893
    post_loop_fn <- matchEnvStack later_ids env2_id post_loop_body
ross's avatar
ross committed
894 895 896

    --- loop (...)

897
    (core_loop, env1_id_set, env1_ids)
898
               <- dsRecCmd ids local_vars stmts later_ids later_rets rec_ids rec_rets
ross's avatar
ross committed
899 900 901 902

    -- pre_loop_fn = \(env_ids) -> ((env1_ids),(env2_ids))

    let
903 904 905 906
        env1_ty = mkBigCoreVarTupTy env1_ids
        pre_pair_ty = mkCorePairTy env1_ty env2_ty
        pre_loop_body = mkCorePairExpr (mkBigCoreVarTup env1_ids)
                                        (mkBigCoreVarTup env2_ids)
ross's avatar
ross committed
907

908
    pre_loop_fn <- matchEnv env_ids pre_loop_body
ross's avatar
ross committed
909 910 911 912

    -- arr pre_loop_fn >>> first (loop (...)) >>> arr post_loop_fn

    let
913 914
        env_ty = mkBigCoreVarTupTy env_ids
        out_ty = mkBigCoreVarTupTy out_ids
915
        core_body = do_premap ids env_ty pre_pair_ty out_ty
916 917 918 919 920 921 922 923
                pre_loop_fn
                (do_compose ids pre_pair_ty post_pair_ty out_ty
                        (do_first ids env1_ty later_ty env2_ty
                                core_loop)
                        (do_arr ids post_pair_ty out_ty
                                post_loop_fn))

    return (core_body, env1_id_set `unionVarSet` env2_id_set)
ross's avatar
ross committed
924

Ian Lynagh's avatar
Ian Lynagh committed
925 926
dsCmdStmt _ _ _ _ s = pprPanic "dsCmdStmt" (ppr s)

927 928
--	loop (premap (\ ((env1_ids), ~(rec_ids)) -> (env_ids))
--	      (ss >>> arr (\ (out_ids) -> ((later_rets),(rec_rets))))) >>>
929 930 931 932

dsRecCmd
        :: DsCmdEnv		-- arrow combinators
        -> IdSet		-- set of local vars available to this statement
933
        -> [CmdLStmt Id]        -- list of statements inside the RecCmd
934 935 936 937 938 939 940 941 942
        -> [Id]			-- list of vars defined here and used later
        -> [HsExpr Id]		-- expressions corresponding to later_ids
        -> [Id]			-- list of vars fed back through the loop
        -> [HsExpr Id]		-- expressions corresponding to rec_ids
        -> DsM (CoreExpr,	-- desugared statement
                IdSet,		-- subset of local vars that occur free
                [Id])		-- same local vars as a list

dsRecCmd ids local_vars stmts later_ids later_rets rec_ids rec_rets = do
943
    let
944
        later_id_set = mkVarSet later_ids
945
        rec_id_set = mkVarSet rec_ids
946
        local_vars' = rec_id_set `unionVarSet` later_id_set `unionVarSet` local_vars
947

948
    -- mk_pair_fn = \ (out_ids) -> ((later_rets),(rec_rets))
949

950 951
    core_later_rets <- mapM dsExpr later_rets
    core_rec_rets <- mapM dsExpr rec_rets
952
    let
953 954 955 956 957
        -- possibly polymorphic version of vars of later_ids and rec_ids
        out_ids = varSetElems (unionVarSets (map exprFreeIds (core_later_rets ++ core_rec_rets)))
        out_ty = mkBigCoreVarTupTy out_ids

        later_tuple = mkBigCoreTup core_later_rets
958
        later_ty = mkBigCoreVarTupTy later_ids
959 960

        rec_tuple = mkBigCoreTup core_rec_rets
961
        rec_ty = mkBigCoreVarTupTy rec_ids
962

963 964 965
        out_pair = mkCorePairExpr later_tuple rec_tuple
        out_pair_ty = mkCorePairTy later_ty rec_ty

966
    mk_pair_fn <- matchEnv out_ids out_pair
967

ross's avatar
ross committed
968 969
    -- ss

twanvl's avatar