+-- 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]) }
+
+ where
+ -- In a monad comprehension expression, pattern-match failure just calls
+ -- the monadic `fail` rather than throwing an exception
+ handle_failure pat match fail_op
+ | matchCanFail match
+ = do { fail_op' <- dsExpr fail_op
+ ; fail_msg <- mkStringExpr (mk_fail_msg pat)
+ ; extractMatchResult match (App fail_op' fail_msg) }
+ | otherwise
+ = extractMatchResult match (error "It can't fail")
+
+ mk_fail_msg :: Located e -> String
+ mk_fail_msg pat = "Pattern match failure in monad comprehension at " ++
+ showSDoc (ppr (getLoc pat))
+
+-- Desugar nested monad comprehensions, for example in `then..` constructs
+-- dsInnerMonadComp quals [a,b,c] ret_op
+-- returns the desugaring of
+-- [ (a,b,c) | quals ]
+
+dsInnerMonadComp :: [LStmt Id]
+ -> [Id] -- Return a tuple of these variables
+ -> LHsExpr Id -- The monomorphic "return" operator
+ -> DsM CoreExpr
+dsInnerMonadComp stmts bndrs ret_op
+ = dsMcStmts (stmts ++ [noLoc (ReturnStmt (mkBigLHsVarTup bndrs) ret_op)])
+
+-- The `unzip` function for `GroupStmt` in a monad comprehensions
+--
+-- unzip :: m (a,b,..) -> (m a,m b,..)
+-- unzip m_tuple = ( liftM selN1 m_tuple
+-- , liftM selN2 m_tuple
+-- , .. )
+--
+-- mkMcUnzipM m [t1, t2]
+-- = (unzip_fn, \ys :: m (t1, t2) ->
+-- ( liftM (selN1 :: (t1, t2) -> t1) ys
+-- , liftM (selN2 :: (t1, t2) -> t2) ys
+-- ))
+--
+mkMcUnzipM :: CoreExpr
+ -> Type -- m
+ -> [Type] -- [a,b,c,..]
+ -> DsM (Id, CoreExpr)
+mkMcUnzipM liftM_op m_ty elt_tys
+ = do { ys <- newSysLocalDs monad_tuple_ty
+ ; xs <- mapM newSysLocalDs elt_tys
+ ; scrut <- newSysLocalDs tuple_tys
+
+ ; unzip_fn <- newSysLocalDs unzip_fn_ty
+
+ ; let -- Select one Id from our tuple
+ selectExpr n = mkLams [scrut] $ mkTupleSelector xs (xs !! n) scrut (Var scrut)
+ -- Apply 'selectVar' and 'ys' to 'liftM'
+ tupleElem n = mkApps liftM_op
+ -- Types (m is figured out by the type checker):
+ -- liftM :: forall a b. (a -> b) -> m a -> m b
+ [ Type tuple_tys, Type (elt_tys !! n)
+ -- Arguments:
+ , selectExpr n, Var ys ]
+ -- The final expression with the big tuple
+ unzip_body = mkBigCoreTup [ tupleElem n | n <- [0..length elt_tys - 1] ]
+
+ ; return (unzip_fn, mkLams [ys] unzip_body) }
+ where monad_tys = map (m_ty `mkAppTy`) elt_tys -- [m a,m b,m c,..]
+ tuple_monad_tys = mkBigCoreTupTy monad_tys -- (m a,m b,m c,..)
+ tuple_tys = mkBigCoreTupTy elt_tys -- (a,b,c,..)
+ monad_tuple_ty = m_ty `mkAppTy` tuple_tys -- m (a,b,c,..)
+ unzip_fn_ty = monad_tuple_ty `mkFunTy` tuple_monad_tys -- m (a,b,c,..) -> (m a,m b,m c,..)
+
+-- Generate the `mzip` function for `ParStmt` in monad comprehensions, for
+-- example:
+--
+-- mzip :: m t1
+-- -> (m t2 -> m t3 -> m (t2, t3))
+-- -> m (t1, (t2, t3))
+--
+-- mkMcZipM m [t1, t2, t3]
+-- = (zip_fn, \(q1::t1) (q2::t2) (q3::t3) ->
+-- mzip q1 (mzip q2 q3))
+--
+mkMcZipM :: CoreExpr
+ -> Type
+ -> [Type]
+ -> DsM (Id, CoreExpr)
+
+mkMcZipM mzip_op m_ty tys@(_:_:_) -- min. 2 types
+ = do { (ids, t1, tuple_ty, zip_body) <- loop tys
+ ; zip_fn <- newSysLocalDs $
+ (m_ty `mkAppTy` t1)
+ `mkFunTy`
+ (m_ty `mkAppTy` tuple_ty)
+ `mkFunTy`
+ (m_ty `mkAppTy` mkBigCoreTupTy [t1, tuple_ty])
+ ; return (zip_fn, mkLams ids zip_body) }
+
+ where
+ -- loop :: [Type] -> DsM ([Id], Type, [Type], CoreExpr)
+ loop [t1, t2] = do -- last run of the `loop`
+ { ids@[a,b] <- newSysLocalsDs (map (m_ty `mkAppTy`) [t1,t2])
+ ; let zip_body = mkApps mzip_op [ Type t1, Type t2 , Var a, Var b ]
+ ; return (ids, t1, t2, zip_body) }
+
+ loop (t1:tr) = do
+ { -- Get ty, ids etc from the "inner" zip
+ (ids', t1', t2', zip_body') <- loop tr
+
+ ; a <- newSysLocalDs $ m_ty `mkAppTy` t1
+ ; let tuple_ty' = mkBigCoreTupTy [t1', t2']
+ zip_body = mkApps mzip_op [ Type t1, Type tuple_ty', Var a, zip_body' ]
+ ; return ((a:ids'), t1, tuple_ty', zip_body) }
+
+-- This case should never happen:
+mkMcZipM _ _ tys = pprPanic "mkMcZipM: unexpected argument" (ppr tys)
+