Add silent superclass parameters to the vectoriser

4 files changed, 85 insertions, 35 deletions

diff --git a/compiler/vectorise/Vectorise/Generic/PADict.hs b/compiler/vectorise/Vectorise/Generic/PADict.hs
index 5bc25194fc..6b7145d3b7 100644
--- a/compiler/vectorise/Vectorise/Generic/PADict.hs
+++ b/compiler/vectorise/Vectorise/Generic/PADict.hs
@@ -19,6 +19,7 @@ import Type
 import Id
 import Var
 import Name
+import FastString
 
 
 -- |Build the PA dictionary function for some type and hoist it to top level.
@@ -33,8 +34,8 @@ import Name
 --      fromArrPRepr :: PData (PRepr a) -> PData a
 --
 -- Example:
---    df :: forall a. PA a -> PA (T a)
---    df = /\a. \(d:PA a). MkPA ($PR_df a d) ($toPRepr a d) ... 
+--    df :: forall a. PR (PRepr a) -> PA a -> PA (T a)
+--    df = /\a. \(c:PR (PRepr a)) (d:PA a). MkPA c ($PR_df a d) ($toPRepr a d) ... 
 --    $dPR_df :: forall a. PA a -> PR (PRepr (T a))
 --    $dPR_df = ....   
 --    $toRepr :: forall a. PA a -> T a -> PRepr (T a)
@@ -52,34 +53,48 @@ buildPADict
         -> VM Var       -- ^ name of the top-level dictionary function.
 
 buildPADict vect_tc prepr_ax pdata_tc pdatas_tc repr
- = polyAbstract tvs $ \args ->    -- The args are the dictionaries we lambda
-                                  -- abstract over; and they are put in the
-                                  -- envt, so when we need a (PA a) we can 
-                                  -- find it in the envt
+ = polyAbstract tvs $ \args ->    -- The args are the dictionaries we lambda abstract over; and they
+                                  -- are put in the envt, so when we need a (PA a) we can find it in
+                                  -- the envt; they don't include the silent superclass args yet
    do { mod <- liftDs getModuleDs
       ; let dfun_name = mkLocalisedOccName mod mkPADFunOcc vect_tc_name
-      
+   
+          -- The superclass dictionary is a (silent) argument if the tycon is polymorphic...
+      ; let mk_super_ty = do { r <- mkPReprType inst_ty
+                             ; pr_cls <- builtin prClass
+                             ; return $ mkClassPred pr_cls [r]
+                             }
+      ; super_tys <- sequence [mk_super_ty | not (null tvs)]
+      ; super_args <- mapM (newLocalVar (fsLit "pr")) super_tys
+      ; let all_args = super_args ++ args
+     
+          -- ...it is constant otherwise
+      ; super_consts <- sequence [prDictOfPReprInstTyCon inst_ty prepr_ax [] | null tvs]
+
           -- Get ids for each of the methods in the dictionary, including superclass
       ; paMethodBuilders <- buildPAScAndMethods
-      ; method_ids       <- mapM (method args dfun_name) paMethodBuilders
+      ; method_ids       <- mapM (method all_args dfun_name) paMethodBuilders
 
           -- Expression to build the dictionary.
       ; pa_dc  <- builtin paDataCon
-      ; let dict = mkLams (tvs ++ args)
+      ; let dict = mkLams (tvs ++ all_args)
                  $ mkConApp pa_dc
                  $ Type inst_ty
-                   : map (method_call args) method_ids
+                 : map Var super_args ++ super_consts -- the superclass dictionary is either lambda-bound or constant
+                   ++ map (method_call all_args) method_ids
 
           -- Build the type of the dictionary function.
       ; pa_cls <- builtin paClass
       ; let dfun_ty = mkForAllTys tvs
-                    $ mkFunTys (map varType args)
+                    $ mkFunTys (map varType all_args)
                                (mkClassPred pa_cls [inst_ty])
 
           -- Set the unfolding for the inliner.
       ; raw_dfun <- newExportedVar dfun_name dfun_ty
       ; let dfun_unf = mkDFunUnfolding dfun_ty $
-                       map (DFunPolyArg . Var) method_ids
+                       map (const $ DFunLamArg 0) super_args
+                       -- ++ map DFunConstArg super_consts
+                       ++ map (DFunPolyArg . Var) method_ids
             dfun = raw_dfun `setIdUnfolding`  dfun_unf
                             `setInlinePragma` dfunInlinePragma
 
diff --git a/compiler/vectorise/Vectorise/Monad/InstEnv.hs b/compiler/vectorise/Vectorise/Monad/InstEnv.hs
index 34d3d75b75..fc12ee567c 100644
--- a/compiler/vectorise/Vectorise/Monad/InstEnv.hs
+++ b/compiler/vectorise/Vectorise/Monad/InstEnv.hs
@@ -40,13 +40,13 @@ lookupInst cls tys
                   cantVectorise dflags "Vectorise.Monad.InstEnv.lookupInst:" err
        }
 
--- Look up the representation tycon of a family instance.
+-- Look up a family instance.
 --
 -- The match must be unique - ie, match exactly one instance - but the 
 -- type arguments used for matching may be more specific than those of 
 -- the family instance declaration.
 --
--- Return the instance tycon and its type instance.  For example, if we have
+-- Return the family instance and its type instance.  For example, if we have
 --
 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 --
@@ -56,13 +56,12 @@ lookupInst cls tys
 --
 -- which implies that :R42T was declared as 'data instance T [a]'.
 --
-lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
+lookupFamInst :: TyCon -> [Type] -> VM (FamInst, [Type])
 lookupFamInst tycon tys
   = ASSERT( isFamilyTyCon tycon )
     do { instEnv <- readGEnv global_fam_inst_env
        ; case lookupFamInstEnv instEnv tycon tys of
-           [(fam_inst, rep_tys)] -> return ( dataFamInstRepTyCon fam_inst
-                                           , rep_tys)
+           [(fam_inst, rep_tys)] -> return ( fam_inst, rep_tys)
            _other                -> 
              do dflags <- getDynFlags
                 cantVectorise dflags "VectMonad.lookupFamInst: not found: "
diff --git a/compiler/vectorise/Vectorise/Utils/Base.hs b/compiler/vectorise/Vectorise/Utils/Base.hs
index 2b47ddfb9b..9ed4e2c60e 100644
--- a/compiler/vectorise/Vectorise/Utils/Base.hs
+++ b/compiler/vectorise/Vectorise/Utils/Base.hs
@@ -21,6 +21,8 @@ module Vectorise.Utils.Base
   , pdataReprTyConExact
   , pdatasReprTyConExact
   , pdataUnwrapScrut
+  
+  , preprSynTyCon
 ) where
 
 import Vectorise.Monad
@@ -29,6 +31,7 @@ import Vectorise.Builtins
 
 import CoreSyn
 import CoreUtils
+import FamInstEnv
 import Coercion
 import Type
 import TyCon
@@ -200,7 +203,11 @@ unwrapNewTypeBodyOfPDatasWrap e ty
 -- a set of distinct type variables.
 -- 
 pdataReprTyCon :: Type -> VM (TyCon, [Type])
-pdataReprTyCon ty = builtin pdataTyCon >>= (`lookupFamInst` [ty])
+pdataReprTyCon ty 
+  = do 
+    { (famInst, tys) <- builtin pdataTyCon >>= (`lookupFamInst` [ty])
+    ; return (dataFamInstRepTyCon famInst, tys)
+    }
 
 -- |Get the representation tycon of the 'PData' data family for a given type constructor.
 --
@@ -225,7 +232,7 @@ pdatasReprTyConExact tycon
   = do {   -- look up the representation tycon; if there is a match at all, it will be be exact
        ;   -- (i.e.,' _tys' will be distinct type variables)
        ; (ptycon, _tys) <- pdatasReprTyCon (tycon `mkTyConApp` mkTyVarTys (tyConTyVars tycon))
-       ; return ptycon
+       ; return $ dataFamInstRepTyCon ptycon
        }
   where
     pdatasReprTyCon ty = builtin pdatasTyCon >>= (`lookupFamInst` [ty])
@@ -240,3 +247,11 @@ pdataUnwrapScrut (ve, le)
        }
   where
     ty = exprType ve
+
+
+-- 'PRepr' representation types ----------------------------------------------
+
+-- |Get the representation tycon of the 'PRepr' type family for a given type.
+--
+preprSynTyCon :: Type -> VM (FamInst, [Type])
+preprSynTyCon ty = builtin preprTyCon >>= (`lookupFamInst` [ty])
diff --git a/compiler/vectorise/Vectorise/Utils/PADict.hs b/compiler/vectorise/Vectorise/Utils/PADict.hs
index de80127c44..85060c477c 100644
--- a/compiler/vectorise/Vectorise/Utils/PADict.hs
+++ b/compiler/vectorise/Vectorise/Utils/PADict.hs
@@ -12,6 +12,7 @@ import Vectorise.Utils.Base
 
 import CoreSyn
 import CoreUtils
+import FamInstEnv
 import Coercion
 import Type
 import TypeRep
@@ -66,25 +67,35 @@ 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 <- maybeCantVectoriseM "No PA dictionary for type variable"
+      = 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
+        ; 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 noPADictErr (ppr tc <+> text "in" <+> ppr ty)
+      = do 
+        { dfun <- maybeCantVectoriseM noPADictErr (ppr tc <+> text "in" <+> ppr ty)
                 $ lookupTyConPA tc
-         dicts <- mapM paDictOfType ty_args
-         return $ Var dfun `mkTyApps` ty_args `mkApps` dicts
-     where
-       noPADictErr = "No PA dictionary for type constructor (did you import 'Data.Array.Parallel'?)"
-
-    paDictOfTyApp _ _ = do dflags <- getDynFlags
-                           failure dflags
+        ; super <- super_dict tc ty_args
+        ; dicts <- mapM paDictOfType ty_args
+        ; return $ Var dfun `mkTyApps` ty_args `mkApps` super `mkApps` dicts
+        }
+      where
+        noPADictErr = "No PA dictionary for type constructor (did you import 'Data.Array.Parallel'?)"
+
+        super_dict _     []      = return []
+        super_dict tycon ty_args
+          = do
+            { pr <- prDictOfPReprInst (TyConApp tycon ty_args)
+            ; return [pr]
+            }
+    
+    paDictOfTyApp _ _ = getDynFlags >>= failure
 
     failure dflags = cantVectorise dflags "Can't construct PA dictionary for type" (ppr ty)
 
@@ -96,11 +107,21 @@ paMethod _ query ty
   = liftM Var $ builtin (query tycon)
 paMethod method _ ty
   = do
-      fn   <- builtin method
-      dict <- paDictOfType ty
-      return $ mkApps (Var fn) [Type ty, dict]
+    { 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_fam, prepr_args) <- preprSynTyCon ty
+    ; prDictOfPReprInstTyCon ty (famInstAxiom prepr_fam) prepr_args
+    }
 
--- | Given a type @ty@, its PRepr synonym tycon and its type arguments,
+-- |Given a type @ty@, its PRepr synonym tycon and its type arguments,
 -- return the PR @PRepr ty@. Suppose we have:
 --
 -- > type instance PRepr (T a1 ... an) = t