X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise%2FUtils%2FPADict.hs;h=d9a00b9a8d8927371df21f72a4a6e2f03159e698;hb=4837a66d2d6a8a6db09067a3cd8c9038a9027a2d;hp=30937b1f72265f4820d1b0b9716721c53b1b6185;hpb=acb9c929a4ab025972027b55b4c18d4410207d29;p=ghc-hetmet.git

diff --git a/compiler/vectorise/Vectorise/Utils/PADict.hs b/compiler/vectorise/Vectorise/Utils/PADict.hs
index 30937b1..d9a00b9 100644
--- a/compiler/vectorise/Vectorise/Utils/PADict.hs
+++ b/compiler/vectorise/Vectorise/Utils/PADict.hs
@@ -4,7 +4,6 @@ module Vectorise.Utils.PADict (
 	paDictArgType,
 	paDictOfType,
 	paDFunType,
-	paDFunApply,
 	paMethod	
 )
 where
@@ -13,6 +12,8 @@ import Vectorise.Builtins
 import Vectorise.Utils.Base
 
 import CoreSyn
+import CoreUtils
+import Coercion
 import Type
 import TypeRep
 import TyCon
@@ -47,36 +48,37 @@ paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
     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
+          dicts <- mapM paDictOfType ty_args
+          return $ dfun `mkTyApps` ty_args `mkApps` dicts
 
-    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
+    -- for tycons, we also need to apply the dfun to the PR dictionary of
+    -- the representation type
+    paDictOfTyApp (TyConApp tc []) ty_args
+     = do
+         dfun <- maybeV $ lookupTyConPA tc
+         pr <- prDictOfPRepr tc ty_args
+         dicts <- mapM paDictOfType ty_args
+         return $ Var dfun `mkTyApps` ty_args `mkApps` (pr:dicts)
 
-	    Just dfun	-> liftM Just $ paDFunApply (Var dfun) ty_args
+    paDictOfTyApp _ _ = failure
 
-    paDictOfTyApp ty _
-     = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
+    failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
 
 
 
@@ -91,12 +93,6 @@ paDFunType tc
     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
@@ -106,7 +102,84 @@ 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]
 
+-- | Get the PR (PRepr t) dictionary, where t is the tycon applied to the type
+-- arguments
+prDictOfPRepr :: TyCon -> [Type] -> VM CoreExpr
+prDictOfPRepr tycon tys
+  = do
+      (prepr_tc, prepr_args) <- preprSynTyCon (mkTyConApp tycon tys)
+      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 $ mkTyConApp prepr_tc prepr_args
+
+-- | 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
+        -- a representation tycon must have a PR instance
+        dfun <- maybeV $ 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)
+