vectoriser: don't always pass superclass dictionaries to PA dfuns

[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
index d9a00b9..93f2297 100644 (file)
--- a/compiler/vectorise/Vectorise/Utils/PADict.hs
+++ b/compiler/vectorise/Vectorise/Utils/PADict.hs
@@ -1,9 +1,7 @@
 
 module Vectorise.Utils.PADict (
-       mkPADictType,
        paDictArgType,
        paDictOfType,
-       paDFunType,
        paMethod        
 )
 where
@@ -23,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
@@ -43,12 +42,13 @@ 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
+--
 paDictOfType :: Type -> VM CoreExpr
 paDictOfType ty 
   = paDictOfTyApp ty_fn ty_args
@@ -63,35 +63,32 @@ paDictOfType ty
     -- 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)
+     = 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
+    -- the representation type if the tycon is polymorphic
     paDictOfTyApp (TyConApp tc []) ty_args
      = do
-         dfun <- maybeV $ lookupTyConPA tc
-         pr <- prDictOfPRepr tc ty_args
+         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` (pr:dicts)
+         return $ Var dfun `mkTyApps` ty_args `mkApps` super `mkApps` dicts
 
     paDictOfTyApp _ _ = failure
 
-    failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
-
-
+    super_dict _ [] = return []
+    super_dict tycon ty_args
+      = do
+          pr <- prDictOfPReprInst (TyConApp tycon ty_args)
+          return [pr]
 
-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
+    failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
 
 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
 paMethod _ name ty
@@ -106,12 +103,11 @@ paMethod method _ ty
       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
+-- | Given a type @ty@, return the PR dictionary for @PRepr ty@.
+prDictOfPReprInst :: Type -> VM CoreExpr
+prDictOfPReprInst ty
   = do
-      (prepr_tc, prepr_args) <- preprSynTyCon (mkTyConApp tycon tys)
+      (prepr_tc, prepr_args) <- preprSynTyCon ty
       case coreView (mkTyConApp prepr_tc prepr_args) of
         Just rhs -> do
                       dict <- prDictOfReprType rhs
@@ -122,7 +118,7 @@ prDictOfPRepr tycon tys
                              $ mkTyConApp arg_co prepr_args
                       return $ mkCoerce co dict
         Nothing  -> cantVectorise "Invalid PRepr type instance"
-                                  $ ppr $ mkTyConApp prepr_tc prepr_args
+                                  $ ppr ty
 
 -- | Get the PR dictionary for a type. The argument must be a representation
 -- type.
@@ -130,9 +126,18 @@ 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
+        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