Use the right dictionary when calling lengthPA

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

diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 76d625c..71ba7a3 100644 (file)
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -1,14 +1,80 @@
 module VectUtils (
-  paDictArgType
+  collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
+  splitClosureTy,
+  mkPADictType, mkPArrayType,
+  paDictArgType, paDictOfType,
+  paMethod, lengthPA, replicatePA, emptyPA,
+  abstractOverTyVars, applyToTypes,
+  lookupPArrayFamInst,
+  hoistExpr, takeHoisted
 ) where
 
 #include "HsVersions.h"
 
 import VectMonad
 
+import CoreSyn
+import CoreUtils
 import Type
 import TypeRep
+import TyCon
 import Var
+import PrelNames
+
+import Outputable
+import FastString
+
+import Control.Monad         ( liftM, zipWithM_ )
+
+collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
+collectAnnTypeArgs expr = go expr []
+  where
+    go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
+    go e                             tys = (e, tys)
+
+collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
+collectAnnTypeBinders expr = go [] expr
+  where
+    go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
+    go bs e                           = (reverse bs, e)
+
+isAnnTypeArg :: AnnExpr b ann -> Bool
+isAnnTypeArg (_, AnnType t) = True
+isAnnTypeArg _              = False
+
+isClosureTyCon :: TyCon -> Bool
+isClosureTyCon tc = tyConName tc == closureTyConName
+
+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)
+
+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
+      tc <- builtin paDictTyCon
+      return $ TyConApp tc [ty]
+
+mkPArrayType :: Type -> VM Type
+mkPArrayType ty
+  = do
+      tc <- builtin parrayTyCon
+      return $ TyConApp tc [ty]
 
 paDictArgType :: TyVar -> VM (Maybe Type)
 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
@@ -26,10 +92,90 @@ paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
 
     go ty k
       | isLiftedTypeKind k
-      = do
-          tc <- builtin paDictTyCon
-          return . Just $ TyConApp tc [ty]
-
+      = liftM Just (mkPADictType ty)
 
     go ty k = return Nothing
 
+paDictOfType :: Type -> VM CoreExpr
+paDictOfType ty = paDictOfTyApp ty_fn ty_args
+  where
+    (ty_fn, ty_args) = splitAppTys ty
+
+paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
+paDictOfTyApp ty_fn ty_args
+  | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
+paDictOfTyApp (TyVarTy tv) ty_args
+  = do
+      dfun <- maybeV (lookupTyVarPA tv)
+      paDFunApply dfun ty_args
+paDictOfTyApp (TyConApp tc _) ty_args
+  = do
+      pa_class <- builtin paClass
+      (dfun, ty_args') <- lookupInst pa_class [TyConApp tc ty_args]
+      paDFunApply (Var dfun) ty_args'
+paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
+
+paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
+paDFunApply dfun tys
+  = do
+      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
+
+abstractOverTyVars :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
+abstractOverTyVars tvs p
+  = 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])
+
+applyToTypes :: CoreExpr -> [Type] -> VM CoreExpr
+applyToTypes expr tys
+  = do
+      dicts <- mapM paDictOfType tys
+      return $ expr `mkTyApps` tys `mkApps` dicts
+
+lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
+lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
+
+hoistExpr :: FastString -> CoreExpr -> VM Var
+hoistExpr fs expr
+  = do
+      var <- newLocalVar fs (exprType expr)
+      updGEnv $ \env ->
+        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
+