DsArrows.lhs 37.8 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 9 10 11 12

\begin{code}
module DsArrows ( dsProcExpr ) where

#include "HsVersions.h"

Simon Marlow's avatar
Simon Marlow committed
13 14
import Match
import DsUtils
15 16
import DsMonad

17
import HsSyn	hiding (collectPatBinders, collectPatsBinders )
Simon Marlow's avatar
Simon Marlow committed
18
import TcHsSyn
19 20 21 22 23 24

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

25
import {-# SOURCE #-} DsExpr ( dsExpr, dsLExpr, dsLocalBinds )
26

Simon Marlow's avatar
Simon Marlow committed
27 28
import TcType
import Type
29
import CoreSyn
Simon Marlow's avatar
Simon Marlow committed
30 31
import CoreFVs
import CoreUtils
32
import MkCore
33

Simon Marlow's avatar
Simon Marlow committed
34
import Name
Ian Lynagh's avatar
Ian Lynagh committed
35
import Var
36
import Id
Simon Marlow's avatar
Simon Marlow committed
37 38 39 40
import DataCon
import TysWiredIn
import BasicTypes
import PrelNames
Ian Lynagh's avatar
Ian Lynagh committed
41
import Outputable
42
import Bag
Simon Marlow's avatar
Simon Marlow committed
43 44
import VarSet
import SrcLoc
45 46

import Data.List
47 48 49 50 51 52 53 54
\end{code}

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

55
mkCmdEnv :: SyntaxTable Id -> DsM DsCmdEnv
56 57
mkCmdEnv ids = do
    (meth_binds, ds_meths) <- dsSyntaxTable ids
58
    return $ DsCmdEnv {
59 60 61 62 63 64 65 66
               meth_binds = meth_binds,
               arr_id     = Var (lookupEvidence ds_meths arrAName),
               compose_id = Var (lookupEvidence ds_meths composeAName),
               first_id   = Var (lookupEvidence ds_meths firstAName),
               app_id     = Var (lookupEvidence ds_meths appAName),
               choice_id  = Var (lookupEvidence ds_meths choiceAName),
               loop_id    = Var (lookupEvidence ds_meths loopAName)
             }
67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106

bindCmdEnv :: DsCmdEnv -> CoreExpr -> CoreExpr
bindCmdEnv ids body = foldr Let body (meth_binds ids)

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

-- map_arrow (f :: b -> c) (g :: a c d) = arr f >>> g :: a b d
do_map_arrow :: DsCmdEnv -> Type -> Type -> Type ->
		CoreExpr -> CoreExpr -> CoreExpr
do_map_arrow ids b_ty c_ty d_ty f c
107
   = do_compose ids b_ty c_ty d_ty (do_arr ids b_ty c_ty f) c
108

109
mkFailExpr :: HsMatchContext Id -> Type -> DsM CoreExpr
110 111 112 113 114
mkFailExpr ctxt ty
  = mkErrorAppDs pAT_ERROR_ID ty (matchContextErrString ctxt)

-- construct CoreExpr for \ (a :: a_ty, b :: b_ty) -> b
mkSndExpr :: Type -> Type -> DsM CoreExpr
115 116 117 118 119 120
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)))
121 122 123 124 125 126 127 128 129 130 131
\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
132 133
coreCaseTuple uniqs scrut_var vars body
  = mkTupleCase uniqs vars body scrut_var (Var scrut_var)
134

ross's avatar
ross committed
135 136
coreCasePair :: Id -> Id -> Id -> CoreExpr -> CoreExpr
coreCasePair scrut_var var1 var2 body
137
  = Case (Var scrut_var) scrut_var (exprType body)
ross's avatar
ross committed
138
         [(DataAlt (tupleCon Boxed 2), [var1, var2], body)]
139 140 141 142
\end{code}

\begin{code}
mkCorePairTy :: Type -> Type -> Type
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
143
mkCorePairTy t1 t2 = mkBoxedTupleTy [t1, t2]
144 145 146 147 148 149 150

mkCorePairExpr :: CoreExpr -> CoreExpr -> CoreExpr
mkCorePairExpr e1 e2 = mkCoreTup [e1, e2]
\end{code}

The input is divided into a local environment, which is a flat tuple
(unless it's too big), and a stack, each element of which is paired
Ross Paterson's avatar
Ross Paterson committed
151
with the environment in turn.  In general, the input has the form
152 153 154 155 156 157 158 159

	(...((x1,...,xn),s1),...sk)

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}
envStackType :: [Id] -> [Type] -> Type
160
envStackType ids stack_tys = foldl mkCorePairTy (mkBigCoreVarTupTy ids) stack_tys
161 162 163 164

----------------------------------------------
--		buildEnvStack
--
ross's avatar
ross committed
165
--	(...((x1,...,xn),s1),...sk)
166 167 168

buildEnvStack :: [Id] -> [Id] -> CoreExpr
buildEnvStack env_ids stack_ids
169
  = foldl mkCorePairExpr (mkBigCoreVarTup env_ids) (map Var stack_ids)
170 171 172 173

----------------------------------------------
-- 		matchEnvStack
--
ross's avatar
ross committed
174
--	\ (...((x1,...,xn),s1),...sk) -> e
175
--	=>
ross's avatar
ross committed
176 177
--	\ zk ->
--	case zk of (zk-1,sk) ->
178 179
--	...
--	case z1 of (z0,s1) ->
ross's avatar
ross committed
180
--	case z0 of (x1,...,xn) ->
181 182
--	e

ross's avatar
ross committed
183 184
matchEnvStack	:: [Id] 	-- x1..xn
		-> [Id] 	-- s1..sk
185 186
		-> CoreExpr 	-- e
		-> DsM CoreExpr
187 188 189 190 191
matchEnvStack env_ids stack_ids body = do
    uniqs <- newUniqueSupply
    tup_var <- newSysLocalDs (mkBigCoreVarTupTy env_ids)
    matchVarStack tup_var stack_ids
               (coreCaseTuple uniqs tup_var env_ids body)
192 193 194 195 196


