- -- Either rules are switched off, or we are ignoring what there are;
- -- Either way foldr/build won't happen, so use the more efficient
- -- Wadler-style desugaring
- || isParallelComp quals
- -- Foldr-style desugaring can't handle
- -- parallel list comprehensions
- then deListComp quals body (mkNilExpr elt_ty)
-
- else -- Foldr/build should be enabled, so desugar
- -- into foldrs and builds
- newTyVarsDs [alphaTyVar] `thenDs` \ [n_tyvar] ->
- let
- n_ty = mkTyVarTy n_tyvar
- c_ty = mkFunTys [elt_ty, n_ty] n_ty
- in
- newSysLocalsDs [c_ty,n_ty] `thenDs` \ [c, n] ->
- dfListComp c n quals body `thenDs` \ result ->
- dsLookupGlobalId buildName `thenDs` \ build_id ->
- returnDs (Var build_id `App` Type elt_ty
- `App` mkLams [n_tyvar, c, n] result)
-
- where isParallelComp (ParStmt bndrstmtss : _) = True
- isParallelComp _ = False
+ -- Either rules are switched off, or we are ignoring what there are;
+ -- Either way foldr/build won't happen, so use the more efficient
+ -- Wadler-style desugaring
+ || isParallelComp quals
+ -- Foldr-style desugaring can't handle parallel list comprehensions
+ then deListComp quals body (mkNilExpr elt_ty)
+ else mkBuildExpr elt_ty (\(c, _) (n, _) -> dfListComp c n quals body)
+ -- Foldr/build should be enabled, so desugar
+ -- into foldrs and builds
+
+ where
+ -- We must test for ParStmt anywhere, not just at the head, because an extension
+ -- to list comprehensions would be to add brackets to specify the associativity
+ -- of qualifier lists. This is really easy to do by adding extra ParStmts into the
+ -- mix of possibly a single element in length, so we do this to leave the possibility open
+ isParallelComp = any isParallelStmt
+
+ isParallelStmt (ParStmt _) = True
+ isParallelStmt _ = False
+
+
+-- This function lets you desugar a inner list comprehension and a list of the binders
+-- of that comprehension that we need in the outer comprehension into such an expression
+-- and the type of the elements that it outputs (tuples of binders)
+dsInnerListComp :: ([LStmt Id], [Id]) -> DsM (CoreExpr, Type)
+dsInnerListComp (stmts, bndrs) = do
+ expr <- dsListComp stmts (mkBigLHsVarTup bndrs) bndrs_tuple_type
+ return (expr, bndrs_tuple_type)
+ where
+ bndrs_types = map idType bndrs
+ bndrs_tuple_type = mkBigCoreTupTy bndrs_types
+
+
+-- This function factors out commonality between the desugaring strategies for TransformStmt.
+-- Given such a statement it gives you back an expression representing how to compute the transformed
+-- list and the tuple that you need to bind from that list in order to proceed with your desugaring
+dsTransformStmt :: Stmt Id -> DsM (CoreExpr, LPat Id)
+dsTransformStmt (TransformStmt (stmts, binders) usingExpr maybeByExpr) = do
+ (expr, binders_tuple_type) <- dsInnerListComp (stmts, binders)
+ usingExpr' <- dsLExpr usingExpr
+
+ using_args <-
+ case maybeByExpr of
+ Nothing -> return [expr]
+ Just byExpr -> do
+ byExpr' <- dsLExpr byExpr
+
+ us <- newUniqueSupply
+ [tuple_binder] <- newSysLocalsDs [binders_tuple_type]
+ let byExprWrapper = mkTupleCase us binders byExpr' tuple_binder (Var tuple_binder)
+
+ return [Lam tuple_binder byExprWrapper, expr]
+
+ let inner_list_expr = mkApps usingExpr' ((Type binders_tuple_type) : using_args)
+
+ let pat = mkBigLHsVarPatTup binders
+ return (inner_list_expr, pat)
+
+-- 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
+-- list and the tuple that you need to bind from that list in order to proceed with your desugaring
+dsGroupStmt :: Stmt Id -> DsM (CoreExpr, LPat Id)
+dsGroupStmt (GroupStmt (stmts, binderMap) groupByClause) = do
+ let (fromBinders, toBinders) = unzip binderMap
+
+ fromBindersTypes = map idType fromBinders
+ toBindersTypes = map idType toBinders
+
+ toBindersTupleType = mkBigCoreTupTy toBindersTypes
+
+ -- Desugar an inner comprehension which outputs a list of tuples of the "from" binders
+ (expr, fromBindersTupleType) <- dsInnerListComp (stmts, fromBinders)
+
+ -- 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', usingArgs) <-
+ case groupByClause of
+ GroupByNothing usingExpr -> liftM2 (,) (dsLExpr usingExpr) (return [expr])
+ GroupBySomething usingExpr byExpr -> do
+ usingExpr' <- dsLExpr (either id noLoc usingExpr)
+
+ byExpr' <- dsLExpr byExpr
+
+ us <- newUniqueSupply
+ [fromBindersTuple] <- newSysLocalsDs [fromBindersTupleType]
+ let byExprWrapper = mkTupleCase us fromBinders byExpr' fromBindersTuple (Var fromBindersTuple)
+
+ return (usingExpr', [Lam fromBindersTuple byExprWrapper, expr])
+
+ -- Create an unzip function for the appropriate arity and element types and find "map"
+ (unzip_fn, unzip_rhs) <- mkUnzipBind fromBindersTypes
+ 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' ((Type fromBindersTupleType) : 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!
+ unzipped_inner_list_expr = mkApps (Var map_id)
+ [Type (mkListTy fromBindersTupleType), Type toBindersTupleType, Var unzip_fn, inner_list_expr]
+ -- Then finally we bind the unzip function around that expression
+ bound_unzipped_inner_list_expr = Let (Rec [(unzip_fn, unzip_rhs)]) unzipped_inner_list_expr
+
+ -- Build a pattern that ensures the consumer binds into the NEW binders, which hold lists rather than single values
+ let pat = mkBigLHsVarPatTup toBinders
+ return (bound_unzipped_inner_list_expr, pat)
+