diff options
Diffstat (limited to 'compiler/deSugar/DsListComp.hs')
| -rw-r--r-- | compiler/deSugar/DsListComp.hs | 118 |
1 files changed, 62 insertions, 56 deletions
diff --git a/compiler/deSugar/DsListComp.hs b/compiler/deSugar/DsListComp.hs index 4d11fa21b8..d835995857 100644 --- a/compiler/deSugar/DsListComp.hs +++ b/compiler/deSugar/DsListComp.hs @@ -81,11 +81,13 @@ dsListComp lquals res_ty = do -- and the type of the elements that it outputs (tuples of binders) dsInnerListComp :: (ParStmtBlock Id Id) -> DsM (CoreExpr, Type) dsInnerListComp (ParStmtBlock stmts bndrs _) - = do { expr <- dsListComp (stmts ++ [noLoc $ mkLastStmt (mkBigLHsVarTupId bndrs)]) + = do { let bndrs_tuple_type = mkBigCoreVarTupTy bndrs + + -- really use original bndrs below! + ; expr <- dsListComp (stmts ++ [noLoc $ mkLastStmt (mkBigLHsVarTupId bndrs)]) (mkListTy bndrs_tuple_type) + ; return (expr, bndrs_tuple_type) } - where - bndrs_tuple_type = mkBigCoreVarTupTy bndrs -- This function factors out commonality between the desugaring strategies for GroupStmt. -- Given such a statement it gives you back an expression representing how to compute the transformed @@ -94,47 +96,50 @@ dsTransStmt :: ExprStmt Id -> DsM (CoreExpr, LPat Id) dsTransStmt (TransStmt { trS_form = form, trS_stmts = stmts, trS_bndrs = binderMap , trS_by = by, trS_using = using }) = do let (from_bndrs, to_bndrs) = unzip binderMap - from_bndrs_tys = map idType from_bndrs + + let from_bndrs_tys = map idType from_bndrs to_bndrs_tys = map idType to_bndrs + to_bndrs_tup_ty = mkBigCoreTupTy to_bndrs_tys -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders - (expr, from_tup_ty) <- dsInnerListComp (ParStmtBlock stmts from_bndrs noSyntaxExpr) + (expr', from_tup_ty) <- dsInnerListComp (ParStmtBlock stmts from_bndrs noSyntaxExpr) -- 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 using - usingArgs <- case by of - Nothing -> return [expr] - Just by_e -> do { by_e' <- dsLExpr by_e - ; lam <- matchTuple from_bndrs by_e' - ; return [lam, expr] } + usingArgs' <- case by of + Nothing -> return [expr'] + Just by_e -> do { by_e' <- dsLExpr by_e + ; lam' <- matchTuple from_bndrs by_e' + ; return [lam', expr'] } -- Create an unzip function for the appropriate arity and element types and find "map" - unzip_stuff <- mkUnzipBind form from_bndrs_tys + unzip_stuff' <- mkUnzipBind form from_bndrs_tys map_id <- dsLookupGlobalId mapName -- Generate the expressions to build the grouped list let -- First we apply the grouping function to the inner list - inner_list_expr = mkApps usingExpr' usingArgs + inner_list_expr' = mkApps usingExpr' usingArgs' -- Then we map our "unzip" across it to turn the lists of tuples into tuples of lists -- We make sure we instantiate the type variable "a" to be a list of "from" tuples and -- the "b" to be a tuple of "to" lists! -- Then finally we bind the unzip function around that expression - bound_unzipped_inner_list_expr - = case unzip_stuff of - Nothing -> inner_list_expr - Just (unzip_fn, unzip_rhs) -> Let (Rec [(unzip_fn, unzip_rhs)]) $ - mkApps (Var map_id) $ - [ Type (mkListTy from_tup_ty) - , Type to_bndrs_tup_ty - , Var unzip_fn - , inner_list_expr] + bound_unzipped_inner_list_expr' + = case unzip_stuff' of + Nothing -> inner_list_expr' + Just (unzip_fn', unzip_rhs') -> + Let (Rec [(unzip_fn', unzip_rhs')]) $ + mkApps (Var map_id) $ + [ Type (mkListTy from_tup_ty) + , Type to_bndrs_tup_ty + , Var unzip_fn' + , inner_list_expr' ] -- Build a pattern that ensures the consumer binds into the NEW binders, -- which hold lists rather than single values - let pat = mkBigLHsVarPatTupId to_bndrs - return (bound_unzipped_inner_list_expr, pat) + let pat = mkBigLHsVarPatTupId to_bndrs -- NB: no '! + return (bound_unzipped_inner_list_expr', pat) dsTransStmt _ = panic "dsTransStmt: Not given a TransStmt" @@ -260,13 +265,12 @@ deBindComp :: OutPat Id -> CoreExpr -> DsM (Expr Id) deBindComp pat core_list1 quals core_list2 = do - let - u3_ty@u1_ty = exprType core_list1 -- two names, same thing + let u3_ty@u1_ty = exprType core_list1 -- two names, same thing -- u1_ty is a [alpha] type, and u2_ty = alpha - u2_ty = hsLPatType pat + let u2_ty = hsLPatType pat - res_ty = exprType core_list2 + let res_ty = exprType core_list2 h_ty = u1_ty `mkFunTy` res_ty [h, u1, u2, u3] <- newSysLocalsDs [h_ty, u1_ty, u2_ty, u3_ty] @@ -309,9 +313,9 @@ TE[ e | p <- l , q ] c n = let \end{verbatim} -} -dfListComp :: Id -> Id -- 'c' and 'n' - -> [ExprStmt Id] -- the rest of the qual's - -> DsM CoreExpr +dfListComp :: Id -> Id -- 'c' and 'n' + -> [ExprStmt Id] -- the rest of the qual's + -> DsM CoreExpr dfListComp _ _ [] = panic "dfListComp" @@ -355,7 +359,7 @@ dfBindComp :: Id -> Id -- 'c' and 'n' dfBindComp c_id n_id (pat, core_list1) quals = do -- find the required type let x_ty = hsLPatType pat - b_ty = idType n_id + let b_ty = idType n_id -- create some new local id's [b, x] <- newSysLocalsDs [b_ty, x_ty] @@ -570,7 +574,7 @@ dePArrComp (BindStmt p e _ _ : qs) pa cea = do -- dePArrComp (LetStmt (L _ ds) : qs) pa cea = do mapP <- dsDPHBuiltin mapPVar - let xs = collectLocalBinders ds + let xs = collectLocalBinders ds ty'cea = parrElemType cea v <- newSysLocalDs ty'cea clet <- dsLocalBinds ds (mkCoreTup (map Var xs)) @@ -629,10 +633,10 @@ dePArrParComp qss quals = do -- generate Core corresponding to `\p -> e' -- -deLambda :: Type -- type of the argument - -> LPat Id -- argument pattern - -> LHsExpr Id -- body - -> DsM (CoreExpr, Type) +deLambda :: Type -- type of the argument + -> LPat Id -- argument pattern + -> LHsExpr Id -- body + -> DsM (CoreExpr, Type) deLambda ty p e = mkLambda ty p =<< dsLExpr e @@ -720,37 +724,39 @@ dsMcStmt (TransStmt { trS_stmts = stmts, trS_bndrs = bndrs , trS_ret = return_op, trS_bind = bind_op , trS_fmap = fmap_op, trS_form = form }) stmts_rest = do { let (from_bndrs, to_bndrs) = unzip bndrs - from_bndr_tys = map idType from_bndrs -- Types ty + + ; let from_bndr_tys = map idType from_bndrs -- Types ty + -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders - ; expr <- dsInnerMonadComp stmts from_bndrs return_op + ; expr' <- dsInnerMonadComp stmts from_bndrs 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 using - ; usingArgs <- case by of - Nothing -> return [expr] - Just by_e -> do { by_e' <- dsLExpr by_e - ; lam <- matchTuple from_bndrs by_e' - ; return [lam, expr] } + ; usingArgs' <- case by of + Nothing -> return [expr'] + Just by_e -> do { by_e' <- dsLExpr by_e + ; lam' <- matchTuple from_bndrs by_e' + ; return [lam', expr'] } -- 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 -- n (a,b,c) - tup_n_ty = mkBigCoreVarTupTy to_bndrs - - ; body <- dsMcStmts stmts_rest - ; n_tup_var <- newSysLocalDs n_tup_ty - ; tup_n_var <- newSysLocalDs tup_n_ty - ; tup_n_expr <- mkMcUnzipM form fmap_op n_tup_var from_bndr_tys - ; us <- newUniqueSupply - ; let rhs' = mkApps usingExpr' usingArgs - body' = mkTupleCase us to_bndrs body tup_n_var tup_n_expr - - ; return (mkApps bind_op' [rhs', Lam n_tup_var body']) } + ; 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' -- n (a,b,c) + tup_n_ty' = mkBigCoreVarTupTy to_bndrs + + ; body <- dsMcStmts stmts_rest + ; n_tup_var' <- newSysLocalDs n_tup_ty' + ; tup_n_var' <- newSysLocalDs tup_n_ty' + ; tup_n_expr' <- mkMcUnzipM form fmap_op n_tup_var' from_bndr_tys + ; us <- newUniqueSupply + ; let rhs' = mkApps usingExpr' usingArgs' + body' = mkTupleCase us to_bndrs body tup_n_var' tup_n_expr' + + ; return (mkApps bind_op' [rhs', Lam n_tup_var' body']) } -- Parallel statements. Use `Control.Monad.Zip.mzip` to zip parallel -- statements, for example: |