Add silent superclass parameters to the vectoriser

1 files changed, 27 insertions, 12 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