Move VectCore to Vectorise tree
[ghc-hetmet.git] / compiler / vectorise / VectUtils.hs
index 8dccd61..37dbecb 100644 (file)
@@ -21,9 +21,9 @@ module VectUtils (
   buildClosure, buildClosures,
   mkClosureApp
 ) where
-
-import VectCore
 import VectMonad
+import Vectorise.Env
+import Vectorise.Vect
 
 import MkCore ( mkCoreTup, mkWildCase )
 import CoreSyn
@@ -38,9 +38,10 @@ import Var
 import MkId               ( unwrapFamInstScrut )
 import Id                 ( setIdUnfolding )
 import TysWiredIn
-import BasicTypes         ( Boxity(..) )
+import BasicTypes         ( Boxity(..), Arity )
 import Literal            ( Literal, mkMachInt )
 
+
 import Outputable
 import FastString
 
@@ -163,6 +164,7 @@ prDFunOfTyCon tycon
   . maybeCantVectoriseM "No PR dictionary for tycon" (ppr tycon)
   $ lookupTyConPR tycon
 
+
 paDictArgType :: TyVar -> VM (Maybe Type)
 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
   where
@@ -183,25 +185,39 @@ paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
 
     go _ _ = return Nothing
 
-paDictOfType :: Type -> VM CoreExpr
-paDictOfType ty = paDictOfTyApp ty_fn ty_args
+
+-- | Get the PA dictionary for some type, or `Nothing` if there isn't one.
+paDictOfType :: Type -> VM (Maybe 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
-      dfun <- maybeCantVectoriseM "No PA dictionary for tycon" (ppr tc)
-            $ lookupTyConPA tc
-      paDFunApply (Var dfun) ty_args
-paDictOfTyApp ty _
-  = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
+    paDictOfTyApp :: Type -> [Type] -> VM (Maybe 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)
+          liftM Just $ paDFunApply dfun ty_args
+
+    paDictOfTyApp (TyConApp tc _) ty_args
+     = do mdfun <- lookupTyConPA tc
+          case mdfun of
+           Nothing     
+            -> pprTrace "VectUtils.paDictOfType"
+                        (vcat [ text "No PA dictionary"
+                              , text "for tycon: " <> ppr tc
+                              , text "in type:   " <> ppr ty])
+            $ return Nothing
+
+           Just dfun   -> liftM Just $ paDFunApply (Var dfun) ty_args
+
+    paDictOfTyApp ty _
+     = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
+
+
 
 paDFunType :: TyCon -> VM Type
 paDFunType tc
@@ -216,10 +232,10 @@ paDFunType tc
 
 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
 paDFunApply dfun tys
-  = do
-      dicts <- mapM paDictOfType tys
+ = do Just dicts <- liftM sequence $ mapM paDictOfType tys
       return $ mkApps (mkTyApps dfun tys) dicts
 
+
 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
 paMethod _ name ty
   | Just tycon <- splitPrimTyCon ty
@@ -229,8 +245,8 @@ paMethod _ name ty
 
 paMethod method _ ty
   = do
-      fn   <- builtin method
-      dict <- paDictOfType ty
+      fn        <- builtin method
+      Just dict <- paDictOfType ty
       return $ mkApps (Var fn) [Type ty, dict]
 
 prDictOfType :: Type -> VM CoreExpr
@@ -256,8 +272,8 @@ prDFunApply dfun tys
 wrapPR :: Type -> VM CoreExpr
 wrapPR ty
   = do
-      pa_dict <- paDictOfType ty
-      pr_dfun <- prDFunOfTyCon =<< builtin wrapTyCon
+      Just  pa_dict <- paDictOfType ty
+      pr_dfun       <- prDFunOfTyCon =<< builtin wrapTyCon
       return $ mkApps pr_dfun [Type ty, pa_dict]
 
 replicatePD :: CoreExpr -> CoreExpr -> VM CoreExpr
@@ -281,6 +297,7 @@ combinePD ty len sel xs
   where
     n = length xs
 
+-- | Like `replicatePD` but use the lifting context in the vectoriser state.
 liftPD :: CoreExpr -> VM CoreExpr
 liftPD x
   = do
@@ -300,8 +317,8 @@ zipScalars arg_tys res_ty
 scalarClosure :: [Type] -> Type -> CoreExpr -> CoreExpr -> VM CoreExpr
 scalarClosure arg_tys res_ty scalar_fun array_fun
   = do
-      ctr <- builtin (closureCtrFun $ length arg_tys)
-      pas <- mapM paDictOfType (init arg_tys)
+      ctr      <- builtin (closureCtrFun $ length arg_tys)
+      Just pas <- liftM sequence $ mapM paDictOfType (init arg_tys)
       return $ Var ctr `mkTyApps` (arg_tys ++ [res_ty])
                        `mkApps`   (pas ++ [scalar_fun, array_fun])
 
@@ -337,27 +354,31 @@ polyArity tvs = do
 
 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
 polyApply expr tys
-  = do
-      dicts <- mapM paDictOfType tys
+ = do Just dicts <- liftM sequence $ mapM paDictOfType tys
       return $ expr `mkTyApps` tys `mkApps` dicts
 
 polyVApply :: VExpr -> [Type] -> VM VExpr
 polyVApply expr tys
-  = do
-      dicts <- mapM paDictOfType tys
-      return $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr
-
+ = do Just dicts <- liftM sequence $ mapM paDictOfType tys
+      return     $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr
 
-data Inline = Inline Int -- arity
-            | DontInline
+-- Inline ---------------------------------------------------------------------
+-- | Records whether we should inline a particular binding.
+data Inline 
+        = Inline Arity
+        | DontInline
 
+-- | Add to the arity contained within an `Inline`, if any.
 addInlineArity :: Inline -> Int -> Inline
 addInlineArity (Inline m) n = Inline (m+n)
 addInlineArity DontInline _ = DontInline
 
+-- | Says to always inline a binding.
 inlineMe :: Inline
 inlineMe = Inline 0
 
+
+-- Hoising --------------------------------------------------------------------
 hoistBinding :: Var -> CoreExpr -> VM ()
 hoistBinding v e = updGEnv $ \env ->
   env { global_bindings = (v,e) : global_bindings env }
@@ -371,7 +392,7 @@ hoistExpr fs expr inl
   where
     mk_inline var = case inl of
                       Inline arity -> var `setIdUnfolding`
-                                      mkInlineRule needSaturated expr arity
+                                      mkInlineRule expr (Just arity)
                       DontInline   -> var
 
 hoistVExpr :: VExpr -> Inline -> VM VVar
@@ -412,24 +433,25 @@ boxExpr ty (vexpr, lexpr)
         Nothing     -> return (vexpr, lexpr)
 -}
 
+-- Closures -------------------------------------------------------------------
 mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr
 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
-  = do
-      dict <- paDictOfType env_ty
-      mkv  <- builtin closureVar
-      mkl  <- builtin liftedClosureVar
+ = do Just dict <- paDictOfType env_ty
+      mkv       <- builtin closureVar
+      mkl       <- builtin liftedClosureVar
       return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
               Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
 
+
 mkClosureApp :: Type -> Type -> VExpr -> VExpr -> VM VExpr
 mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg)
-  = do
-      vapply <- builtin applyVar
+ = do vapply <- builtin applyVar
       lapply <- builtin liftedApplyVar
       lc     <- builtin liftingContext
       return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
               Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [Var lc, lclo, larg])
 
+
 buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
 buildClosures _   _    [] _ mk_body
   = mk_body
@@ -471,6 +493,8 @@ buildClosure tvs vars arg_ty res_ty mk_body
 
       mkClosure arg_ty res_ty env_ty fn env
 
+
+-- Environments ---------------------------------------------------------------
 buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr)
 buildEnv [] = do
              ty    <- voidType