+-- Entry point for monad comprehension desugaring
+--
+dsMonadComp :: [LStmt Id] -- the statements
+ -> Type -- the final type
+ -> DsM CoreExpr
+dsMonadComp stmts res_ty
+ = dsMcStmts stmts (DsMonadComp (Left return_op) body m_ty)
+ where
+ (m_ty, _) = tcSplitAppTy res_ty
+
+
+dsMcStmts :: [LStmt Id]
+ -> DsMonadComp
+ -> DsM CoreExpr
+
+-- No statements left for desugaring. Desugar the body after calling "return"
+-- on it.
+dsMcStmts [] DsMonadComp { mc_return, mc_body }
+ = case mc_return of
+ Left ret -> dsLExpr $ noLoc ret `nlHsApp` mc_body
+ Right ret' -> do
+ { body' <- dsLExpr mc_body
+ ; return $ mkApps ret' [body'] }
+
+-- Otherwise desugar each statement step by step
+dsMcStmts ((L loc stmt) : lstmts) mc
+ = putSrcSpanDs loc (dsMcStmt stmt lstmts mc)
+
+
+dsMcStmt :: Stmt Id -> [LStmt Id] -> DsM CoreExpr
+
+dsMcStmt (LastStmt body ret_op) stmts
+ = ASSERT( null stmts )
+ do { body' <- dsLExpr body
+ ; ret_op' <- dsExpr ret_op
+ ; return (App ret_op' body') }
+
+-- [ .. | let binds, stmts ]
+dsMcStmt (LetStmt binds) stmts
+ = do { rest <- dsMcStmts stmts
+ ; dsLocalBinds binds rest }
+
+-- [ .. | a <- m, stmts ]
+dsMcStmt (BindStmt pat rhs bind_op fail_op) stmts
+ = do { rhs' <- dsLExpr rhs
+ ; dsMcBindStmt pat rhs' bind_op fail_op stmts }
+
+-- Apply `guard` to the `exp` expression
+--
+-- [ .. | exp, stmts ]
+--
+dsMcStmt (ExprStmt exp then_exp guard_exp _) stmts
+ = do { exp' <- dsLExpr exp
+ ; guard_exp' <- dsExpr guard_exp
+ ; then_exp' <- dsExpr then_exp
+ ; rest <- dsMcStmts stmts
+ ; return $ mkApps then_exp' [ mkApps guard_exp' [exp']
+ , rest ] }
+
+-- Transform statements desugar like this:
+--
+-- [ .. | qs, then f by e ] -> f (\q_v -> e) [| qs |]
+--
+-- where [| qs |] is the desugared inner monad comprehenion generated by the
+-- statements `qs`.
+dsMcStmt (TransformStmt stmts binders usingExpr maybeByExpr return_op bind_op) stmts_rest
+ = do { expr <- dsInnerMonadComp stmts binders return_op
+ ; let binders_tup_type = mkBigCoreTupTy $ map idType binders
+ ; usingExpr' <- dsLExpr usingExpr
+ ; using_args <- case maybeByExpr of
+ Nothing -> return [expr]
+ Just byExpr -> do
+ byExpr' <- dsLExpr byExpr
+ us <- newUniqueSupply
+ tup_binder <- newSysLocalDs binders_tup_type
+ let byExprWrapper = mkTupleCase us binders byExpr' tup_binder (Var tup_binder)
+ return [Lam tup_binder byExprWrapper, expr]
+
+ ; let pat = mkBigLHsVarPatTup binders
+ rhs = mkApps usingExpr' ((Type binders_tup_type) : using_args)
+
+ ; dsMcBindStmt pat rhs bind_op noSyntaxExpr stmts_rest }
+
+-- Group statements desugar like this:
+--
+-- [| (q, then group by e using f); rest |]
+-- ---> f {qt} (\qv -> e) [| q; return qv |] >>= \ n_tup ->
+-- case unzip n_tup of qv -> [| rest |]
+--
+-- where variables (v1:t1, ..., vk:tk) are bound by q
+-- qv = (v1, ..., vk)
+-- qt = (t1, ..., tk)
+-- (>>=) :: m2 a -> (a -> m3 b) -> m3 b
+-- f :: forall a. (a -> t) -> m1 a -> m2 (n a)
+-- n_tup :: n qt
+-- unzip :: n qt -> (n t1, ..., n tk) (needs Functor n)
+--
+-- [| q, then group by e using f |] -> (f (\q_v -> e) [| q |]) >>= (return . (unzip q_v))
+--
+-- which is equal to
+--
+-- [| q, then group by e using f |] -> liftM (unzip q_v) (f (\q_v -> e) [| q |])
+--
+-- where unzip is of the form
+--
+-- unzip :: n (a,b,c,..) -> (n a,n b,n c,..)
+-- unzip m_tuple = ( fmap selN1 m_tuple
+-- , fmap selN2 m_tuple
+-- , .. )
+-- where selN1 (a,b,c,..) = a
+-- selN2 (a,b,c,..) = b
+-- ..
+--
+dsMcStmt (GroupStmt stmts binderMap by using return_op bind_op fmap_op) stmts_rest
+ = do { let (fromBinders, toBinders) = unzip binderMap
+ fromBindersTypes = map idType fromBinders -- Types ty
+ fromBindersTupleTy = mkBigCoreTupTy fromBindersTypes
+ toBindersTypes = map idType toBinders -- Types (n ty)
+ toBindersTupleTy = mkBigCoreTupTy toBindersTypes
+
+ -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders
+ ; expr <- dsInnerMonadComp stmts fromBinders return_op
+
+ -- Work out what arguments should be supplied to that expression: i.e. is an extraction
+ -- function required? If so, create that desugared function and add to arguments
+ ; usingExpr' <- dsLExpr (either id noLoc using)
+ ; usingArgs <- case by of
+ Nothing -> return [expr]
+ Just by_e -> do { by_e' <- dsLExpr by_e
+ ; lam <- matchTuple fromBinders by_e'
+ ; return [lam, expr] }
+
+ -- Create an unzip function for the appropriate arity and element types
+ ; fmap_op' <- dsExpr fmap_op
+ ; (unzip_fn, unzip_rhs) <- mkMcUnzipM fmap_op' m_ty fromBindersTypes
+
+ -- Generate the expressions to build the grouped list
+ -- Build a pattern that ensures the consumer binds into the NEW binders,
+ -- which hold monads rather than single values
+ ; bind_op' <- dsExpr bind_op
+ ; let bind_ty = exprType bind_op' -- m2 (n (a,b,c)) -> (n (a,b,c) -> r1) -> r2
+ n_tup_ty = funArgTy $ funArgTy $ funResultTy bind_ty
+
+ ; body <- dsMcStmts stmts_rest
+ ; n_tup_var <- newSysLocalDs n_tup_ty
+ ; tup_n_var <- newSysLocalDs (mkBigCoreVarTupTy toBinders)
+ ; us <- newUniqueSupply
+ ; let unzip_n_tup = Let (Rec [(unzip_fn, unzip_rhs)]) $
+ App (Var unzip_fn) (Var n_tup_var)
+ -- unzip_n_tup :: (n a, n b, n c)
+ body' = mkTupleCase us toBinders body unzip_n_tup (Var tup_n_var)
+
+ ; return (mkApps bind_op' [rhs', Lam n_tup_var body']) }
+
+-- Parallel statements. Use `Control.Monad.Zip.mzip` to zip parallel
+-- statements, for example:
+--
+-- [ body | qs1 | qs2 | qs3 ]
+-- -> [ body | (bndrs1, (bndrs2, bndrs3))
+-- <- [bndrs1 | qs1] `mzip` ([bndrs2 | qs2] `mzip` [bndrs3 | qs3]) ]
+--
+-- where `mzip` has type
+-- mzip :: forall a b. m a -> m b -> m (a,b)
+-- NB: we need a polymorphic mzip because we call it several times
+
+dsMcStmt (ParStmt pairs mzip_op bind_op return_op) stmts_rest
+ = do { exps <- mapM ds_inner pairs
+ ; let qual_tys = map (mkBigCoreVarTupTy . snd) pairs
+ ; mzip_op' <- dsExpr mzip_op
+ ; (zip_fn, zip_rhs) <- mkMcZipM mzip_op' (mc_m_ty mc) qual_tys
+
+ ; let -- The pattern variables
+ vars = map (mkBigLHsVarPatTup . snd) pairs
+ -- Pattern with tuples of variables
+ -- [v1,v2,v3] => (v1, (v2, v3))
+ pat = foldr (\tn tm -> mkBigLHsPatTup [tn, tm]) (last vars) (init vars)
+ rhs = Let (Rec [(zip_fn, zip_rhs)]) (mkApps (Var zip_fn) exps)
+
+ ; dsMcBindStmt pat rhs bind_op noSyntaxExpr stmts_rest }
+ where
+ ds_inner (stmts, bndrs) = dsInnerMonadComp stmts bndrs mono_ret_op
+ where
+ mono_ret_op = HsWrap (WpTyApp (mkBigCoreVarTupTy bndrs)) return_op
+
+dsMcStmt stmt _ = pprPanic "dsMcStmt: unexpected stmt" (ppr stmt)
+
+
+matchTuple :: [Id] -> CoreExpr -> DsM CoreExpr
+-- (matchTuple [a,b,c] body)
+-- returns the Core term
+-- \x. case x of (a,b,c) -> body
+matchTuple ids body
+ = do { us <- newUniqueSupply
+ ; tup_id <- newSysLocalDs (mkBigLHsVarPatTup ids)
+ ; return (Lam tup_id $ mkTupleCase us ids body tup_id (Var tup_id)) }
+
+
+-- general `rhs' >>= \pat -> stmts` desugaring where `rhs'` is already a
+-- desugared `CoreExpr`
+dsMcBindStmt :: LPat Id
+ -> CoreExpr -- ^ the desugared rhs of the bind statement
+ -> SyntaxExpr Id
+ -> SyntaxExpr Id
+ -> [LStmt Id]
+ -> DsM CoreExpr
+dsMcBindStmt pat rhs' bind_op fail_op stmts
+ = do { body <- dsMcStmts stmts
+ ; bind_op' <- dsExpr bind_op
+ ; var <- selectSimpleMatchVarL pat
+ ; let bind_ty = exprType bind_op' -- rhs -> (pat -> res1) -> res2
+ res1_ty = funResultTy (funArgTy (funResultTy bind_ty))
+ ; match <- matchSinglePat (Var var) (StmtCtxt DoExpr) pat
+ res1_ty (cantFailMatchResult body)
+ ; match_code <- handle_failure pat match fail_op
+ ; return (mkApps bind_op' [rhs', Lam var match_code]) }