Finish breaking up vectoriser utils

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

diff --git a/compiler/vectorise/Vectorise/Utils/PADict.hs b/compiler/vectorise/Vectorise/Utils/PADict.hs
new file mode 100644 (file)
index 0000000..44faa2e
--- /dev/null
+++ b/compiler/vectorise/Vectorise/Utils/PADict.hs
@@ -0,0 +1,113 @@
+
+module Vectorise.Utils.PADict (
+       mkPADictType,
+       paDictArgType,
+       paDictOfType,
+       paDFunType,
+       paDFunApply,
+       paMethod        
+)
+where
+import Vectorise.Monad
+import Vectorise.Builtins
+import Vectorise.Utils.Base
+
+import CoreSyn
+import Coercion
+import Type
+import TypeRep
+import TyCon
+import Var
+import Outputable
+import FastString
+import Control.Monad
+
+
+mkPADictType :: Type -> VM Type
+mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
+
+
+paDictArgType :: TyVar -> VM (Maybe Type)
+paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
+  where
+    go ty k | Just k' <- kindView k = go ty k'
+    go ty (FunTy k1 k2)
+      = do
+          tv   <- newTyVar (fsLit "a") k1
+          mty1 <- go (TyVarTy tv) k1
+          case mty1 of
+            Just ty1 -> do
+                          mty2 <- go (AppTy ty (TyVarTy tv)) k2
+                          return $ fmap (ForAllTy tv . FunTy ty1) mty2
+            Nothing  -> go ty k2
+
+    go ty k
+      | isLiftedTypeKind k
+      = liftM Just (mkPADictType ty)
+
+    go _ _ = return Nothing
+
+
+-- | 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 (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
+  = do
+      margs <- mapM paDictArgType tvs
+      res   <- mkPADictType (mkTyConApp tc arg_tys)
+      return . mkForAllTys tvs
+             $ mkFunTys [arg | Just arg <- margs] res
+  where
+    tvs = tyConTyVars tc
+    arg_tys = mkTyVarTys tvs
+
+paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
+paDFunApply dfun 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
+  = liftM Var
+  . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
+  $ lookupPrimMethod tycon name
+
+paMethod method _ ty
+  = do
+      fn        <- builtin method
+      Just dict <- paDictOfType ty
+      return $ mkApps (Var fn) [Type ty, dict]
+