2 collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
4 mkPADictType, mkPArrayType,
5 paDictArgType, paDictOfType,
6 paMethod, lengthPA, replicatePA, emptyPA,
7 polyAbstract, polyApply,
12 #include "HsVersions.h"
22 import DataCon ( dataConWrapId )
24 import Id ( mkWildId )
25 import MkId ( unwrapFamInstScrut )
28 import BasicTypes ( Boxity(..) )
33 import Control.Monad ( liftM, zipWithM_ )
35 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
36 collectAnnTypeArgs expr = go expr []
38 go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
41 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
42 collectAnnTypeBinders expr = go [] expr
44 go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
45 go bs e = (reverse bs, e)
47 isAnnTypeArg :: AnnExpr b ann -> Bool
48 isAnnTypeArg (_, AnnType t) = True
49 isAnnTypeArg _ = False
51 isClosureTyCon :: TyCon -> Bool
52 isClosureTyCon tc = tyConName tc == closureTyConName
54 splitClosureTy :: Type -> (Type, Type)
56 | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
60 | otherwise = pprPanic "splitClosureTy" (ppr ty)
62 isPArrayTyCon :: TyCon -> Bool
63 isPArrayTyCon tc = tyConName tc == parrayTyConName
65 splitPArrayTy :: Type -> Type
67 | Just (tc, [arg_ty]) <- splitTyConApp_maybe ty
71 | otherwise = pprPanic "splitPArrayTy" (ppr ty)
73 mkPADictType :: Type -> VM Type
76 tc <- builtin paDictTyCon
77 return $ TyConApp tc [ty]
79 mkPArrayType :: Type -> VM Type
82 tc <- builtin parrayTyCon
83 return $ TyConApp tc [ty]
85 paDictArgType :: TyVar -> VM (Maybe Type)
86 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
88 go ty k | Just k' <- kindView k = go ty k'
91 tv <- newTyVar FSLIT("a") k1
92 mty1 <- go (TyVarTy tv) k1
95 mty2 <- go (AppTy ty (TyVarTy tv)) k2
96 return $ fmap (ForAllTy tv . FunTy ty1) mty2
101 = liftM Just (mkPADictType ty)
103 go ty k = return Nothing
105 paDictOfType :: Type -> VM CoreExpr
106 paDictOfType ty = paDictOfTyApp ty_fn ty_args
108 (ty_fn, ty_args) = splitAppTys ty
110 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
111 paDictOfTyApp ty_fn ty_args
112 | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
113 paDictOfTyApp (TyVarTy tv) ty_args
115 dfun <- maybeV (lookupTyVarPA tv)
116 paDFunApply dfun ty_args
117 paDictOfTyApp (TyConApp tc _) ty_args
119 pa_class <- builtin paClass
120 (dfun, ty_args') <- lookupInst pa_class [TyConApp tc ty_args]
121 paDFunApply (Var dfun) ty_args'
122 paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
124 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
127 dicts <- mapM paDictOfType tys
128 return $ mkApps (mkTyApps dfun tys) dicts
130 paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr
134 dict <- paDictOfType ty
135 return $ mkApps (Var fn) [Type ty, dict]
137 lengthPA :: CoreExpr -> VM CoreExpr
138 lengthPA x = liftM (`App` x) (paMethod lengthPAVar ty)
140 ty = splitPArrayTy (exprType x)
142 replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
143 replicatePA len x = liftM (`mkApps` [len,x])
144 (paMethod replicatePAVar (exprType x))
146 emptyPA :: Type -> VM CoreExpr
147 emptyPA = paMethod emptyPAVar
149 polyAbstract :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
153 mdicts <- mapM mk_dict_var tvs
154 zipWithM_ (\tv -> maybe (defLocalTyVar tv) (defLocalTyVarWithPA tv . Var)) tvs mdicts
158 r <- paDictArgType tv
160 Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty)
161 Nothing -> return Nothing
163 mk_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts])
165 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
168 dicts <- mapM paDictOfType tys
169 return $ expr `mkTyApps` tys `mkApps` dicts
171 lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
172 lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
174 hoistExpr :: FastString -> CoreExpr -> VM Var
177 var <- newLocalVar fs (exprType expr)
179 env { global_bindings = (var, expr) : global_bindings env }
182 hoistPolyExpr :: FastString -> [TyVar] -> CoreExpr -> VM CoreExpr
183 hoistPolyExpr fs tvs expr
185 poly_expr <- closedV . polyAbstract tvs $ \abstract -> return (abstract expr)
186 fn <- hoistExpr fs poly_expr
187 polyApply (Var fn) (mkTyVarTys tvs)
189 takeHoisted :: VM [(Var, CoreExpr)]
193 setGEnv $ env { global_bindings = [] }
194 return $ global_bindings env
197 mkClosure :: Type -> Type -> Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr
198 mkClosure arg_ty res_ty env_ty pa_dict vfn lfn env
200 mk <- builtin mkClosureVar
201 return $ Var mk `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [pa_dict, vfn, lfn, env]
203 mkClosureP :: Type -> Type -> Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr
204 mkClosureP arg_ty res_ty env_ty pa_dict vfn lfn env
206 mk <- builtin mkClosurePVar
207 return $ Var mk `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [pa_dict, vfn, lfn, env]
209 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
211 -- f = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
212 -- f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
214 buildClosure :: [TyVar] -> [(Var,Var)] -> (Var,Var) -> Var -> CoreExpr -> CoreExpr -> VM (CoreExpr, CoreExpr)
215 buildClosure tvs env (varg, larg) lv vbody lbody
217 let (venv_ty, venv, bind_venv) = mkVectEnv tys vs
218 (lenv, bind_lenv) <- mkLiftEnv (Var lv) tys ls
219 lenv_ty <- mkPArrayType venv_ty
221 venv_bndr <- newLocalVar FSLIT("env") venv_ty
222 lenv_bndr <- newLocalVar FSLIT("env") lenv_ty
224 let mono_vfn = mkLams [venv_bndr, varg]
225 . bind_venv (Var venv_bndr)
226 $ vbody `mkVarApps` vs `mkVarApps` [varg]
227 mono_lfn = mkLams [lenv_bndr, larg]
228 . bind_lenv (Var lenv_bndr) lv
229 $ lbody `mkVarApps` (lv:ls) `mkVarApps` [larg]
231 vfn <- hoistPolyExpr FSLIT("vfn") tvs mono_vfn
232 lfn <- hoistPolyExpr FSLIT("lfn") tvs mono_lfn
234 pa_dict <- paDictOfType venv_ty
236 vclo <- mkClosure arg_ty res_ty venv_ty pa_dict vfn lfn venv
237 lclo <- mkClosureP arg_ty res_ty venv_ty pa_dict vfn lfn lenv
247 res_ty = exprType vbody
249 mkVectEnv :: [Type] -> [Var] -> (Type, CoreExpr, CoreExpr -> CoreExpr -> CoreExpr)
250 mkVectEnv [] [] = (unitTy, Var unitDataConId, \env body -> body)
251 mkVectEnv [ty] [v] = (ty, Var v, \env body -> Let (NonRec v env) body)
252 mkVectEnv tys vs = (ty, mkCoreTup (map Var vs),
253 \env body -> Case env (mkWildId ty) (exprType body)
254 [(DataAlt (tupleCon Boxed (length vs)), vs, body)])
258 mkLiftEnv :: CoreExpr -> [Type] -> [Var]
259 -> VM (CoreExpr, CoreExpr -> Var -> CoreExpr -> CoreExpr)
260 mkLiftEnv lc [ty] [v]
262 len <- lengthPA (Var v)
263 return (Var v, \env lv body -> Let (NonRec v env)
264 $ Case len lv (exprType body) [(DEFAULT, [], body)])
266 -- NOTE: this transparently deals with empty environments
269 (env_tc, env_tyargs) <- lookupPArrayFamInst vty
270 let [env_con] = tyConDataCons env_tc
272 env = Var (dataConWrapId env_con)
273 `mkTyApps` env_tyargs
274 `mkApps` (lc : map Var vs)
276 bind env lv body = let scrut = unwrapFamInstScrut env_tc env_tyargs env
278 Case scrut (mkWildId (exprType scrut)) (exprType body)
279 [(DataAlt env_con, lv : vs, body)]
282 vty = mkCoreTupTy tys