Fix the vectorisation of workers of data constructors

6 files changed, 103 insertions, 53 deletions

diff --git a/compiler/vectorise/Vectorise/Env.hs b/compiler/vectorise/Vectorise/Env.hs
index 0020d67412..64ab075cef 100644
--- a/compiler/vectorise/Vectorise/Env.hs
+++ b/compiler/vectorise/Vectorise/Env.hs
@@ -211,9 +211,10 @@ modVectInfo env mg_ids mg_tyCons vectDecls info
     vectTypeTyCons  = [tycon | VectType _ tycon _ <- vectDecls] ++
                       [tycon | VectClass tycon    <- vectDecls]
     vectDataCons    = concatMap tyConDataCons vectTypeTyCons
-    ids             = mg_ids ++ vectIds ++ selIds
+    ids             = mg_ids ++ vectIds ++ dataConIds ++ selIds
     tyCons          = mg_tyCons ++ vectTypeTyCons
     dataCons        = concatMap tyConDataCons mg_tyCons ++ vectDataCons
+    dataConIds      = map dataConWorkId dataCons
     selIds          = concat [ classAllSelIds cls 
                              | tycon <- tyCons
                              , cls <- maybeToList . tyConClass_maybe $ tycon]
diff --git a/compiler/vectorise/Vectorise/Exp.hs b/compiler/vectorise/Vectorise/Exp.hs
index ca7b13f866..8afe149496 100644
--- a/compiler/vectorise/Vectorise/Exp.hs
+++ b/compiler/vectorise/Vectorise/Exp.hs
@@ -522,7 +522,8 @@ unVectDict ty e
 -- |Vectorise an 'n'-ary lambda abstraction by building a set of 'n' explicit closures.
 --
 -- All non-dictionary free variables go into the closure's environment, whereas the dictionary
--- variables are passed explicit (as conventional arguments) into the body during closure construction.
+-- variables are passed explicit (as conventional arguments) into the body during closure
+-- construction.
 --
 vectLam :: Bool             -- ^ When the RHS of a binding, whether that binding should be inlined.
         -> Bool             -- ^ Whether the binding is a loop breaker.
diff --git a/compiler/vectorise/Vectorise/Generic/PAMethods.hs b/compiler/vectorise/Vectorise/Generic/PAMethods.hs
index 4712e21ee1..ad74ab34f8 100644
--- a/compiler/vectorise/Vectorise/Generic/PAMethods.hs
+++ b/compiler/vectorise/Vectorise/Generic/PAMethods.hs
@@ -287,7 +287,7 @@ buildToArrPRepr vect_tc prepr_tc pdata_tc _ r
     to_comp expr (Wrap ty)
       = do
           wrap_tc  <- builtin wrapTyCon
