X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise%2FUtils%2FPADict.hs;h=93f2297d2276495f99220fc8b1637f8d762b822a;hb=65ea7d89555c33391dc1729e9742216a0e3d171c;hp=30937b1f72265f4820d1b0b9716721c53b1b6185;hpb=ddfe30b656d1211f77a1eef551ed475b4441eb6a;p=ghc-hetmet.git

diff --git a/compiler/vectorise/Vectorise/Utils/PADict.hs b/compiler/vectorise/Vectorise/Utils/PADict.hs
index 30937b1..93f2297 100644
--- a/compiler/vectorise/Vectorise/Utils/PADict.hs
+++ b/compiler/vectorise/Vectorise/Utils/PADict.hs
@@ -1,10 +1,7 @@
 
 module Vectorise.Utils.PADict (
-	mkPADictType,
 	paDictArgType,
 	paDictOfType,
-	paDFunType,
-	paDFunApply,
 	paMethod	
 )
 where
@@ -13,6 +10,8 @@ import Vectorise.Builtins
 import Vectorise.Utils.Base
 
 import CoreSyn
+import CoreUtils
+import Coercion
 import Type
 import TypeRep
 import TyCon
@@ -22,10 +21,11 @@ import FastString
 import Control.Monad
 
 
-mkPADictType :: Type -> VM Type
-mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
-
-
+-- | Construct the PA argument type for the tyvar. For the tyvar (v :: *) it's
+-- just PA v. For (v :: (* -> *) -> *) it's
+--
+-- > forall (a :: * -> *). (forall (b :: *). PA b -> PA (a b)) -> PA (v a)
+--
 paDictArgType :: TyVar -> VM (Maybe Type)
 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
   where
@@ -42,60 +42,53 @@ paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
 
     go ty k
       | isLiftedTypeKind k
-      = liftM Just (mkPADictType ty)
+      = liftM Just (mkBuiltinTyConApp paTyCon [ty])
 
     go _ _ = return Nothing
 
 
--- | Get the PA dictionary for some type, or `Nothing` if there isn't one.
-paDictOfType :: Type -> VM (Maybe CoreExpr)
+-- | Get the PA dictionary for some type
+--
+paDictOfType :: Type -> VM CoreExpr
 paDictOfType ty 
   = paDictOfTyApp ty_fn ty_args
   where
     (ty_fn, ty_args) = splitAppTys ty
 
-    paDictOfTyApp :: Type -> [Type] -> VM (Maybe CoreExpr)
+    paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
     paDictOfTyApp ty_fn ty_args
         | Just ty_fn' <- coreView ty_fn 
         = paDictOfTyApp ty_fn' ty_args
 
+    -- for type variables, look up the dfun and apply to the PA dictionaries
+    -- of the type arguments
     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
+     = do dfun <- maybeCantVectoriseM "No PA dictionary for type variable"
+                                      (ppr tv <+> text "in" <+> ppr ty)
+                $ lookupTyVarPA tv
+          dicts <- mapM paDictOfType ty_args
+          return $ dfun `mkTyApps` ty_args `mkApps` dicts
+
+    -- for tycons, we also need to apply the dfun to the PR dictionary of
+    -- the representation type if the tycon is polymorphic
+    paDictOfTyApp (TyConApp tc []) ty_args
+     = do
+         dfun <- maybeCantVectoriseM "No PA dictionary for type constructor"
+                                      (ppr tc <+> text "in" <+> ppr ty)
+                $ lookupTyConPA tc
+         super <- super_dict tc ty_args
+         dicts <- mapM paDictOfType ty_args
+         return $ Var dfun `mkTyApps` ty_args `mkApps` super `mkApps` dicts
+
+    paDictOfTyApp _ _ = failure
+
+    super_dict _ [] = return []
+    super_dict tycon ty_args
+      = do
+          pr <- prDictOfPReprInst (TyConApp tycon ty_args)
+          return [pr]
 
