2 collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
5 mkPADictType, mkPArrayType,
6 paDictArgType, paDictOfType,
7 paMethod, lengthPA, replicatePA, emptyPA,
8 polyAbstract, polyApply, polyVApply,
10 hoistExpr, hoistPolyVExpr, takeHoisted,
11 buildClosure, buildClosures
14 #include "HsVersions.h"
25 import DataCon ( dataConWrapId )
27 import Id ( mkWildId )
28 import MkId ( unwrapFamInstScrut )
31 import BasicTypes ( Boxity(..) )
36 import Control.Monad ( liftM, zipWithM_ )
38 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
39 collectAnnTypeArgs expr = go expr []
41 go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
44 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
45 collectAnnTypeBinders expr = go [] expr
47 go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
48 go bs e = (reverse bs, e)
50 collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
51 collectAnnValBinders expr = go [] expr
53 go bs (_, AnnLam b e) | isId b = go (b:bs) e
54 go bs e = (reverse bs, e)
56 isAnnTypeArg :: AnnExpr b ann -> Bool
57 isAnnTypeArg (_, AnnType t) = True
58 isAnnTypeArg _ = False
60 isClosureTyCon :: TyCon -> Bool
61 isClosureTyCon tc = tyConName tc == closureTyConName
63 splitClosureTy :: Type -> (Type, Type)
65 | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
69 | otherwise = pprPanic "splitClosureTy" (ppr ty)
71 isPArrayTyCon :: TyCon -> Bool
72 isPArrayTyCon tc = tyConName tc == parrayTyConName
74 splitPArrayTy :: Type -> Type
76 | Just (tc, [arg_ty]) <- splitTyConApp_maybe ty
80 | otherwise = pprPanic "splitPArrayTy" (ppr ty)
82 mkClosureType :: Type -> Type -> VM Type
83 mkClosureType arg_ty res_ty
85 tc <- builtin closureTyCon
86 return $ mkTyConApp tc [arg_ty, res_ty]
88 mkClosureTypes :: [Type] -> Type -> VM Type
89 mkClosureTypes arg_tys res_ty
91 tc <- builtin closureTyCon
92 return $ foldr (mk tc) res_ty arg_tys
94 mk tc arg_ty res_ty = mkTyConApp tc [arg_ty, res_ty]
96 mkPADictType :: Type -> VM Type
99 tc <- builtin paDictTyCon
100 return $ TyConApp tc [ty]
102 mkPArrayType :: Type -> VM Type
105 tc <- builtin parrayTyCon
106 return $ TyConApp tc [ty]
108 paDictArgType :: TyVar -> VM (Maybe Type)
109 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
111 go ty k | Just k' <- kindView k = go ty k'
114 tv <- newTyVar FSLIT("a") k1
115 mty1 <- go (TyVarTy tv) k1
118 mty2 <- go (AppTy ty (TyVarTy tv)) k2
119 return $ fmap (ForAllTy tv . FunTy ty1) mty2
124 = liftM Just (mkPADictType ty)
126 go ty k = return Nothing
128 paDictOfType :: Type -> VM CoreExpr
129 paDictOfType ty = paDictOfTyApp ty_fn ty_args
131 (ty_fn, ty_args) = splitAppTys ty
133 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
134 paDictOfTyApp ty_fn ty_args
135 | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
136 paDictOfTyApp (TyVarTy tv) ty_args
138 dfun <- maybeV (lookupTyVarPA tv)
139 paDFunApply dfun ty_args
140 paDictOfTyApp (TyConApp tc _) ty_args
142 pa_class <- builtin paClass
143 (dfun, ty_args') <- lookupInst pa_class [TyConApp tc ty_args]
144 paDFunApply (Var dfun) ty_args'
145 paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
147 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
150 dicts <- mapM paDictOfType tys
151 return $ mkApps (mkTyApps dfun tys) dicts
153 paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr
157 dict <- paDictOfType ty
158 return $ mkApps (Var fn) [Type ty, dict]
160 lengthPA :: CoreExpr -> VM CoreExpr
161 lengthPA x = liftM (`App` x) (paMethod lengthPAVar ty)
163 ty = splitPArrayTy (exprType x)
165 replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
166 replicatePA len x = liftM (`mkApps` [len,x])
167 (paMethod replicatePAVar (exprType x))
169 emptyPA :: Type -> VM CoreExpr
170 emptyPA = paMethod emptyPAVar
172 newLocalVVar :: FastString -> Type -> VM VVar
175 lty <- mkPArrayType vty
176 vv <- newLocalVar fs vty
177 lv <- newLocalVar fs lty
180 polyAbstract :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
184 mdicts <- mapM mk_dict_var tvs
185 zipWithM_ (\tv -> maybe (defLocalTyVar tv) (defLocalTyVarWithPA tv . Var)) tvs mdicts
189 r <- paDictArgType tv
191 Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty)
192 Nothing -> return Nothing
194 mk_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts])
196 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
199 dicts <- mapM paDictOfType tys
200 return $ expr `mkTyApps` tys `mkApps` dicts
202 polyVApply :: VExpr -> [Type] -> VM VExpr
205 dicts <- mapM paDictOfType tys
206 return $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr
208 lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
209 lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
211 hoistExpr :: FastString -> CoreExpr -> VM Var
214 var <- newLocalVar fs (exprType expr)
216 env { global_bindings = (var, expr) : global_bindings env }
219 hoistVExpr :: FastString -> VExpr -> VM VVar
220 hoistVExpr fs (ve, le)
222 vv <- hoistExpr ('v' `consFS` fs) ve
223 lv <- hoistExpr ('l' `consFS` fs) le
226 hoistPolyVExpr :: FastString -> [TyVar] -> VM VExpr -> VM VExpr
227 hoistPolyVExpr fs tvs p
229 expr <- closedV . polyAbstract tvs $ \abstract ->
230 liftM (mapVect abstract) p
231 fn <- hoistVExpr fs expr
232 polyVApply (vVar fn) (mkTyVarTys tvs)
234 takeHoisted :: VM [(Var, CoreExpr)]
238 setGEnv $ env { global_bindings = [] }
239 return $ global_bindings env
242 mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr
243 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
245 dict <- paDictOfType env_ty
246 mkv <- builtin mkClosureVar
247 mkl <- builtin mkClosurePVar
248 return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
249 Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
251 buildClosures :: [TyVar] -> Var -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
252 buildClosures tvs lc vars [arg_ty] res_ty mk_body
253 = buildClosure tvs lc vars arg_ty res_ty mk_body
254 buildClosures tvs lc vars (arg_ty : arg_tys) res_ty mk_body
256 res_ty' <- mkClosureTypes arg_tys res_ty
257 arg <- newLocalVVar FSLIT("x") arg_ty
258 buildClosure tvs lc vars arg_ty res_ty'
259 . hoistPolyVExpr FSLIT("fn") tvs
261 clo <- buildClosures tvs lc (vars ++ [arg]) arg_tys res_ty mk_body
262 return $ vLams lc (vars ++ [arg]) clo
264 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
266 -- f = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
267 -- f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
269 buildClosure :: [TyVar] -> Var -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
270 buildClosure tvs lv vars arg_ty res_ty mk_body
272 (env_ty, env, bind) <- buildEnv lv vars
273 env_bndr <- newLocalVVar FSLIT("env") env_ty
274 arg_bndr <- newLocalVVar FSLIT("arg") arg_ty
276 fn <- hoistPolyVExpr FSLIT("fn") tvs
279 body' <- bind (vVar env_bndr)
280 (vVarApps lv body (vars ++ [arg_bndr]))
281 return (vLamsWithoutLC [env_bndr, arg_bndr] body')
283 mkClosure arg_ty res_ty env_ty fn env
285 buildEnv :: Var -> [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VM VExpr)
288 let (ty, venv, vbind) = mkVectEnv tys vs
289 (lenv, lbind) <- mkLiftEnv lv tys ls
290 return (ty, (venv, lenv),
291 \(venv,lenv) (vbody,lbody) ->
293 let vbody' = vbind venv vbody
294 lbody' <- lbind lenv lbody
295 return (vbody', lbody'))
300 mkVectEnv :: [Type] -> [Var] -> (Type, CoreExpr, CoreExpr -> CoreExpr -> CoreExpr)
301 mkVectEnv [] [] = (unitTy, Var unitDataConId, \env body -> body)
302 mkVectEnv [ty] [v] = (ty, Var v, \env body -> Let (NonRec v env) body)
303 mkVectEnv tys vs = (ty, mkCoreTup (map Var vs),
304 \env body -> Case env (mkWildId ty) (exprType body)
305 [(DataAlt (tupleCon Boxed (length vs)), vs, body)])
309 mkLiftEnv :: Var -> [Type] -> [Var] -> VM (CoreExpr, CoreExpr -> CoreExpr -> VM CoreExpr)
310 mkLiftEnv lv [ty] [v]
311 = return (Var v, \env body ->
313 len <- lengthPA (Var v)
314 return . Let (NonRec v env)
315 $ Case len lv (exprType body) [(DEFAULT, [], body)])
317 -- NOTE: this transparently deals with empty environments
320 (env_tc, env_tyargs) <- lookupPArrayFamInst vty
321 let [env_con] = tyConDataCons env_tc
323 env = Var (dataConWrapId env_con)
324 `mkTyApps` env_tyargs
325 `mkVarApps` (lv : vs)
327 bind env body = let scrut = unwrapFamInstScrut env_tc env_tyargs env
329 return $ Case scrut (mkWildId (exprType scrut))
331 [(DataAlt env_con, lv : bndrs, body)]
334 vty = mkCoreTupTy tys
336 bndrs | null vs = [mkWildId unitTy]