-          (pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
+          pwrap_tc <- pdataReprTyConExact (mkTyConApp wrap_tc [ty])
           return $ wrapNewTypeBody pwrap_tc [ty] expr
 
 
@@ -355,8 +355,8 @@ buildFromArrPRepr vect_tc prepr_tc pdata_tc _ r
 
     from_comp _ res expr (Keep _ _) = return (res, [expr])
     from_comp _ res expr (Wrap ty)
-     = do wrap_tc       <- builtin wrapTyCon
-          (pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
+     = do wrap_tc  <- builtin wrapTyCon
+          pwrap_tc <- pdataReprTyConExact (mkTyConApp wrap_tc [ty])
           return (res, [unwrapNewTypeBody pwrap_tc [ty]
                         $ unwrapFamInstScrut pwrap_tc [ty] expr])
 
diff --git a/compiler/vectorise/Vectorise/Type/Env.hs b/compiler/vectorise/Vectorise/Type/Env.hs
index 87d071770c..1b806c3138 100644
--- a/compiler/vectorise/Vectorise/Type/Env.hs
+++ b/compiler/vectorise/Vectorise/Type/Env.hs
@@ -164,7 +164,6 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
            -- Furthermore, 'drop_tcs' are those type constructors that we cannot vectorise.
        ; let maybeVectoriseTyCons           = filter notLocalScalarTyCon tycons ++ impVectTyCons
              (conv_tcs, keep_tcs, drop_tcs) = classifyTyCons vectTyConFlavour maybeVectoriseTyCons
-             orig_tcs                       = keep_tcs ++ conv_tcs
              
        ; traceVt " VECT SCALAR    : " $ ppr localScalarTyCons
        ; traceVt " VECT [class]   : " $ ppr impVectTyCons
@@ -206,8 +205,8 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
            -- We don't need new representation types for dictionary constructors. The constructors
            -- are always fully applied, and we don't need to lift them to arrays as a dictionary
            -- of a particular type always has the same value.
-       ; let vect_tcs = filter (not . isClassTyCon) 
-                      $ keep_tcs ++ new_tcs
+       ; let orig_tcs = filter (not . isClassTyCon) $ keep_tcs ++ conv_tcs
+             vect_tcs = filter (not . isClassTyCon) $ keep_tcs ++ new_tcs
 
            -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
            -- type constructors with vectorised representations.
@@ -220,18 +219,36 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
              fam_insts = map mkLocalFamInst inst_tcs
        ; updGEnv $ extendFamEnv fam_insts
 
-           -- Generate dfuns for the 'PA' instances of the vectorised type constructors and
-           -- associate the type constructors with their dfuns in the global environment.  We get
-           -- back the dfun bindings (which we will subsequently inject into the modules toplevel).
+           -- Generate workers for the vectorised data constructors, dfuns for the 'PA' instances of
+           -- the vectorised type constructors, and associate the type constructors with their dfuns
+           -- in the global environment.  We get back the dfun bindings (which we will subsequently
+           -- inject into the modules toplevel).
        ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
            do { defTyConPAs (zipLazy vect_tcs dfuns)
-              ; dfuns <- sequence 
-                      $  zipWith5 buildTyConBindings
-                                  orig_tcs
-                                  vect_tcs
-                                  repr_tcs
-                                  pdata_tcs
-                                  pdatas_tcs
+
+                  -- query the 'PData' instance type constructors for type constructors that have a
+                  -- VECTORISE pragma with an explicit right-hand side (this is Item (3) of
+                  -- "Note [Pragmas to vectorise tycons]" above)
+              ; pdata_withRHS_tcs <- mapM pdataReprTyConExact
+                                          [ mkTyConApp tycon tys
+                                          | (tycon, _) <- vectTyConsWithRHS
+                                          , let tys = mkTyVarTys (tyConTyVars tycon)
+                                          ]
+
+                  -- build workers for all vectorised data constructors (except scalar ones)
+              ; sequence_ $
+                  zipWith3 vectDataConWorkers (orig_tcs  ++ map fst vectTyConsWithRHS)
+                                              (vect_tcs  ++ map snd vectTyConsWithRHS)
+                                              (pdata_tcs ++ pdata_withRHS_tcs)
+
+                  -- build a 'PA' dictionary for all type constructors (except scalar ones and those
+                  -- defined with an explicit right-hand side where the dictionary is user-supplied)
+              ; dfuns <- sequence $
+                           zipWith4 buildTyConPADict
+                                    vect_tcs
+                                    repr_tcs
+                                    pdata_tcs
+                                    pdatas_tcs
 
               ; binds <- takeHoisted
               ; return (dfuns, binds)
@@ -244,23 +261,32 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
 
 
 -- Helpers --------------------------------------------------------------------
-buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> TyCon -> VM Var
-buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc pdatas_tc
- = do   vectDataConWorkers orig_tc vect_tc pdata_tc
-        repr <- tyConRepr vect_tc
-        buildPADict vect_tc prepr_tc pdata_tc pdatas_tc repr
-      
 
+buildTyConPADict :: TyCon -> TyCon -> TyCon -> TyCon -> VM Var
+buildTyConPADict vect_tc prepr_tc pdata_tc pdatas_tc
+ = tyConRepr vect_tc >>= buildPADict vect_tc prepr_tc pdata_tc pdatas_tc
+
+-- Produce a custom-made worker for the data constructors of a vectorised data type.  This includes
+-- all data constructors that may be used in vetcorised code — i.e., all data constructors of data
+-- types other than scalar ones.  Also adds a mapping from the original to vectorised worker into
+-- the vectorisation map.
+--
+-- FIXME: It's not nice that we need create a special worker after the data constructors has
+--   already been constructed.  Also, I don't think the worker is properly added to the data
+--   constructor.  Seems messy.
 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
 vectDataConWorkers orig_tc vect_tc arr_tc
- = do bs <- sequence
-          . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
-          $ zipWith4 mk_data_con (tyConDataCons vect_tc)
-                                 rep_tys
-                                 (inits rep_tys)
-                                 (tail $ tails rep_tys)
-      mapM_ (uncurry hoistBinding) bs
- where
+  = do { traceVt "Building vectorised worker for datatype" (ppr orig_tc)
+  
+       ; bs <- sequence
+             . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
+             $ zipWith4 mk_data_con (tyConDataCons vect_tc)
+                                    rep_tys
+                                    (inits rep_tys)
+                                    (tail $ tails rep_tys)
+        ; mapM_ (uncurry hoistBinding) bs
+        }
+  where
     tyvars   = tyConTyVars vect_tc
     var_tys  = mkTyVarTys tyvars
     ty_args  = map Type var_tys
@@ -272,7 +298,6 @@ vectDataConWorkers orig_tc vect_tc arr_tc
 
     rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc
 
-
     mk_data_con con tys pre post
       = liftM2 (,) (vect_data_con con)
                    (lift_data_con tys pre post (mkDataConTag con))
diff --git a/compiler/vectorise/Vectorise/Type/PRepr.hs b/compiler/vectorise/Vectorise/Type/PRepr.hs
index 977815f51f..6e427ccec4 100644
--- a/compiler/vectorise/Vectorise/Type/PRepr.hs
+++ b/compiler/vectorise/Vectorise/Type/PRepr.hs
@@ -276,8 +276,8 @@ buildToArrPRepr vect_tc prepr_tc pdata_tc r
     to_comp expr (Wrap ty)
       = do
           wrap_tc  <- builtin wrapTyCon
-          (pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
-          return $ wrapNewTypeBody pwrap_tc [ty] expr
+          pwrap_tc <- pdataReprTyConExact (mkTyConApp wrap_tc [ty])
+        return $ wrapNewTypeBody pwrap_tc [ty] expr
 
 
 buildFromArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
@@ -358,7 +358,7 @@ buildFromArrPRepr vect_tc prepr_tc pdata_tc r
     from_comp _ res expr (Wrap ty)
       = do
           wrap_tc  <- builtin wrapTyCon
-          (pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
+          pwrap_tc <- pdataReprTyConExact (mkTyConApp wrap_tc [ty])
           return (res, [unwrapNewTypeBody pwrap_tc [ty]
                         $ unwrapFamInstScrut pwrap_tc [ty] expr])
 
diff --git a/compiler/vectorise/Vectorise/Utils/Base.hs b/compiler/vectorise/Vectorise/Utils/Base.hs
index c4dfe5c96b..a08174d513 100644
--- a/compiler/vectorise/Vectorise/Utils/Base.hs
+++ b/compiler/vectorise/Vectorise/Utils/Base.hs
@@ -15,7 +15,7 @@ module Vectorise.Utils.Base (
   mkBuiltinCo,
   mkVScrut,
 
-  pdataReprTyCon,   pdatasReprTyCon,
+  pdataReprTyCon, pdataReprTyConExact, pdatasReprTyCon,
   pdataReprDataCon, pdatasReprDataCon,
   prDFunOfTyCon
 ) where
@@ -28,12 +28,14 @@ import CoreSyn
 import CoreUtils
 import Coercion
 import Type
+import TypeRep
 import TyCon
 import DataCon
 import MkId
 import Literal
 import Outputable
 import FastString
+import ListSetOps
 
 import Control.Monad (liftM)
 
@@ -150,6 +152,7 @@ mkBuiltinCo get_tc
 
 
 -------------------------------------------------------------------------------
+
 mkVScrut :: VExpr -> VM (CoreExpr, CoreExpr, TyCon, [Type])
 mkVScrut (ve, le)
   = do
@@ -158,37 +161,57 @@ mkVScrut (ve, le)
   where
     ty = exprType ve
 
-
--- | Get the PData tycon that represents this type.
---   This tycon does not appear explicitly in the source program.
---   See Note [PData TyCons] in Vectorise.PRepr
+-- |Get the representation tycon of the 'PData' data family for a given type.
+--
+-- This tycon does not appear explicitly in the source program — see Note [PData TyCons] in
+-- 'Vectorise.Generic.Description':
 --
 --   @pdataReprTyCon {Sum2} = {PDataSum2}@
 --
+-- The type for which we look up a 'PData' instance may be more specific than the type in the
+-- instance declaration.  In that case the second component of the result will be more specific than
+-- a set of distinct type variables.
+-- 
 pdataReprTyCon :: Type -> VM (TyCon, [Type])
-pdataReprTyCon ty
-        = builtin pdataTyCon >>= (`lookupFamInst` [ty])
+pdataReprTyCon ty = builtin pdataTyCon >>= (`lookupFamInst` [ty])
+
+-- |Get the representation tycon of the 'PData' data family for a given type which must match the
+-- type index in the looked up 'PData' instance exactly.
+--
+pdataReprTyConExact :: Type -> VM TyCon
+pdataReprTyConExact ty
+  = do { (tycon, tys) <- pdataReprTyCon ty
+       ; if uniqueTyVars tys
+         then
+           return tycon
+         else
+           cantVectorise "No exact 'PData' family instance for" (ppr ty)
+       } 
+  where
+    uniqueTyVars tys = all isTyVarTy tys && hasNoDups (map extractTyVar tys)
+      where
+        extractTyVar (TyVarTy tv) = tv
+        extractTyVar _            = panic "Vectorise.Utils.Base: extractTyVar"
 
 pdataReprDataCon :: Type -> VM (DataCon, [Type])
 pdataReprDataCon ty
- = do   (tc, arg_tys) <- pdataReprTyCon ty
-        let [dc] = tyConDataCons tc
-        return (dc, arg_tys)
+  = do { (tc, arg_tys) <- pdataReprTyCon ty
+       ; let [dc] = tyConDataCons tc
+       ; return (dc, arg_tys)
+       }
 
 pdatasReprTyCon :: Type -> VM (TyCon, [Type])
-pdatasReprTyCon ty
-        = builtin pdatasTyCon >>= (`lookupFamInst` [ty])
+pdatasReprTyCon ty = builtin pdatasTyCon >>= (`lookupFamInst` [ty])
 
 pdatasReprDataCon :: Type -> VM (DataCon, [Type])
 pdatasReprDataCon ty
- = do   (tc, arg_tys) <- pdatasReprTyCon ty
-        let [dc] = tyConDataCons tc
-        return (dc, arg_tys)
-
+  = do { (tc, arg_tys) <- pdatasReprTyCon ty
+       ; let [dc] = tyConDataCons tc
+       ; return (dc, arg_tys)
+       }
 
 prDFunOfTyCon :: TyCon -> VM CoreExpr
 prDFunOfTyCon tycon
   = liftM Var
   . maybeCantVectoriseM "No PR dictionary for tycon" (ppr tycon)
   $ lookupTyConPR tycon
-