----------------------------------------------
-- 		matchVarStack
--
ross's avatar
ross committed
197
--	\ (...(z0,s1),...sk) -> e
198
--	=>
ross's avatar
ross committed
199 200
--	\ zk ->
--	case zk of (zk-1,sk) ->
201 202 203 204 205
--	...
--	case z1 of (z0,s1) ->
--	e

matchVarStack :: Id 		-- z0
ross's avatar
ross committed
206
	      -> [Id] 		-- s1..sk
207 208 209
	      -> CoreExpr 	-- e
	      -> DsM CoreExpr
matchVarStack env_id [] body
210 211 212 213 214
  = return (Lam env_id body)
matchVarStack env_id (stack_id:stack_ids) body = do
    pair_id <- newSysLocalDs (mkCorePairTy (idType env_id) (idType stack_id))
    matchVarStack pair_id stack_ids
               (coreCasePair pair_id env_id stack_id body)
215 216 217
\end{code}

\begin{code}
218
mkHsEnvStackExpr :: [Id] -> [Id] -> LHsExpr Id
219
mkHsEnvStackExpr env_ids stack_ids
220 221 222
  = foldl (\a b -> mkLHsTupleExpr [a,b]) 
	  (mkLHsVarTuple env_ids) 
	  (map nlHsVar stack_ids)
223 224 225 226 227 228 229 230 231 232 233 234 235
\end{code}

Translation of arrow abstraction

\begin{code}

--	A | xs |- c :: [] t'  	    ---> c'
--	--------------------------
--	A |- proc p -> c :: a t t'  ---> arr (\ p -> (xs)) >>> c'
--
--		where (xs) is the tuple of variables bound by p

dsProcExpr
236 237
	:: LPat Id
	-> LHsCmdTop Id
238
	-> DsM CoreExpr
239 240 241
dsProcExpr pat (L _ (HsCmdTop cmd [] cmd_ty ids)) = do
    meth_ids <- mkCmdEnv ids
    let locals = mkVarSet (collectPatBinders pat)
Ian Lynagh's avatar
Ian Lynagh committed
242
    (core_cmd, _free_vars, env_ids) <- dsfixCmd meth_ids locals [] cmd_ty cmd
243 244 245 246 247 248 249 250 251
    let env_ty = mkBigCoreVarTupTy env_ids
    fail_expr <- mkFailExpr ProcExpr env_ty
    var <- selectSimpleMatchVarL pat
    match_code <- matchSimply (Var var) ProcExpr pat (mkBigCoreVarTup env_ids) fail_expr
    let pat_ty = hsLPatType pat
        proc_code = do_map_arrow meth_ids pat_ty env_ty cmd_ty
                    (Lam var match_code)
                    core_cmd
    return (bindCmdEnv meth_ids proc_code)
Ian Lynagh's avatar
Ian Lynagh committed
252
dsProcExpr _ c = pprPanic "dsProcExpr" (ppr c)
253 254 255 256 257 258 259
\end{code}

Translation of command judgements of the form

	A | xs |- c :: [ts] t

\begin{code}
Ian Lynagh's avatar
Ian Lynagh committed
260 261
dsLCmd :: DsCmdEnv -> IdSet -> [Id] -> [Type] -> Type -> LHsCmd Id
       -> DsM (CoreExpr, IdSet)
262 263
dsLCmd ids local_vars env_ids stack res_ty cmd
  = dsCmd ids local_vars env_ids stack res_ty (unLoc cmd)
264

265
dsCmd   :: DsCmdEnv		-- arrow combinators
266 267 268 269 270 271
	-> IdSet		-- set of local vars available to this command
	-> [Id]			-- list of vars in the input to this command
				-- This is typically fed back,
				-- so don't pull on it too early
	-> [Type]		-- type of the stack
	-> Type			-- return type of the command
272
	-> HsCmd Id		-- command to desugar
273 274 275
	-> DsM (CoreExpr,	-- desugared expression
		IdSet)		-- set of local vars that occur free

ross's avatar
ross committed
276
--	A |- f :: a (t*ts) t'
277
--	A, xs |- arg :: t
ross's avatar
ross committed
278 279 280 281
--	-----------------------------
--	A | xs |- f -< arg :: [ts] t'
--
--		---> arr (\ ((xs)*ts) -> (arg*ts)) >>> f
282

