2 collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
6 mkPRepr, mkToPRepr, mkToArrPRepr, mkFromPRepr, mkFromArrPRepr,
7 mkPADictType, mkPArrayType, mkPReprType,
8 parrayCoerce, parrayReprTyCon, parrayReprDataCon, mkVScrut,
9 prDictOfType, prCoerce,
10 paDictArgType, paDictOfType, paDFunType,
11 paMethod, lengthPA, replicatePA, emptyPA, liftPA,
12 polyAbstract, polyApply, polyVApply,
13 hoistBinding, hoistExpr, hoistPolyVExpr, takeHoisted,
14 buildClosure, buildClosures,
18 #include "HsVersions.h"
30 import DataCon ( DataCon, dataConWrapId, dataConTag )
32 import Id ( mkWildId )
33 import MkId ( unwrapFamInstScrut )
37 import BasicTypes ( Boxity(..) )
42 import Data.List ( zipWith4 )
43 import Control.Monad ( liftM, liftM2, zipWithM_ )
45 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
46 collectAnnTypeArgs expr = go expr []
48 go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
51 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
52 collectAnnTypeBinders expr = go [] expr
54 go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
55 go bs e = (reverse bs, e)
57 collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
58 collectAnnValBinders expr = go [] expr
60 go bs (_, AnnLam b e) | isId b = go (b:bs) e
61 go bs e = (reverse bs, e)
63 isAnnTypeArg :: AnnExpr b ann -> Bool
64 isAnnTypeArg (_, AnnType t) = True
65 isAnnTypeArg _ = False
67 mkDataConTag :: DataCon -> CoreExpr
68 mkDataConTag dc = mkConApp intDataCon [mkIntLitInt $ dataConTag dc]
70 splitUnTy :: String -> Name -> Type -> Type
72 | Just (tc, [ty']) <- splitTyConApp_maybe ty
73 , tyConName tc == name
76 | otherwise = pprPanic s (ppr ty)
78 splitBinTy :: String -> Name -> Type -> (Type, Type)
80 | Just (tc, [ty1, ty2]) <- splitTyConApp_maybe ty
81 , tyConName tc == name
84 | otherwise = pprPanic s (ppr ty)
86 splitFixedTyConApp :: TyCon -> Type -> [Type]
87 splitFixedTyConApp tc ty
88 | Just (tc', tys) <- splitTyConApp_maybe ty
92 | otherwise = pprPanic "splitFixedTyConApp" (ppr tc <+> ppr ty)
94 splitEmbedTy :: Type -> Type
95 splitEmbedTy = splitUnTy "splitEmbedTy" embedTyConName
97 splitClosureTy :: Type -> (Type, Type)
98 splitClosureTy = splitBinTy "splitClosureTy" closureTyConName
100 splitPArrayTy :: Type -> Type
101 splitPArrayTy = splitUnTy "splitPArrayTy" parrayTyConName
103 mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type
104 mkBuiltinTyConApp get_tc tys
107 return $ mkTyConApp tc tys
109 mkBuiltinTyConApps :: (Builtins -> TyCon) -> [Type] -> Type -> VM Type
110 mkBuiltinTyConApps get_tc tys ty
113 return $ foldr (mk tc) ty tys
115 mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
117 mkBuiltinTyConApps1 :: (Builtins -> TyCon) -> Type -> [Type] -> VM Type
118 mkBuiltinTyConApps1 get_tc dft [] = return dft
119 mkBuiltinTyConApps1 get_tc dft tys
123 [] -> pprPanic "mkBuiltinTyConApps1" (ppr tc)
124 _ -> return $ foldr1 (mk tc) tys
126 mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
128 mkPRepr :: [[Type]] -> VM Type
131 embed_tc <- builtin embedTyCon
132 sum_tcs <- builtins sumTyCon
133 prod_tcs <- builtins prodTyCon
135 let mk_sum [] = unitTy
137 mk_sum tys = mkTyConApp (sum_tcs $ length tys) tys
141 mk_prod tys = mkTyConApp (prod_tcs $ length tys) tys
143 mk_embed ty = mkTyConApp embed_tc [ty]
146 . map (mk_prod . map mk_embed)
149 mkToPRepr :: [[CoreExpr]] -> VM ([CoreExpr], Type)
152 embed_tc <- builtin embedTyCon
153 embed_dc <- builtin embedDataCon
154 sum_tcs <- builtins sumTyCon
155 prod_tcs <- builtins prodTyCon
157 let mk_sum [] = ([Var unitDataConId], unitTy)
158 mk_sum [(expr, ty)] = ([expr], ty)
159 mk_sum es = (zipWith mk_alt (tyConDataCons sum_tc) exprs,
160 mkTyConApp sum_tc tys)
162 (exprs, tys) = unzip es
163 sum_tc = sum_tcs (length es)
164 mk_alt dc expr = mkConApp dc (map Type tys ++ [expr])
166 mk_prod [] = (Var unitDataConId, unitTy)
167 mk_prod [(expr, ty)] = (expr, ty)
168 mk_prod es = (mkConApp prod_dc (map Type tys ++ exprs),
169 mkTyConApp prod_tc tys)
171 (exprs, tys) = unzip es
172 prod_tc = prod_tcs (length es)
173 [prod_dc] = tyConDataCons prod_tc
175 mk_embed expr = (mkConApp embed_dc [Type ty, expr],
176 mkTyConApp embed_tc [ty])
177 where ty = exprType expr
179 return . mk_sum $ map (mk_prod . map mk_embed) ess
181 mkToArrPRepr :: CoreExpr -> CoreExpr -> [[CoreExpr]] -> VM CoreExpr
182 mkToArrPRepr len sel ess
184 embed_tc <- builtin embedTyCon
185 (embed_rtc, _) <- parrayReprTyCon (mkTyConApp embed_tc [unitTy])
186 let [embed_rdc] = tyConDataCons embed_rtc
188 let mk_sum [(expr, ty)] = return (expr, ty)
191 sum_tc <- builtin . sumTyCon $ length es
192 (sum_rtc, _) <- parrayReprTyCon (mkTyConApp sum_tc tys)
193 let [sum_rdc] = tyConDataCons sum_rtc
195 return (mkConApp sum_rdc (map Type tys ++ (len : sel : exprs)),
196 mkTyConApp sum_tc tys)
198 (exprs, tys) = unzip es
200 mk_prod [(expr, ty)] = return (expr, ty)
203 prod_tc <- builtin . prodTyCon $ length es
204 (prod_rtc, _) <- parrayReprTyCon (mkTyConApp prod_tc tys)
205 let [prod_rdc] = tyConDataCons prod_rtc
207 return (mkConApp prod_rdc (map Type tys ++ (len : exprs)),
208 mkTyConApp prod_tc tys)
210 (exprs, tys) = unzip es
212 mk_embed expr = (mkConApp embed_rdc [Type ty, expr],
213 mkTyConApp embed_tc [ty])
214 where ty = splitPArrayTy (exprType expr)
216 liftM fst (mk_sum =<< mapM (mk_prod . map mk_embed) ess)
218 mkFromPRepr :: CoreExpr -> Type -> [([Var], CoreExpr)] -> VM CoreExpr
219 mkFromPRepr scrut res_ty alts
221 embed_dc <- builtin embedDataCon
222 sum_tcs <- builtins sumTyCon
223 prod_tcs <- builtins prodTyCon
225 let un_sum expr ty [(vars, res)] = un_prod expr ty vars res
228 ps <- mapM (newLocalVar FSLIT("p")) tys
230 $ zipWith4 un_prod (map Var ps) tys vars rs
231 return . Case expr (mkWildId ty) res_ty
232 $ zipWith3 mk_alt sum_dcs ps bodies
234 (vars, rs) = unzip bs
235 tys = splitFixedTyConApp sum_tc ty
236 sum_tc = sum_tcs $ length bs
237 sum_dcs = tyConDataCons sum_tc
239 mk_alt dc p body = (DataAlt dc, [p], body)
241 un_prod expr ty [] r = return r
242 un_prod expr ty [var] r = return $ un_embed expr ty var r
243 un_prod expr ty vars r
245 xs <- mapM (newLocalVar FSLIT("x")) tys
246 let body = foldr (\(e,t,v) r -> un_embed e t v r) r
247 $ zip3 (map Var xs) tys vars
248 return $ Case expr (mkWildId ty) res_ty
249 [(DataAlt prod_dc, xs, body)]
251 tys = splitFixedTyConApp prod_tc ty
252 prod_tc = prod_tcs $ length vars
253 [prod_dc] = tyConDataCons prod_tc
255 un_embed expr ty var r
256 = Case expr (mkWildId ty) res_ty
257 [(DataAlt embed_dc, [var], r)]
259 un_sum scrut (exprType scrut) alts
261 mkFromArrPRepr :: CoreExpr -> Type -> Var -> Var -> [[Var]] -> CoreExpr
263 mkFromArrPRepr scrut res_ty len sel vars res
264 = return (Var unitDataConId)
266 mkClosureType :: Type -> Type -> VM Type
267 mkClosureType arg_ty res_ty = mkBuiltinTyConApp closureTyCon [arg_ty, res_ty]
269 mkClosureTypes :: [Type] -> Type -> VM Type
270 mkClosureTypes = mkBuiltinTyConApps closureTyCon
272 mkPReprType :: Type -> VM Type
273 mkPReprType ty = mkBuiltinTyConApp preprTyCon [ty]
275 mkPADictType :: Type -> VM Type
276 mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
278 mkPArrayType :: Type -> VM Type
279 mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty]
281 parrayCoerce :: TyCon -> [Type] -> CoreExpr -> VM CoreExpr
282 parrayCoerce repr_tc args expr
283 | Just arg_co <- tyConFamilyCoercion_maybe repr_tc
285 parray <- builtin parrayTyCon
287 let co = mkAppCoercion (mkTyConApp parray [])
288 (mkSymCoercion (mkTyConApp arg_co args))
290 return $ mkCoerce co expr
292 parrayReprTyCon :: Type -> VM (TyCon, [Type])
293 parrayReprTyCon ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
295 parrayReprDataCon :: Type -> VM (DataCon, [Type])
298 (tc, arg_tys) <- parrayReprTyCon ty
299 let [dc] = tyConDataCons tc
302 mkVScrut :: VExpr -> VM (VExpr, TyCon, [Type])
305 (tc, arg_tys) <- parrayReprTyCon (exprType ve)
306 return ((ve, unwrapFamInstScrut tc arg_tys le), tc, arg_tys)
308 prDictOfType :: Type -> VM CoreExpr
310 | Just (tycon, ty_args) <- splitTyConApp_maybe orig_ty
312 dfun <- traceMaybeV "prDictOfType" (ppr tycon) (lookupTyConPR tycon)
313 prDFunApply (Var dfun) ty_args
315 prDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
318 args <- mapM mkDFunArg arg_tys
319 return $ mkApps mono_dfun args
321 mono_dfun = mkTyApps dfun tys
322 (arg_tys, _) = splitFunTys (exprType mono_dfun)
324 mkDFunArg :: Type -> VM CoreExpr
326 | Just (tycon, [arg]) <- splitTyConApp_maybe ty
328 = let name = tyConName tycon
330 get_dict | name == paTyConName = paDictOfType
331 | name == prTyConName = prDictOfType
332 | otherwise = pprPanic "mkDFunArg" (ppr ty)
336 mkDFunArg ty = pprPanic "mkDFunArg" (ppr ty)
338 prCoerce :: TyCon -> [Type] -> CoreExpr -> VM CoreExpr
339 prCoerce repr_tc args expr
340 | Just arg_co <- tyConFamilyCoercion_maybe repr_tc
342 pr_tc <- builtin prTyCon
344 let co = mkAppCoercion (mkTyConApp pr_tc [])
345 (mkSymCoercion (mkTyConApp arg_co args))
347 return $ mkCoerce co expr
349 paDictArgType :: TyVar -> VM (Maybe Type)
350 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
352 go ty k | Just k' <- kindView k = go ty k'
355 tv <- newTyVar FSLIT("a") k1
356 mty1 <- go (TyVarTy tv) k1
359 mty2 <- go (AppTy ty (TyVarTy tv)) k2
360 return $ fmap (ForAllTy tv . FunTy ty1) mty2
365 = liftM Just (mkPADictType ty)
367 go ty k = return Nothing
369 paDictOfType :: Type -> VM CoreExpr
370 paDictOfType ty = paDictOfTyApp ty_fn ty_args
372 (ty_fn, ty_args) = splitAppTys ty
374 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
375 paDictOfTyApp ty_fn ty_args
376 | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
377 paDictOfTyApp (TyVarTy tv) ty_args
379 dfun <- maybeV (lookupTyVarPA tv)
380 paDFunApply dfun ty_args
381 paDictOfTyApp (TyConApp tc _) ty_args
383 dfun <- traceMaybeV "paDictOfTyApp" (ppr tc) (lookupTyConPA tc)
384 paDFunApply (Var dfun) ty_args
385 paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
387 paDFunType :: TyCon -> VM Type
390 margs <- mapM paDictArgType tvs
391 res <- mkPADictType (mkTyConApp tc arg_tys)
392 return . mkForAllTys tvs
393 $ mkFunTys [arg | Just arg <- margs] res
396 arg_tys = mkTyVarTys tvs
398 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
401 dicts <- mapM paDictOfType tys
402 return $ mkApps (mkTyApps dfun tys) dicts
404 paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr
408 dict <- paDictOfType ty
409 return $ mkApps (Var fn) [Type ty, dict]
411 lengthPA :: CoreExpr -> VM CoreExpr
412 lengthPA x = liftM (`App` x) (paMethod lengthPAVar ty)
414 ty = splitPArrayTy (exprType x)
416 replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
417 replicatePA len x = liftM (`mkApps` [len,x])
418 (paMethod replicatePAVar (exprType x))
420 emptyPA :: Type -> VM CoreExpr
421 emptyPA = paMethod emptyPAVar
423 liftPA :: CoreExpr -> VM CoreExpr
426 lc <- builtin liftingContext
427 replicatePA (Var lc) x
429 newLocalVVar :: FastString -> Type -> VM VVar
432 lty <- mkPArrayType vty
433 vv <- newLocalVar fs vty
434 lv <- newLocalVar fs lty
437 polyAbstract :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
441 mdicts <- mapM mk_dict_var tvs
442 zipWithM_ (\tv -> maybe (defLocalTyVar tv) (defLocalTyVarWithPA tv . Var)) tvs mdicts
446 r <- paDictArgType tv
448 Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty)
449 Nothing -> return Nothing
451 mk_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts])
453 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
456 dicts <- mapM paDictOfType tys
457 return $ expr `mkTyApps` tys `mkApps` dicts
459 polyVApply :: VExpr -> [Type] -> VM VExpr
462 dicts <- mapM paDictOfType tys
463 return $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr
465 hoistBinding :: Var -> CoreExpr -> VM ()
466 hoistBinding v e = updGEnv $ \env ->
467 env { global_bindings = (v,e) : global_bindings env }
469 hoistExpr :: FastString -> CoreExpr -> VM Var
472 var <- newLocalVar fs (exprType expr)
473 hoistBinding var expr
476 hoistVExpr :: VExpr -> VM VVar
480 vv <- hoistExpr ('v' `consFS` fs) ve
481 lv <- hoistExpr ('l' `consFS` fs) le
484 hoistPolyVExpr :: [TyVar] -> VM VExpr -> VM VExpr
487 expr <- closedV . polyAbstract tvs $ \abstract ->
488 liftM (mapVect abstract) p
489 fn <- hoistVExpr expr
490 polyVApply (vVar fn) (mkTyVarTys tvs)
492 takeHoisted :: VM [(Var, CoreExpr)]
496 setGEnv $ env { global_bindings = [] }
497 return $ global_bindings env
499 mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr
500 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
502 dict <- paDictOfType env_ty
503 mkv <- builtin mkClosureVar
504 mkl <- builtin mkClosurePVar
505 return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
506 Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
508 mkClosureApp :: VExpr -> VExpr -> VM VExpr
509 mkClosureApp (vclo, lclo) (varg, larg)
511 vapply <- builtin applyClosureVar
512 lapply <- builtin applyClosurePVar
513 return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
514 Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [lclo, larg])
516 (arg_ty, res_ty) = splitClosureTy (exprType vclo)
518 buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
519 buildClosures tvs vars [] res_ty mk_body
521 buildClosures tvs vars [arg_ty] res_ty mk_body
522 = buildClosure tvs vars arg_ty res_ty mk_body
523 buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body
525 res_ty' <- mkClosureTypes arg_tys res_ty
526 arg <- newLocalVVar FSLIT("x") arg_ty
527 buildClosure tvs vars arg_ty res_ty'
530 lc <- builtin liftingContext
531 clo <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body
532 return $ vLams lc (vars ++ [arg]) clo
534 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
536 -- f = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
537 -- f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
539 buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
540 buildClosure tvs vars arg_ty res_ty mk_body
542 (env_ty, env, bind) <- buildEnv vars
543 env_bndr <- newLocalVVar FSLIT("env") env_ty
544 arg_bndr <- newLocalVVar FSLIT("arg") arg_ty
546 fn <- hoistPolyVExpr tvs
548 lc <- builtin liftingContext
550 body' <- bind (vVar env_bndr)
551 (vVarApps lc body (vars ++ [arg_bndr]))
552 return (vLamsWithoutLC [env_bndr, arg_bndr] body')
554 mkClosure arg_ty res_ty env_ty fn env
556 buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VM VExpr)
559 lc <- builtin liftingContext
560 let (ty, venv, vbind) = mkVectEnv tys vs
561 (lenv, lbind) <- mkLiftEnv lc tys ls
562 return (ty, (venv, lenv),
563 \(venv,lenv) (vbody,lbody) ->
565 let vbody' = vbind venv vbody
566 lbody' <- lbind lenv lbody
567 return (vbody', lbody'))
572 mkVectEnv :: [Type] -> [Var] -> (Type, CoreExpr, CoreExpr -> CoreExpr -> CoreExpr)
573 mkVectEnv [] [] = (unitTy, Var unitDataConId, \env body -> body)
574 mkVectEnv [ty] [v] = (ty, Var v, \env body -> Let (NonRec v env) body)
575 mkVectEnv tys vs = (ty, mkCoreTup (map Var vs),
576 \env body -> Case env (mkWildId ty) (exprType body)
577 [(DataAlt (tupleCon Boxed (length vs)), vs, body)])
581 mkLiftEnv :: Var -> [Type] -> [Var] -> VM (CoreExpr, CoreExpr -> CoreExpr -> VM CoreExpr)
582 mkLiftEnv lc [ty] [v]
583 = return (Var v, \env body ->
585 len <- lengthPA (Var v)
586 return . Let (NonRec v env)
587 $ Case len lc (exprType body) [(DEFAULT, [], body)])
589 -- NOTE: this transparently deals with empty environments
592 (env_tc, env_tyargs) <- parrayReprTyCon vty
593 let [env_con] = tyConDataCons env_tc
595 env = Var (dataConWrapId env_con)
596 `mkTyApps` env_tyargs
597 `mkVarApps` (lc : vs)
599 bind env body = let scrut = unwrapFamInstScrut env_tc env_tyargs env
601 return $ Case scrut (mkWildId (exprType scrut))
603 [(DataAlt env_con, lc : bndrs, body)]
606 vty = mkCoreTupTy tys
608 bndrs | null vs = [mkWildId unitTy]