Break out type vectorisation into own module

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

diff --git a/compiler/vectorise/Vectorise/Type/Type.hs b/compiler/vectorise/Vectorise/Type/Type.hs
new file mode 100644 (file)
index 0000000..00df5d5
--- /dev/null
+++ b/compiler/vectorise/Vectorise/Type/Type.hs
@@ -0,0 +1,117 @@
+
+module Vectorise.Type.Type
+       ( vectTyCon
+       , vectAndLiftType
+       , vectType)
+where
+import VectUtils
+import Vectorise.Monad
+import Vectorise.Builtins
+import TypeRep
+import Type
+import TyCon
+import Var
+import Outputable
+import Control.Monad
+import Data.List
+import Data.Maybe
+
+
+-- | Vectorise a type constructor.
+vectTyCon :: TyCon -> VM TyCon
+vectTyCon tc
+  | isFunTyCon tc        = builtin closureTyCon
+  | isBoxedTupleTyCon tc = return tc
+  | isUnLiftedTyCon tc   = return tc
+  | otherwise            
+  = maybeCantVectoriseM "Tycon not vectorised: " (ppr tc)
+       $ lookupTyCon tc
+
+
+-- | Produce the vectorised and lifted versions of a type.
+vectAndLiftType :: Type -> VM (Type, Type)
+vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
+vectAndLiftType ty
+  = do
+      mdicts   <- mapM paDictArgType tyvars
+      let dicts = [dict | Just dict <- mdicts]
+      vmono_ty <- vectType mono_ty
+      lmono_ty <- mkPDataType vmono_ty
+      return (abstractType tyvars dicts vmono_ty,
+              abstractType tyvars dicts lmono_ty)
+  where
+    (tyvars, mono_ty) = splitForAllTys ty
+
+
+-- | Vectorise a type.
+vectType :: Type -> VM Type
+vectType ty
+       | Just ty'      <- coreView ty
+       = vectType ty'
+       
+vectType (TyVarTy tv)          = return $ TyVarTy tv
+vectType (AppTy ty1 ty2)       = liftM2 AppTy    (vectType ty1) (vectType ty2)
+vectType (TyConApp tc tys)     = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
+vectType (FunTy ty1 ty2)       = liftM2 TyConApp (builtin closureTyCon)
+                                                 (mapM vectAndBoxType [ty1,ty2])
+
+-- For each quantified var we need to add a PA dictionary out the front of the type.
+-- So          forall a. C  a => a -> a   
+-- turns into  forall a. Cv a => PA a => a :-> a
+vectType ty@(ForAllTy _ _)
+ = do
+      -- split the type into the quantified vars, its dictionaries and the body.
+      let (tyvars, tyBody)   = splitForAllTys ty
+      let (tyArgs, tyResult) = splitFunTys    tyBody
+
+      let (tyArgs_dict, tyArgs_regular) 
+                  = partition isDictType tyArgs
+
+      -- vectorise the body.
+      let tyBody' = mkFunTys tyArgs_regular tyResult
+      tyBody''    <- vectType tyBody'
+
+      -- vectorise the dictionary parameters.
+      dictsVect   <- mapM vectType tyArgs_dict
+
+      -- make a PA dictionary for each of the type variables.
+      dictsPA     <- liftM catMaybes $ mapM paDictArgType tyvars
+
+      -- pack it all back together.
+      return $ abstractType tyvars (dictsVect ++ dictsPA) tyBody''
+
+vectType ty = cantVectorise "Can't vectorise type" (ppr ty)
+
+
+-- | Add quantified vars and dictionary parameters to the front of a type.
+abstractType :: [TyVar] -> [Type] -> Type -> Type
+abstractType tyvars dicts = mkForAllTys tyvars . mkFunTys dicts
+
+
+-- | Check if some type is a type class dictionary.
+isDictType :: Type -> Bool
+isDictType ty
+ = case splitTyConApp_maybe ty of
+       Just (tyCon, _)         -> isClassTyCon tyCon
+       _                       -> False
+
+
+-- | Create the boxed version of a vectorised type.
+vectAndBoxType :: Type -> VM Type
+vectAndBoxType ty = vectType ty >>= boxType
+
+
+-- | Create the boxed version of a type.
+boxType :: Type -> VM Type
+boxType ty
+  | Just (tycon, []) <- splitTyConApp_maybe ty
+  , isUnLiftedTyCon tycon
+  = do
+      r <- lookupBoxedTyCon tycon
+      case r of
+        Just tycon' -> return $ mkTyConApp tycon' []
+        Nothing     -> return ty
+
+  | otherwise  = return ty
+
+