2 collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
4 mkPADictType, mkPArrayType,
5 paDictArgType, paDictOfType,
6 paMethod, lengthPA, replicatePA, emptyPA,
7 polyAbstract, polyApply,
12 #include "HsVersions.h"
27 import Control.Monad ( liftM, zipWithM_ )
29 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
30 collectAnnTypeArgs expr = go expr []
32 go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
35 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
36 collectAnnTypeBinders expr = go [] expr
38 go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
39 go bs e = (reverse bs, e)
41 isAnnTypeArg :: AnnExpr b ann -> Bool
42 isAnnTypeArg (_, AnnType t) = True
43 isAnnTypeArg _ = False
45 isClosureTyCon :: TyCon -> Bool
46 isClosureTyCon tc = tyConName tc == closureTyConName
48 splitClosureTy :: Type -> (Type, Type)
50 | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
54 | otherwise = pprPanic "splitClosureTy" (ppr ty)
56 isPArrayTyCon :: TyCon -> Bool
57 isPArrayTyCon tc = tyConName tc == parrayTyConName
59 splitPArrayTy :: Type -> Type
61 | Just (tc, [arg_ty]) <- splitTyConApp_maybe ty
65 | otherwise = pprPanic "splitPArrayTy" (ppr ty)
67 mkPADictType :: Type -> VM Type
70 tc <- builtin paDictTyCon
71 return $ TyConApp tc [ty]
73 mkPArrayType :: Type -> VM Type
76 tc <- builtin parrayTyCon
77 return $ TyConApp tc [ty]
79 paDictArgType :: TyVar -> VM (Maybe Type)
80 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
82 go ty k | Just k' <- kindView k = go ty k'
85 tv <- newTyVar FSLIT("a") k1
86 mty1 <- go (TyVarTy tv) k1
89 mty2 <- go (AppTy ty (TyVarTy tv)) k2
90 return $ fmap (ForAllTy tv . FunTy ty1) mty2
95 = liftM Just (mkPADictType ty)
97 go ty k = return Nothing
99 paDictOfType :: Type -> VM CoreExpr
100 paDictOfType ty = paDictOfTyApp ty_fn ty_args
102 (ty_fn, ty_args) = splitAppTys ty
104 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
105 paDictOfTyApp ty_fn ty_args
106 | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
107 paDictOfTyApp (TyVarTy tv) ty_args
109 dfun <- maybeV (lookupTyVarPA tv)
110 paDFunApply dfun ty_args
111 paDictOfTyApp (TyConApp tc _) ty_args
113 pa_class <- builtin paClass
114 (dfun, ty_args') <- lookupInst pa_class [TyConApp tc ty_args]
115 paDFunApply (Var dfun) ty_args'
116 paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
118 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
121 dicts <- mapM paDictOfType tys
122 return $ mkApps (mkTyApps dfun tys) dicts
124 paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr
128 dict <- paDictOfType ty
129 return $ mkApps (Var fn) [Type ty, dict]
131 lengthPA :: CoreExpr -> VM CoreExpr
132 lengthPA x = liftM (`App` x) (paMethod lengthPAVar ty)
134 ty = splitPArrayTy (exprType x)
136 replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
137 replicatePA len x = liftM (`mkApps` [len,x])
138 (paMethod replicatePAVar (exprType x))
140 emptyPA :: Type -> VM CoreExpr
141 emptyPA = paMethod emptyPAVar
143 polyAbstract :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
147 mdicts <- mapM mk_dict_var tvs
148 zipWithM_ (\tv -> maybe (defLocalTyVar tv) (defLocalTyVarWithPA tv . Var)) tvs mdicts
152 r <- paDictArgType tv
154 Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty)
155 Nothing -> return Nothing
157 mk_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts])
159 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
162 dicts <- mapM paDictOfType tys
163 return $ expr `mkTyApps` tys `mkApps` dicts
165 lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
166 lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
168 hoistExpr :: FastString -> CoreExpr -> VM Var
171 var <- newLocalVar fs (exprType expr)
173 env { global_bindings = (var, expr) : global_bindings env }
176 takeHoisted :: VM [(Var, CoreExpr)]
180 setGEnv $ env { global_bindings = [] }
181 return $ global_bindings env