+    failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
 
 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
 paMethod _ name ty
@@ -106,7 +99,92 @@ paMethod _ name ty
 
 paMethod method _ ty
   = do
-      fn        <- builtin method
-      Just dict <- paDictOfType ty
+      fn   <- builtin method
+      dict <- paDictOfType ty
       return $ mkApps (Var fn) [Type ty, dict]
 
+-- | Given a type @ty@, return the PR dictionary for @PRepr ty@.
+prDictOfPReprInst :: Type -> VM CoreExpr
+prDictOfPReprInst ty
+  = do
+      (prepr_tc, prepr_args) <- preprSynTyCon ty
+      case coreView (mkTyConApp prepr_tc prepr_args) of
+        Just rhs -> do
+                      dict <- prDictOfReprType rhs
+                      pr_co <- mkBuiltinCo prTyCon
+                      let Just arg_co = tyConFamilyCoercion_maybe prepr_tc
+                      let co = mkAppCoercion pr_co
+                             $ mkSymCoercion
+                             $ mkTyConApp arg_co prepr_args
+                      return $ mkCoerce co dict
+        Nothing  -> cantVectorise "Invalid PRepr type instance"
+                                  $ ppr ty
+
+-- | Get the PR dictionary for a type. The argument must be a representation
+-- type.
+prDictOfReprType :: Type -> VM CoreExpr
+prDictOfReprType ty
+  | Just (tycon, tyargs) <- splitTyConApp_maybe ty
+    = do
+        prepr <- builtin preprTyCon
+        if tycon == prepr
+          then do
+                 [ty'] <- return tyargs
+                 prDictOfPReprInst ty'
+          else do 
+                 -- a representation tycon must have a PR instance
+                 dfun <- maybeCantVectoriseM
+                           "No PR dictionary for type constructor"
+                           (ppr tycon <+> text "in" <+> ppr ty)
+                       $ lookupTyConPR tycon
+                 prDFunApply dfun tyargs
+
+  | otherwise
+    = do
+        -- it is a tyvar or an application of a tyvar
+        -- determine the PR dictionary from its PA dictionary
+        --
+        -- NOTE: This assumes that PRepr t ~ t is for all representation types
+        -- t
+        --
+        -- FIXME: This doesn't work for kinds other than * at the moment. We'd
+        -- have to simply abstract the term over the missing type arguments.
+        pa    <- paDictOfType ty
+        prsel <- builtin paPRSel
+        return $ Var prsel `mkApps` [Type ty, pa]
+
+-- | Apply a tycon's PR dfun to dictionary arguments (PR or PA) corresponding
+-- to the argument types.
+prDFunApply :: Var -> [Type] -> VM CoreExpr
+prDFunApply dfun tys
+  | Just [] <- ctxs    -- PR (a :-> b) doesn't have a context
+  = return $ Var dfun `mkTyApps` tys
+
+  | Just tycons <- ctxs
+  , length tycons == length tys
+  = do
+      pa <- builtin paTyCon
+      pr <- builtin prTyCon 
+      args <- zipWithM (dictionary pa pr) tys tycons
+      return $ Var dfun `mkTyApps` tys `mkApps` args
+
+  | otherwise = invalid
+  where
+    -- the dfun's contexts - if its type is (PA a, PR b) => PR (C a b) then
+    -- ctxs is Just [PA, PR]
+    ctxs = fmap (map fst)
+         $ sequence
+         $ map splitTyConApp_maybe
+         $ fst
+         $ splitFunTys
+         $ snd
+         $ splitForAllTys
+         $ varType dfun
+
+    dictionary pa pr ty tycon
+      | tycon == pa = paDictOfType ty
+      | tycon == pr = prDictOfReprType ty
+      | otherwise   = invalid
+
+    invalid = cantVectorise "Invalid PR dfun type" (ppr (varType dfun) <+> ppr tys)
+