mkPADictType, mkPArrayType,
paDictArgType, paDictOfType,
paMethod, lengthPA, replicatePA, emptyPA,
+ polyAbstract, polyApply,
lookupPArrayFamInst,
hoistExpr, takeHoisted
) where
import Outputable
import FastString
-import Control.Monad ( liftM )
+import Control.Monad ( liftM, zipWithM_ )
collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
collectAnnTypeArgs expr = go expr []
isAnnTypeArg _ = False
isClosureTyCon :: TyCon -> Bool
-isClosureTyCon tc = tyConUnique tc == closureTyConKey
+isClosureTyCon tc = tyConName tc == closureTyConName
splitClosureTy :: Type -> (Type, Type)
splitClosureTy ty
| otherwise = pprPanic "splitClosureTy" (ppr ty)
+isPArrayTyCon :: TyCon -> Bool
+isPArrayTyCon tc = tyConName tc == parrayTyConName
+
+splitPArrayTy :: Type -> Type
+splitPArrayTy ty
+ | Just (tc, [arg_ty]) <- splitTyConApp_maybe ty
+ , isPArrayTyCon tc
+ = arg_ty
+
+ | otherwise = pprPanic "splitPArrayTy" (ppr ty)
+
mkPADictType :: Type -> VM Type
mkPADictType ty
= do
return $ mkApps (Var fn) [Type ty, dict]
lengthPA :: CoreExpr -> VM CoreExpr
-lengthPA x = liftM (`App` x) (paMethod lengthPAVar (exprType x))
+lengthPA x = liftM (`App` x) (paMethod lengthPAVar ty)
+ where
+ ty = splitPArrayTy (exprType x)
replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
replicatePA len x = liftM (`mkApps` [len,x])
emptyPA :: Type -> VM CoreExpr
emptyPA = paMethod emptyPAVar
+polyAbstract :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
+polyAbstract tvs p
+ = localV
+ $ do
+ mdicts <- mapM mk_dict_var tvs
+ zipWithM_ (\tv -> maybe (defLocalTyVar tv) (defLocalTyVarWithPA tv . Var)) tvs mdicts
+ p (mk_lams mdicts)
+ where
+ mk_dict_var tv = do
+ r <- paDictArgType tv
+ case r of
+ Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty)
+ Nothing -> return Nothing
+
+ mk_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts])
+
+polyApply :: CoreExpr -> [Type] -> VM CoreExpr
+polyApply expr tys
+ = do
+ dicts <- mapM paDictOfType tys
+ return $ expr `mkTyApps` tys `mkApps` dicts
+
lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])