Refactoring

[ghc-hetmet.git] / compiler / vectorise / VectUtils.hs

diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 74a3405..7b0e4af 100644 (file)
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -3,8 +3,10 @@ module VectUtils (
   splitClosureTy,
   mkPADictType, mkPArrayType,
   paDictArgType, paDictOfType,
+  paMethod, lengthPA, replicatePA, emptyPA,
+  polyAbstract, polyApply,
   lookupPArrayFamInst,
-  hoistExpr
+  hoistExpr, takeHoisted
 ) where
 
 #include "HsVersions.h"
@@ -22,7 +24,7 @@ import PrelNames
 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 []
@@ -41,7 +43,7 @@ isAnnTypeArg (_, AnnType t) = True
 isAnnTypeArg _              = False
 
 isClosureTyCon :: TyCon -> Bool
-isClosureTyCon tc = tyConUnique tc == closureTyConKey
+isClosureTyCon tc = tyConName tc == closureTyConName
 
 splitClosureTy :: Type -> (Type, Type)
 splitClosureTy ty
@@ -51,6 +53,17 @@ 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
@@ -108,6 +121,47 @@ paDFunApply dfun tys
       dicts <- mapM paDictOfType tys
       return $ mkApps (mkTyApps dfun tys) dicts
 
+paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr
+paMethod method ty
+  = do
+      fn   <- builtin method
+      dict <- paDictOfType ty
+      return $ mkApps (Var fn) [Type ty, dict]
+
+lengthPA :: CoreExpr -> VM CoreExpr
+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])
+                          (paMethod replicatePAVar (exprType 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])
 
@@ -119,3 +173,10 @@ hoistExpr fs expr
         env { global_bindings = (var, expr) : global_bindings env }
       return var
 
+takeHoisted :: VM [(Var, CoreExpr)]
+takeHoisted
+  = do
+      env <- readGEnv id
+      setGEnv $ env { global_bindings = [] }
+      return $ global_bindings env
+