ross's avatar
ross committed
283
dsCmd ids local_vars env_ids stack res_ty
284 285 286
	(HsArrApp arrow arg arrow_ty HsFirstOrderApp _)= do
    let
        (a_arg_ty, _res_ty') = tcSplitAppTy arrow_ty
ross's avatar
ross committed
287
        (_a_ty, arg_ty) = tcSplitAppTy a_arg_ty
288 289 290 291 292 293 294 295 296 297 298 299
    core_arrow <- dsLExpr arrow
    core_arg   <- dsLExpr arg
    stack_ids  <- mapM newSysLocalDs stack
    core_make_arg <- matchEnvStack env_ids stack_ids
                      (foldl mkCorePairExpr core_arg (map Var stack_ids))
    return (do_map_arrow ids
              (envStackType env_ids stack)
              arg_ty
              res_ty
              core_make_arg
              core_arrow,
               exprFreeVars core_arg `intersectVarSet` local_vars)
300

ross's avatar
ross committed
301
--	A, xs |- f :: a (t*ts) t'
302
--	A, xs |- arg :: t
ross's avatar
ross committed
303 304 305 306
--	------------------------------
--	A | xs |- f -<< arg :: [ts] t'
--
--		---> arr (\ ((xs)*ts) -> (f,(arg*ts))) >>> app
307

ross's avatar
ross committed
308
dsCmd ids local_vars env_ids stack res_ty
309 310 311
	(HsArrApp arrow arg arrow_ty HsHigherOrderApp _) = do
    let
        (a_arg_ty, _res_ty') = tcSplitAppTy arrow_ty
ross's avatar
ross committed
312
        (_a_ty, arg_ty) = tcSplitAppTy a_arg_ty
313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328
    
    core_arrow <- dsLExpr arrow
    core_arg   <- dsLExpr arg
    stack_ids  <- mapM newSysLocalDs stack
    core_make_pair <- matchEnvStack env_ids stack_ids
          (mkCorePairExpr core_arrow
             (foldl mkCorePairExpr core_arg (map Var stack_ids)))
                             
    return (do_map_arrow ids
              (envStackType env_ids stack)
              (mkCorePairTy arrow_ty arg_ty)
              res_ty
              core_make_pair
              (do_app ids arg_ty res_ty),
            (exprFreeVars core_arrow `unionVarSet` exprFreeVars core_arg)
              `intersectVarSet` local_vars)
329

ross's avatar
ross committed
330 331 332 333 334 335 336
--	A | ys |- c :: [t:ts] t'
--	A, xs  |- e :: t
--	------------------------
--	A | xs |- c e :: [ts] t'
--
--		---> arr (\ ((xs)*ts) -> let z = e in (((ys),z)*ts)) >>> c

337 338
dsCmd ids local_vars env_ids stack res_ty (HsApp cmd arg) = do
    core_arg <- dsLExpr arg
ross's avatar
ross committed
339
    let
340 341 342 343 344 345
        arg_ty = exprType core_arg
        stack' = arg_ty:stack
    (core_cmd, free_vars, env_ids')
             <- dsfixCmd ids local_vars stack' res_ty cmd
    stack_ids <- mapM newSysLocalDs stack
    arg_id <- newSysLocalDs arg_ty
ross's avatar
ross committed
346 347
    -- push the argument expression onto the stack
    let
348 349
        core_body = bindNonRec arg_id core_arg
                        (buildEnvStack env_ids' (arg_id:stack_ids))
ross's avatar
ross committed
350
    -- match the environment and stack against the input
351 352 353 354 355 356 357 358 359
    core_map <- matchEnvStack env_ids stack_ids core_body
    return (do_map_arrow ids
                      (envStackType env_ids stack)
                      (envStackType env_ids' stack')
                      res_ty
                      core_map
                      core_cmd,
      (exprFreeVars core_arg `intersectVarSet` local_vars)
              `unionVarSet` free_vars)
ross's avatar
ross committed
360

361 362 363 364 365 366 367
--	A | ys |- c :: [ts] t'
--	-----------------------------------------------
--	A | xs |- \ p1 ... pk -> c :: [t1:...:tk:ts] t'
--
--		---> arr (\ ((((xs), p1), ... pk)*ts) -> ((ys)*ts)) >>> c

dsCmd ids local_vars env_ids stack res_ty
368 369 370 371 372 373 374
    (HsLam (MatchGroup [L _ (Match pats _ (GRHSs [L _ (GRHS [] body)] _ ))] _)) = do
    let
        pat_vars = mkVarSet (collectPatsBinders pats)
        local_vars' = local_vars `unionVarSet` pat_vars
        stack' = drop (length pats) stack
    (core_body, free_vars, env_ids') <- dsfixCmd ids local_vars' stack' res_ty body
    stack_ids <- mapM newSysLocalDs stack
375 376 377 378 379 380

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

    let
        (actual_ids, stack_ids') = splitAt (length pats) stack_ids
381 382 383 384 385 386
        -- build a new environment, plus what's left of the stack
        core_expr = buildEnvStack env_ids' stack_ids'
        in_ty = envStackType env_ids stack
        in_ty' = envStackType env_ids' stack'
    
    fail_expr <- mkFailExpr LambdaExpr in_ty'
387
    -- match the patterns against the top of the old stack
388
    match_code <- matchSimplys (map Var actual_ids) LambdaExpr pats core_expr fail_expr
389
    -- match the old environment and stack against the input
390 391 392
    select_code <- matchEnvStack env_ids stack_ids match_code
    return (do_map_arrow ids in_ty in_ty' res_ty select_code core_body,
            free_vars `minusVarSet` pat_vars)
393 394

dsCmd ids local_vars env_ids stack res_ty (HsPar cmd)
395
  = dsLCmd ids local_vars env_ids stack res_ty cmd
396

ross's avatar
ross committed
397 398 399 400 401 402 403 404 405 406
--	A, xs |- e :: Bool
--	A | xs1 |- c1 :: [ts] t
--	A | xs2 |- c2 :: [ts] t
--	----------------------------------------
--	A | xs |- if e then c1 else c2 :: [ts] t
--
--		---> arr (\ ((xs)*ts) ->
--			if e then Left ((xs1)*ts) else Right ((xs2)*ts)) >>>
--		     c1 ||| c2

407
dsCmd ids local_vars env_ids stack res_ty (HsIf mb_fun cond then_cmd else_cmd) = do
408 409 410 411 412 413 414
    core_cond <- dsLExpr cond
    (core_then, fvs_then, then_ids) <- dsfixCmd ids local_vars stack res_ty then_cmd
    (core_else, fvs_else, else_ids) <- dsfixCmd ids local_vars stack res_ty else_cmd
    stack_ids  <- mapM newSysLocalDs stack
    either_con <- dsLookupTyCon eitherTyConName
    left_con   <- dsLookupDataCon leftDataConName
    right_con  <- dsLookupDataCon rightDataConName
415 416 417

    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]
418 419 420 421 422 423

        in_ty = envStackType env_ids stack
        then_ty = envStackType then_ids stack
        else_ty = envStackType else_ids stack
        sum_ty = mkTyConApp either_con [then_ty, else_ty]
        fvs_cond = exprFreeVars core_cond `intersectVarSet` local_vars
424 425 426 427 428 429 430 431 432 433 434
        
        core_left  = mk_left_expr  then_ty else_ty (buildEnvStack then_ids stack_ids)
        core_right = mk_right_expr then_ty else_ty (buildEnvStack else_ids stack_ids)

    core_if <- case mb_fun of 
       Just fun -> do { core_fun <- dsExpr fun
                      ; matchEnvStack env_ids stack_ids $
                        mkCoreApps core_fun [core_cond, core_left, core_right] }
       Nothing  -> matchEnvStack env_ids stack_ids $
                   mkIfThenElse core_cond core_left core_right

435 436 437 438
    return (do_map_arrow ids in_ty sum_ty res_ty
                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
439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456
\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

	arr (\ ((xs)*ts) -> case e of
		p1 -> (Left (Left (xs1)*ts))
		p2 -> Left ((Right (xs2)*ts))
		p3 -> Right ((xs3)*ts)) >>>
	(c1 ||| c2) ||| c3

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.
457
To build all this, we use triples describing segments of the list of
ross's avatar
ross committed
458
case bodies, containing the following fields:
459
 * a list of expressions of the form (Left|Right)* ((xs)*ts), to be put
ross's avatar
ross committed
460
   into the case replacing the commands
461
 * a sum type that is the common type of these expressions, and also the
ross's avatar
ross committed
462
   input type of the arrow
463
 * a CoreExpr for an arrow built by combining the translated command
ross's avatar
ross committed
464 465 466
   bodies with |||.

\begin{code}
467 468
dsCmd ids local_vars env_ids stack res_ty (HsCase exp (MatchGroup matches match_ty)) = do
    stack_ids <- mapM newSysLocalDs stack
469 470 471 472 473 474

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

    let
        leaves = concatMap leavesMatch matches
475
        make_branch (leaf, bound_vars) = do
476
            (core_leaf, _fvs, leaf_ids) <-
477
                  dsfixCmd ids (local_vars `unionVarSet` bound_vars) stack res_ty leaf
478
            return ([mkHsEnvStackExpr leaf_ids stack_ids],
479 480 481 482 483 484 485
                    envStackType leaf_ids stack,
                    core_leaf)
    
    branches <- mapM make_branch leaves
    either_con <- dsLookupTyCon eitherTyConName
    left_con <- dsLookupDataCon leftDataConName
    right_con <- dsLookupDataCon rightDataConName
486
    let
487 488 489 490
        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
491

492 493
        -- Prefix each tuple with a distinct series of Left's and Right's,
        -- in a balanced way, keeping track of the types.
494

495 496 497
        merge_branches (builds1, in_ty1, core_exp1)
                       (builds2, in_ty2, core_exp2)
          = (map (left_expr in_ty1 in_ty2) builds1 ++
498 499 500
                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)
501
        (leaves', sum_ty, core_choices) = foldb merge_branches branches
502 503 504 505 506 507 508 509 510 511 512 513 514 515 516

        -- 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
        in_ty = envStackType env_ids stack

        pat_ty    = funArgTy match_ty
        match_ty' = mkFunTy pat_ty sum_ty
        -- Note that we replace the HsCase result type by sum_ty,
        -- which is the type of matches'
    
    core_body <- dsExpr (HsCase exp (MatchGroup matches' match_ty'))
    core_matches <- matchEnvStack env_ids stack_ids core_body
    return (do_map_arrow ids in_ty sum_ty res_ty core_matches core_choices,
517
            exprFreeVars core_body `intersectVarSet` local_vars)
518 519 520 521 522 523 524

--	A | ys |- c :: [ts] t
--	----------------------------------
--	A | xs |- let binds in c :: [ts] t
--
--		---> arr (\ ((xs)*ts) -> let binds in ((ys)*ts)) >>> c

525 526
dsCmd ids local_vars env_ids stack res_ty (HsLet binds body) = do
    let
527
        defined_vars = mkVarSet (collectLocalBinders binds)
528 529
        local_vars' = local_vars `unionVarSet` defined_vars
    
Ian Lynagh's avatar
Ian Lynagh committed
530
    (core_body, _free_vars, env_ids') <- dsfixCmd ids local_vars' stack res_ty body
531
    stack_ids <- mapM newSysLocalDs stack
532
    -- build a new environment, plus the stack, using the let bindings
533
    core_binds <- dsLocalBinds binds (buildEnvStack env_ids' stack_ids)
534
    -- match the old environment and stack against the input
535 536 537 538 539 540 541 542
    core_map <- matchEnvStack env_ids stack_ids core_binds
    return (do_map_arrow ids
                        (envStackType env_ids stack)
                        (envStackType env_ids' stack)
                        res_ty
                        core_map
                        core_body,
        exprFreeVars core_binds `intersectVarSet` local_vars)
543

544 545
dsCmd ids local_vars env_ids [] res_ty (HsDo _ctxt stmts _)
  = dsCmdDo ids local_vars env_ids res_ty stmts 
546 547 548 549

--	A |- e :: forall e. a1 (e*ts1) t1 -> ... an (e*tsn) tn -> a (e*ts) t
--	A | xs |- ci :: [tsi] ti
--	-----------------------------------
ross's avatar
ross committed
550
--	A | xs |- (|e c1 ... cn|) :: [ts] t	---> e [t_xs] c1 ... cn
551

552 553 554 555 556 557
dsCmd _ids local_vars env_ids _stack _res_ty (HsArrForm op _ args) = do
    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)
558

559

560 561 562
dsCmd ids local_vars env_ids stack res_ty (HsTick ix vars expr) = do
    (expr1,id_set) <- dsLCmd ids local_vars env_ids stack res_ty expr
    expr2 <- mkTickBox ix vars expr1
563 564
    return (expr2,id_set)

Ian Lynagh's avatar
Ian Lynagh committed
565 566
dsCmd _ _ _ _ _ c = pprPanic "dsCmd" (ppr c)

567 568 569 570 571 572 573
--	A | ys |- c :: [ts] t	(ys <= xs)
--	---------------------
--	A | xs |- c :: [ts] t	---> arr_ts (\ (xs) -> (ys)) >>> c

dsTrimCmdArg
	:: IdSet		-- set of local vars available to this command
	-> [Id]			-- list of vars in the input to this command
574
	-> LHsCmdTop Id	-- command argument to desugar
575 576
	-> DsM (CoreExpr,	-- desugared expression
		IdSet)		-- set of local vars that occur free
577 578 579 580 581
dsTrimCmdArg local_vars env_ids (L _ (HsCmdTop cmd stack cmd_ty ids)) = do
    meth_ids <- mkCmdEnv ids
    (core_cmd, free_vars, env_ids') <- dsfixCmd meth_ids local_vars stack cmd_ty cmd
    stack_ids <- mapM newSysLocalDs stack
    trim_code <- matchEnvStack env_ids stack_ids (buildEnvStack env_ids' stack_ids)
582
    let
583 584 585 586 587
        in_ty = envStackType env_ids stack
        in_ty' = envStackType env_ids' stack
        arg_code = if env_ids' == env_ids then core_cmd else
                do_map_arrow meth_ids in_ty in_ty' cmd_ty trim_code core_cmd
    return (bindCmdEnv meth_ids arg_code, free_vars)
588 589 590 591 592 593 594 595 596

-- Given A | xs |- c :: [ts] t, builds c with xs fed back.
-- Typically needs to be prefixed with arr (\p -> ((xs)*ts))

dsfixCmd
	:: DsCmdEnv		-- arrow combinators
	-> IdSet		-- set of local vars available to this command
	-> [Type]		-- type of the stack
	-> Type			-- return type of the command
597
	-> LHsCmd Id		-- command to desugar
598 599 600 601
	-> DsM (CoreExpr,	-- desugared expression
		IdSet,		-- set of local vars that occur free
		[Id])		-- set as a list, fed back
dsfixCmd ids local_vars stack cmd_ty cmd
602 603 604
  = fixDs (\ ~(_,_,env_ids') -> do
        (core_cmd, free_vars) <- dsLCmd ids local_vars env_ids' stack cmd_ty cmd
        return (core_cmd, free_vars, varSetElems free_vars))
605 606 607 608 609 610 611 612 613 614 615 616 617 618 619

\end{code}

Translation of command judgements of the form

	A | xs |- do { ss } :: [] t

\begin{code}

dsCmdDo :: DsCmdEnv		-- arrow combinators
	-> IdSet		-- set of local vars available to this statement
	-> [Id]			-- list of vars in the input to this statement
				-- This is typically fed back,
				-- so don't pull on it too early
	-> Type			-- return type of the statement
620
	-> [LStmt Id]		-- statements to desugar
621 622 623 624 625 626 627
	-> DsM (CoreExpr,	-- desugared expression
		IdSet)		-- set of local vars that occur free

--	A | xs |- c :: [] t
--	--------------------------
--	A | xs |- do { c } :: [] t

628 629 630
dsCmdDo _ _ _ _ [] = panic "dsCmdDo"

dsCmdDo ids local_vars env_ids res_ty [L _ (LastStmt body _)]
631
  = dsLCmd ids local_vars env_ids [] res_ty body
632

633
dsCmdDo ids local_vars env_ids res_ty (stmt:stmts) = do
634
    let
635
        bound_vars = mkVarSet (collectLStmtBinders stmt)
636
        local_vars' = local_vars `unionVarSet` bound_vars
Ian Lynagh's avatar
Ian Lynagh committed
637
    (core_stmts, _, env_ids') <- fixDs (\ ~(_,_,env_ids') -> do
638
        (core_stmts, fv_stmts) <- dsCmdDo ids local_vars' env_ids' res_ty stmts 
639 640 641 642 643 644 645 646 647
        return (core_stmts, fv_stmts, varSetElems fv_stmts))
    (core_stmt, fv_stmt) <- dsCmdLStmt ids local_vars env_ids env_ids' stmt
    return (do_compose ids
                (mkBigCoreVarTupTy env_ids)
                (mkBigCoreVarTupTy env_ids')
                res_ty
                core_stmt
                core_stmts,
              fv_stmt)
648

ross's avatar
ross committed
649 650 651 652 653
\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}
Ian Lynagh's avatar
Ian Lynagh committed
654 655
dsCmdLStmt :: DsCmdEnv -> IdSet -> [Id] -> [Id] -> LStmt Id
           -> DsM (CoreExpr, IdSet)
656 657
dsCmdLStmt ids local_vars env_ids out_ids cmd
  = dsCmdStmt ids local_vars env_ids out_ids (unLoc cmd)
ross's avatar
ross committed
658

659 660 661 662 663 664 665
dsCmdStmt
	:: DsCmdEnv		-- arrow combinators
	-> IdSet		-- set of local vars available to this statement
	-> [Id]			-- list of vars in the input to this statement
				-- This is typically fed back,
				-- so don't pull on it too early
	-> [Id]			-- list of vars in the output of this statement
666
	-> Stmt Id	-- statement to desugar
667 668 669 670
	-> DsM (CoreExpr,	-- desugared expression
		IdSet)		-- set of local vars that occur free

--	A | xs1 |- c :: [] t
ross's avatar
ross committed
671
--	A | xs' |- do { ss } :: [] t'
672
--	------------------------------
ross's avatar
ross committed
673
--	A | xs |- do { c; ss } :: [] t'
674 675 676 677
--
--		---> arr (\ (xs) -> ((xs1),(xs'))) >>> first c >>>
--			arr snd >>> ss

678
dsCmdStmt ids local_vars env_ids out_ids (ExprStmt cmd _ _ c_ty) = do
679 680 681
    (core_cmd, fv_cmd, env_ids1) <- dsfixCmd ids local_vars [] c_ty cmd
    core_mux <- matchEnvStack env_ids []
        (mkCorePairExpr (mkBigCoreVarTup env_ids1) (mkBigCoreVarTup out_ids))
682
    let
683 684 685
	in_ty = mkBigCoreVarTupTy env_ids
	in_ty1 = mkBigCoreVarTupTy env_ids1
	out_ty = mkBigCoreVarTupTy out_ids
686 687
	before_c_ty = mkCorePairTy in_ty1 out_ty
	after_c_ty = mkCorePairTy c_ty out_ty
688 689
    snd_fn <- mkSndExpr c_ty out_ty
    return (do_map_arrow ids in_ty before_c_ty out_ty core_mux $
690 691 692
		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,
693
	      extendVarSetList fv_cmd out_ids)
694 695 696
  where

--	A | xs1 |- c :: [] t
ross's avatar
ross committed
697
--	A | xs' |- do { ss } :: [] t'		xs2 = xs' - defs(p)
698
--	-----------------------------------
ross's avatar
ross committed
699
--	A | xs |- do { p <- c; ss } :: [] t'
700 701 702 703 704 705 706
--
--		---> arr (\ (xs) -> ((xs1),(xs2))) >>> first c >>>
--			arr (\ (p, (xs2)) -> (xs')) >>> ss
--
-- It would be simpler and more consistent to do this using second,
-- but that's likely to be defined in terms of first.

707 708
dsCmdStmt ids local_vars env_ids out_ids (BindStmt pat cmd _ _) = do
    (core_cmd, fv_cmd, env_ids1) <- dsfixCmd ids local_vars [] (hsLPatType pat) cmd
709
    let
710
	pat_ty = hsLPatType pat
711 712
	pat_vars = mkVarSet (collectPatBinders pat)
	env_ids2 = varSetElems (mkVarSet out_ids `minusVarSet` pat_vars)
713
	env_ty2 = mkBigCoreVarTupTy env_ids2
714 715 716 717

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

718
    core_mux <- matchEnvStack env_ids []
719
	(mkCorePairExpr (mkBigCoreVarTup env_ids1) (mkBigCoreVarTup env_ids2))
720 721 722 723

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

724 725
    env_id <- newSysLocalDs env_ty2
    uniqs <- newUniqueSupply
726
    let
ross's avatar
ross committed
727
	after_c_ty = mkCorePairTy pat_ty env_ty2
728 729
	out_ty = mkBigCoreVarTupTy out_ids
	body_expr = coreCaseTuple uniqs env_id env_ids2 (mkBigCoreVarTup out_ids)
730 731 732 733 734
    
    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
735
    let
ross's avatar
ross committed
736
	proj_expr = Lam pair_id (coreCasePair pair_id pat_id env_id match_code)
737

ross's avatar
ross committed
738
    -- put it all together
739
    let
740 741 742
	in_ty = mkBigCoreVarTupTy env_ids
	in_ty1 = mkBigCoreVarTupTy env_ids1
	in_ty2 = mkBigCoreVarTupTy env_ids2
743
	before_c_ty = mkCorePairTy in_ty1 in_ty2
744
    return (do_map_arrow ids in_ty before_c_ty out_ty core_mux $
745 746 747 748 749 750 751 752 753 754 755
		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))

--	A | xs' |- do { ss } :: [] t
--	--------------------------------------
--	A | xs |- do { let binds; ss } :: [] t
--
--		---> arr (\ (xs) -> let binds in (xs')) >>> ss

756
dsCmdStmt ids local_vars env_ids out_ids (LetStmt binds) = do
757
    -- build a new environment using the let bindings
758
    core_binds <- dsLocalBinds binds (mkBigCoreVarTup out_ids)
759
    -- match the old environment against the input
760 761
    core_map <- matchEnvStack env_ids [] core_binds
    return (do_arr ids
762 763
			(mkBigCoreVarTupTy env_ids)
			(mkBigCoreVarTupTy out_ids)
764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780
			core_map,
	exprFreeVars core_binds `intersectVarSet` local_vars)

--	A | ys |- do { ss; returnA -< ((xs1), (ys2)) } :: [] ...
--	A | xs' |- do { ss' } :: [] t
--	------------------------------------
--	A | xs |- do { rec ss; ss' } :: [] t
--
--			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'

781 782
dsCmdStmt ids local_vars env_ids out_ids 
          (RecStmt { recS_stmts = stmts, recS_later_ids = later_ids, recS_rec_ids = rec_ids
783 784
                   , recS_rec_rets = rhss }) = do
    let
785 786 787
        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
788 789 790

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

791 792
    uniqs <- newUniqueSupply
    env2_id <- newSysLocalDs env2_ty
ross's avatar
ross committed
793
    let
794 795 796 797 798
        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)

    post_loop_fn <- matchEnvStack later_ids [env2_id] post_loop_body
ross's avatar
ross committed
799 800 801

    --- loop (...)

802 803
    (core_loop, env1_id_set, env1_ids)
               <- dsRecCmd ids local_vars stmts later_ids rec_ids rhss
ross's avatar
ross committed
804 805 806 807

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

    let
808 809 810 811
        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
812

813
    pre_loop_fn <- matchEnvStack env_ids [] pre_loop_body
ross's avatar
ross committed
814 815 816 817

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

    let
818 819 820 821 822 823 824 825 826 827 828
        env_ty = mkBigCoreVarTupTy env_ids
        out_ty = mkBigCoreVarTupTy out_ids
        core_body = do_map_arrow ids env_ty pre_pair_ty out_ty
                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
829

Ian Lynagh's avatar
Ian Lynagh committed
830 831
dsCmdStmt _ _ _ _ s = pprPanic "dsCmdStmt" (ppr s)

ross's avatar
ross committed
832 833 834 835
--	loop (arr (\ ((env1_ids), ~(rec_ids)) -> (env_ids)) >>>
--	      ss >>>
--	      arr (\ (out_ids) -> ((later_ids),(rhss))) >>>

Ian Lynagh's avatar
Ian Lynagh committed
836 837
dsRecCmd :: DsCmdEnv -> VarSet -> [LStmt Id] -> [Var] -> [Var] -> [HsExpr Id]
         -> DsM (CoreExpr, VarSet, [Var])
838 839 840 841 842 843
dsRecCmd ids local_vars stmts later_ids rec_ids rhss = do
    let
        rec_id_set = mkVarSet rec_ids
        out_ids = varSetElems (mkVarSet later_ids `unionVarSet` rec_id_set)
        out_ty = mkBigCoreVarTupTy out_ids
        local_vars' = local_vars `unionVarSet` rec_id_set
844 845 846

    -- mk_pair_fn = \ (out_ids) -> ((later_ids),(rhss))

847
    core_rhss <- mapM dsExpr rhss
848
    let
849 850 851 852 853 854 855 856
        later_tuple = mkBigCoreVarTup later_ids
        later_ty = mkBigCoreVarTupTy later_ids
        rec_tuple = mkBigCoreTup core_rhss
        rec_ty = mkBigCoreVarTupTy rec_ids
        out_pair = mkCorePairExpr later_tuple rec_tuple
        out_pair_ty = mkCorePairTy later_ty rec_ty

    mk_pair_fn <- matchEnvStack out_ids [] out_pair
857

ross's avatar
ross committed
858 859
    -- ss

860
    (core_stmts, fv_stmts, env_ids) <- dsfixCmdStmts ids local_vars' out_ids stmts
861

ross's avatar
ross committed
862
    -- squash_pair_fn = \ ((env1_ids), ~(rec_ids)) -> (env_ids)
863

864
    rec_id <- newSysLocalDs rec_ty
865
    let
866 867 868 869 870 871 872 873 874 875 876 877
        env1_id_set = fv_stmts `minusVarSet` rec_id_set
        env1_ids = varSetElems env1_id_set
        env1_ty = mkBigCoreVarTupTy env1_ids
        in_pair_ty = mkCorePairTy env1_ty rec_ty
        core_body = mkBigCoreTup (map selectVar env_ids)
          where
            selectVar v
                | v `elemVarSet` rec_id_set
                  = mkTupleSelector rec_ids v rec_id (Var rec_id)
                | otherwise = Var v

    squash_pair_fn <- matchEnvStack env1_ids [rec_id] core_body
878 879 880 881

    -- loop (arr squash_pair_fn >>> ss >>> arr mk_pair_fn)

    let
882 883 884 885 886 887 888 889 890
        env_ty = mkBigCoreVarTupTy env_ids
        core_loop = do_loop ids env1_ty later_ty rec_ty
                (do_map_arrow ids in_pair_ty env_ty out_pair_ty
                        squash_pair_fn
                        (do_compose ids env_ty out_ty out_pair_ty
                                core_stmts
                                (do_arr ids out_ty out_pair_ty mk_pair_fn)))

    return (core_loop, env1_id_set, env1_ids)
891 892

\end{code}
ross's avatar
ross committed
893
A sequence of statements (as in a rec) is desugared to an arrow between
894 895 896 897 898 899 900
two environments
\begin{code}

dsfixCmdStmts
	:: DsCmdEnv		-- arrow combinators
	-> IdSet		-- set of local vars available to this statement
	-> [Id]			-- output vars of these statements
901
	-> [LStmt Id]	-- statements to desugar
902 903 904 905 906
	-> DsM (CoreExpr,	-- desugared expression
		IdSet,		-- set of local vars that occur free
		[Id])		-- input vars

dsfixCmdStmts ids local_vars out_ids stmts
907 908 909
  = fixDs (\ ~(_,_,env_ids) -> do
        (core_stmts, fv_stmts) <- dsCmdStmts ids local_vars env_ids out_ids stmts
	return (core_stmts, fv_stmts, varSetElems fv_stmts))
910 911 912 913 914 915

dsCmdStmts
	:: DsCmdEnv		-- arrow combinators
	-> IdSet		-- set of local vars available to this statement
	-> [Id]			-- list of vars in the input to these statements
	-> [Id]			-- output vars of these statements
916
	-> [LStmt Id]	-- statements to desugar
917 918 919 920
	-> DsM (CoreExpr,	-- desugared expression
		IdSet)		-- set of local vars that occur free

dsCmdStmts ids local_vars env_ids out_ids [stmt]
921
  = dsCmdLStmt ids local_vars env_ids out_ids stmt
922

923 924
dsCmdStmts ids local_vars env_ids out_ids (stmt:stmts) = do
    let
925
        bound_vars = mkVarSet (collectLStmtBinders stmt)
926
        local_vars' = local_vars `unionVarSet` bound_vars
Ian Lynagh's avatar
Ian Lynagh committed
927
    (core_stmts, _fv_stmts, env_ids') <- dsfixCmdStmts ids local_vars' out_ids stmts
928 929 930 931 932 933 934 935
    (core_stmt, fv_stmt) <- dsCmdLStmt ids local_vars env_ids env_ids' stmt
    return (do_compose ids
                (mkBigCoreVarTupTy env_ids)
                (mkBigCoreVarTupTy env_ids')
                (mkBigCoreVarTupTy out_ids)
                core_stmt
                core_stmts,
              fv_stmt)
936

Ian Lynagh's avatar
Ian Lynagh committed
937 938
dsCmdStmts _ _ _ _ [] = panic "dsCmdStmts []"

939 940 941 942 943
\end{code}

Match a list of expressions against a list of patterns, left-to-right.

\begin{code}
944 945 946 947 948
matchSimplys :: [CoreExpr]              -- Scrutinees
	     -> HsMatchContext Name	-- Match kind
	     -> [LPat Id]         	-- Patterns they should match
	     -> CoreExpr                -- Return this if they all match
	     -> CoreExpr                -- Return this if they don't
949
	     -> DsM CoreExpr
950 951 952
matchSimplys [] _ctxt [] result_expr _fail_expr = return result_expr
matchSimplys (exp:exps) ctxt (pat:pats) result_expr fail_expr = do
    match_code <- matchSimplys exps ctxt pats result_expr fail_expr
953
    matchSimply exp ctxt pat match_code fail_expr
Ian Lynagh's avatar
Ian Lynagh committed
954
matchSimplys _ _ _ _ _ = panic "matchSimplys"
955 956
\end{code}

ross's avatar
ross committed
957
List of leaf expressions, with set of variables bound in each
958

ross's avatar
ross committed
959
\begin{code}
960
leavesMatch :: LMatch Id -> [(LHsExpr Id, IdSet)]
961
leavesMatch (L _ (Match pats _ (GRHSs grhss binds)))
962
  = let
963 964
	defined_vars = mkVarSet (collectPatsBinders pats)
			`unionVarSet`
965
		       mkVarSet (collectLocalBinders binds)
966
    in
967
    [(expr, 
968
      mkVarSet (collectLStmtsBinders stmts) 
969
	`unionVarSet` defined_vars) 
970
    | L _ (GRHS stmts expr) <- grhss]
ross's avatar
ross committed
971
\end{code}
972

ross's avatar
ross committed
973
Replace the leaf commands in a match
974

ross's avatar
ross committed
975
\begin{code}
976 977
replaceLeavesMatch
	:: Type			-- new result type
978 979 980 981
	-> [LHsExpr Id]	-- replacement leaf expressions of that type
	-> LMatch Id	-- the matches of a case command
	-> ([LHsExpr Id],-- remaining leaf expressions
	    LMatch Id)	-- updated match
Ian Lynagh's avatar
Ian Lynagh committed
982
replaceLeavesMatch _res_ty leaves (L loc (Match pat mt (GRHSs grhss binds)))
983
  = let
ross's avatar
ross committed
984
	(leaves', grhss') = mapAccumL replaceLeavesGRHS leaves grhss
985
    in
986
    (leaves', L loc (Match pat mt (GRHSs grhss' binds)))
987 988

replaceLeavesGRHS
989 990 991 992
	:: [LHsExpr Id]	-- replacement leaf expressions of that type
	-> LGRHS Id	-- rhss of a case command
	-> ([LHsExpr Id],-- remaining leaf expressions
	    LGRHS Id)	-- updated GRHS
Ian Lynagh's avatar
Ian Lynagh committed
993
replaceLeavesGRHS (leaf:leaves) (L loc (GRHS stmts _))
994
  = (leaves, L loc (GRHS stmts leaf))
Ian Lynagh's avatar
Ian Lynagh committed
995
replaceLeavesGRHS [] _ = panic "replaceLeavesGRHS []"
996 997 998 999 1000 1001
\end{code}

Balanced fold of a non-empty list.

\begin{code}
foldb :: (a -> a -> a) -> [a] -> a
ross's avatar
ross committed
1002 1003
foldb _ [] = error "foldb of empty list"
foldb _ [x] = x
1004 1005 1006 1007 1008 1009
foldb f xs = foldb f (fold_pairs xs)
  where
    fold_pairs [] = []
    fold_pairs [x] = [x]
    fold_pairs (x1:x2:xs) = f x1 x2:fold_pairs xs
\end{code}
Ross Paterson's avatar
Ross Paterson committed
1010

1011 1012
Note [Dictionary binders in ConPatOut] See also same Note in HsUtils
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Ross Paterson's avatar
Ross Paterson committed
1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031
The following functions to collect value variables from patterns are
copied from HsUtils, with one change: we also collect the dictionary
bindings (pat_binds) from ConPatOut.  We need them for cases like

h :: Arrow a => Int -> a (Int,Int) Int
h x = proc (y,z) -> case compare x y of
                GT -> returnA -< z+x

The type checker turns the case into

                case compare x y of
                  GT { p77 = plusInt } -> returnA -< p77 z x

Here p77 is a local binding for the (+) operation.

See comments in HsUtils for why the other version does not include
these bindings.

\begin{code}
1032
collectPatBinders :: LPat Id -> [Id]
1033
collectPatBinders pat = collectl pat []
Ross Paterson's avatar
Ross Paterson committed
1034

1035
collectPatsBinders :: [LPat Id] -> [Id]
1036
collectPatsBinders pats = foldr collectl [] pats
Ross Paterson's avatar
Ross Paterson committed
1037 1038

---------------------
1039
collectl :: LPat Id -> [Id] -> [Id]
1040 1041
-- See Note [Dictionary binders in ConPatOut]
collectl (L _ pat) bndrs
Ross Paterson's avatar
Ross Paterson committed
1042 1043
  = go pat
  where
1044
    go (VarPat var)               = var : bndrs
Ross Paterson's avatar
Ross Paterson committed
1045 1046 1047
    go (WildPat _)                = bndrs
    go (LazyPat pat)              = collectl pat bndrs
    go (BangPat pat)              = collectl pat bndrs
1048
    go (AsPat (L _ a) pat)        = a : collectl pat bndrs
Ross Paterson's avatar
Ross Paterson committed
1049 1050 1051 1052 1053 1054
    go (ParPat  pat)              = collectl pat bndrs

    go (ListPat pats _)           = foldr collectl bndrs pats
    go (PArrPat pats _)           = foldr collectl bndrs pats
    go (TuplePat pats _ _)        = foldr collectl bndrs pats

Ian Lynagh's avatar
Ian Lynagh committed
1055
    go (ConPatIn _ ps)            = foldr collectl bndrs (hsConPatArgs ps)
Ross Paterson's avatar
Ross Paterson committed
1056
    go (ConPatOut {pat_args=ps, pat_binds=ds}) =
1057
                                    collectEvBinders ds
Ross Paterson's avatar
Ross Paterson committed
1058 1059
                                    ++ foldr collectl bndrs (hsConPatArgs ps)
    go (LitPat _)                 = bndrs
1060
    go (NPat _ _ _)               = bndrs
1061
    go (NPlusKPat (L _ n) _ _ _)  = n : bndrs
Ross Paterson's avatar
Ross Paterson committed
1062 1063 1064

    go (SigPatIn pat _)           = collectl pat bndrs
    go (SigPatOut pat _)          = collectl pat bndrs
Ian Lynagh's avatar
Ian Lynagh committed
1065 1066
    go (TypePat _)                = bndrs
    go (CoPat _ pat _)            = collectl (noLoc pat) bndrs
1067 1068
    go (ViewPat _ pat _)          = collectl pat bndrs
    go p@(QuasiQuotePat {})       = pprPanic "collectl/go" (ppr p)
1069 1070 1071 1072 1073 1074 1075 1076 1077

collectEvBinders :: TcEvBinds -> [Id]
collectEvBinders (EvBinds bs)   = foldrBag add_ev_bndr [] bs
collectEvBinders (TcEvBinds {}) = panic "ToDo: collectEvBinders"

add_ev_bndr :: EvBind -> [Id] -> [Id]
add_ev_bndr (EvBind b _) bs | isId b    = b:bs
                            | otherwise = bs
  -- A worry: what about coercion variable binders??
Ross Paterson's avatar
Ross Paterson committed
1078
\end{code}