Move some vectorisation utility functions

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

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 29774d1..6f9db0a 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -4,6 +4,7 @@ where
 #include "HsVersions.h"
 
 import VectMonad
+import VectUtils
 
 import DynFlags
 import HscTypes
@@ -53,7 +54,7 @@ vectBndr :: Var -> VM (Var, Var)
 vectBndr v
   = do
       vty <- vectType (idType v)
-      lty <- mkPArrayTy vty
+      lty <- mkPArrayType vty
       let vv = v `Id.setIdType` vty
           lv = v `Id.setIdType` lty
       updLEnv (mapTo vv lv)
@@ -83,9 +84,9 @@ vectBndrsIn vs p
 replicateP :: CoreExpr -> CoreExpr -> VM CoreExpr
 replicateP expr len
   = do
-      pa  <- paOfType ty
-      rep <- builtin replicatePAVar
-      return $ mkApps (Var rep) [Type ty, pa, expr, len]
+      dict <- paDictOfType ty
+      rep  <- builtin replicatePAVar
+      return $ mkApps (Var rep) [Type ty, dict, expr, len]
   where
     ty = exprType expr
 
@@ -152,72 +153,16 @@ vectExpr lc (_, AnnLet (AnnRec prs) body)
 vectExpr lc (_, AnnLam bndr body)
   | isTyVar bndr
   = do
-      pa_ty          <- paArgType' (TyVarTy bndr) (tyVarKind bndr)
-      pa_var         <- newLocalVar FSLIT("dPA") pa_ty
-      (vbody, lbody) <- localV
-                      $ do
-                          extendTyVarPA bndr (Var pa_var)
-                          -- FIXME: what about shadowing here (bndr in lc)?
-                          vectExpr lc body
-      return (mkLams [bndr, pa_var] vbody,
-              mkLams [bndr, pa_var] lbody)
-
--- ----------------------------------------------------------------------------
--- PA dictionaries
-
-paArgType :: Type -> Kind -> VM (Maybe Type)
-paArgType ty k
-  | Just k' <- kindView k = paArgType ty k'
-
--- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
--- be made up of * and (->), i.e., they can't be coercion kinds or #.
-paArgType ty (FunTy k1 k2)
-  = do
-      tv  <- newTyVar FSLIT("a") k1
-      ty1 <- paArgType' (TyVarTy tv) k1
-      ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
-      return . Just $ ForAllTy tv (FunTy ty1 ty2)
-
-paArgType ty k
-  | isLiftedTypeKind k
-  = do
-      tc <- builtin paDictTyCon
-      return . Just $ TyConApp tc [ty]
-
-  | otherwise
-  = return Nothing 
-
-paArgType' :: Type -> Kind -> VM Type
-paArgType' ty k
-  = do
-      r <- paArgType ty k
-      case r of
-        Just ty' -> return ty'
-        Nothing  -> pprPanic "paArgType'" (ppr ty)
-
-paOfTyCon :: TyCon -> VM CoreExpr
--- FIXME: just for now
-paOfTyCon tc = maybeV (readGEnv $ \env -> lookupNameEnv (global_tycon_pa env) (tyConName tc))
-
-paOfType :: Type -> VM CoreExpr
-paOfType ty | Just ty' <- coreView ty = paOfType ty'
-
-paOfType (TyVarTy tv) = maybeV (readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv)
-paOfType (AppTy ty1 ty2)
-  = do
-      e1 <- paOfType ty1
-      e2 <- paOfType ty2
-      return $ mkApps e1 [Type ty2, e2]
-paOfType (TyConApp tc tys)
-  = do
-      e  <- paOfTyCon tc
-      es <- mapM paOfType tys
-      return $ mkApps e [arg | (t,e) <- zip tys es, arg <- [Type t, e]]
-paOfType (FunTy ty1 ty2) = paOfType (TyConApp funTyCon [ty1,ty2])
-paOfType t@(ForAllTy tv ty) = pprPanic "paOfType:" (ppr t)
-paOfType ty = pprPanic "paOfType:" (ppr ty)
-        
-
+      r <- paDictArgType bndr
+      (upd_env, add_lam) <- get_upd r
+      (vbody, lbody) <- localV (upd_env >> vectExpr lc body)
+      return (Lam bndr (add_lam vbody), Lam bndr (add_lam lbody))
+  where
+    get_upd Nothing = return (deleteTyVarPA bndr, id)
+    get_upd (Just pa_ty) = do
+                             pa_var <- newLocalVar FSLIT("dPA") pa_ty
+                             return (extendTyVarPA bndr (Var pa_var),
+                                     Lam pa_var)
 
 -- ----------------------------------------------------------------------------
 -- Types
@@ -244,25 +189,12 @@ vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
                                              (mapM vectType [ty1,ty2])
 vectType (ForAllTy tv ty)
   = do
-      r   <- paArgType (TyVarTy tv) (tyVarKind tv)
+      r   <- paDictArgType tv
       ty' <- vectType ty
-      return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
+      return $ ForAllTy tv (wrap r ty')
+  where
+    wrap Nothing      = id
+    wrap (Just pa_ty) = FunTy pa_ty
 
 vectType ty = pprPanic "vectType:" (ppr ty)
 
-isClosureTyCon :: TyCon -> Bool
-isClosureTyCon tc = tyConUnique tc == closureTyConKey
-
-splitClosureTy :: Type -> (Type, Type)
-splitClosureTy ty
-  | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
-  , isClosureTyCon tc
-  = (arg_ty, res_ty)
-
-  | otherwise = pprPanic "splitClosureTy" (ppr ty)
-
-mkPArrayTy :: Type -> VM Type
-mkPArrayTy ty = do
-                  tc <- builtin parrayTyCon
-                  return $ TyConApp tc [ty]
-