vectoriser: Use Sels2 type for vector of selectors in PDatas Sum2 instance

6 files changed, 31 insertions, 18 deletions

diff --git a/compiler/vectorise/Vectorise/Builtins.hs b/compiler/vectorise/Vectorise/Builtins.hs
index bf0fae1c11..9eb24f1df2 100644
--- a/compiler/vectorise/Vectorise/Builtins.hs
+++ b/compiler/vectorise/Vectorise/Builtins.hs
@@ -9,7 +9,7 @@ module Vectorise.Builtins (
   
   -- * Wrapped selectors
   parray_PrimTyCon,
-  selTy,
+  selTy, selsTy,
   selReplicate,
   selTags,
   selElements,
diff --git a/compiler/vectorise/Vectorise/Builtins/Base.hs b/compiler/vectorise/Vectorise/Builtins/Base.hs
index 4ed351d120..2f3990e9d9 100644
--- a/compiler/vectorise/Vectorise/Builtins/Base.hs
+++ b/compiler/vectorise/Vectorise/Builtins/Base.hs
@@ -14,7 +14,7 @@ module Vectorise.Builtins.Base (
   
   -- * Projections
   parray_PrimTyCon,
-  selTy,
+  selTy, selsTy,
   selReplicate,
   selTags,
   selElements,
@@ -105,6 +105,7 @@ data Builtins
         , liftedApplyVar       :: Var                       -- ^ liftedApply
         , closureCtrFuns       :: Array Int Var             -- ^ closure1 .. closure3
         , selTys               :: Array Int Type            -- ^ Sel2
+        , selsTys              :: Array Int Type            -- ^ Sel2s
         , selReplicates        :: Array Int CoreExpr        -- ^ replicate2
         , selTagss             :: Array Int CoreExpr        -- ^ tagsSel2
         , selElementss         :: Array (Int, Int) CoreExpr -- ^ elementsSel2_0 .. elementsSel_2_1
@@ -122,6 +123,9 @@ parray_PrimTyCon tc bi = lookupEnvBuiltin "parray_PrimTyCon" (parray_PrimTyCons
 selTy :: Int -> Builtins -> Type
 selTy           = indexBuiltin "selTy" selTys
 
+selsTy :: Int -> Builtins -> Type
+selsTy          = indexBuiltin "selsTy" selsTys
+
 selReplicate :: Int -> Builtins -> CoreExpr
 selReplicate    = indexBuiltin "selReplicate" selReplicates 
 
diff --git a/compiler/vectorise/Vectorise/Builtins/Initialise.hs b/compiler/vectorise/Vectorise/Builtins/Initialise.hs
index 1d48aa369b..5a38d73d52 100644
--- a/compiler/vectorise/Vectorise/Builtins/Initialise.hs
+++ b/compiler/vectorise/Vectorise/Builtins/Initialise.hs
@@ -102,11 +102,13 @@ initBuiltins
       ; let closureCtrFuns = listArray (1, mAX_DPH_COMBINE) closures
 
           -- Types and functions for selectors
-      ; sel_tys          <- mapM externalType (numbered "Sel" 2 mAX_DPH_SUM)
+      ; sel_tys          <- mapM externalType (numbered "Sel"  2 mAX_DPH_SUM)
+      ; sels_tys         <- mapM externalType (numbered "Sels" 2 mAX_DPH_SUM)
       ; sel_replicates   <- mapM externalFun  (numbered_hash "replicateSel" 2 mAX_DPH_SUM)
       ; sel_tags         <- mapM externalFun  (numbered "tagsSel" 2 mAX_DPH_SUM)
       ; sel_elements     <- mapM mk_elements [(i,j) | i <- [2..mAX_DPH_SUM], j <- [0..i-1]]
       ; let selTys        = listArray (2, mAX_DPH_SUM) sel_tys
+            selsTys       = listArray (2, mAX_DPH_SUM) sels_tys
             selReplicates = listArray (2, mAX_DPH_SUM) sel_replicates
             selTagss      = listArray (2, mAX_DPH_SUM) sel_tags
             selElementss  = array     ((2, 0), (mAX_DPH_SUM, mAX_DPH_SUM)) sel_elements
@@ -150,6 +152,7 @@ initBuiltins
                , liftedApplyVar       = liftedApplyVar
                , closureCtrFuns       = closureCtrFuns
                , selTys               = selTys
+               , selsTys              = selsTys
                , selReplicates        = selReplicates
                , selTagss             = selTagss
                , selElementss         = selElementss
diff --git a/compiler/vectorise/Vectorise/Generic/Description.hs b/compiler/vectorise/Vectorise/Generic/Description.hs
index 8a60d57f79..9bffff1302 100644
--- a/compiler/vectorise/Vectorise/Generic/Description.hs
+++ b/compiler/vectorise/Vectorise/Generic/Description.hs
@@ -46,9 +46,12 @@ data SumRepr
                -- | PDatas version of the sum tycon    (eg PDatasSum2)
                , repr_psums_tc  :: TyCon
 
-               -- | Type of the selector (eg Sel2)
+               -- | Type of the selector               (eg Sel2)
                , repr_sel_ty    :: Type
 
+               -- | Type of multi-selector             (eg Sel2s)
+               , repr_sels_ty   :: Type
+
                -- | Type of each data constructor.
                , repr_con_tys   :: [Type]
 
@@ -128,11 +131,13 @@ tyConRepr tc
            psums_tc     <- liftM fst $ pdatasReprTyCon sumapp
            
            sel_ty       <- builtin (selTy arity)
+           sels_ty      <- builtin (selsTy arity)
            return $ Sum 
                   { repr_sum_tc   = sum_tc
                   , repr_psum_tc  = psum_tc
                   , repr_psums_tc = psums_tc
                   , repr_sel_ty   = sel_ty
+                  , repr_sels_ty  = sels_ty
                   , repr_con_tys  = tys
                   , repr_cons     = rs
                   }
@@ -217,12 +222,13 @@ instance Outputable SumRepr where
         UnarySum con
          -> sep [text "UnarySum", ppr con]
 
-        Sum sumtc psumtc psumstc selty contys cons
+        Sum sumtc psumtc psumstc selty selsty contys cons
          -> text "Sum" $+$ braces (nest 4 
                 $ sep   [ text "repr_sum_tc   = " <> ppr sumtc
                         , text "repr_psum_tc  = " <> ppr psumtc
                         , text "repr_psums_tc = " <> ppr psumstc
                         , text "repr_sel_ty   = " <> ppr selty
+                        , text "repr_sels_ty  = " <> ppr selsty
                         , text "repr_con_tys  = " <> ppr contys
                         , text "repr_cons     = " <> ppr cons])
 
diff --git a/compiler/vectorise/Vectorise/Generic/PAMethods.hs b/compiler/vectorise/Vectorise/Generic/PAMethods.hs
index 6330dddf64..7287a6db5e 100644
--- a/compiler/vectorise/Vectorise/Generic/PAMethods.hs
+++ b/compiler/vectorise/Vectorise/Generic/PAMethods.hs
@@ -410,22 +410,22 @@ buildToArrPReprs vect_tc prepr_tc _ pdatas_tc r
     [pdatas_dc] = tyConDataCons pdatas_tc
          
     to_sum ss
-     = case ss of       -- BROKEN: should be
-        EmptySum    -> builtin pvoidVar >>= \pvoid -> return ([], Var pvoid) 
+     = case ss of
+        EmptySum    -> builtin pvoidsVar >>= \pvoids -> return ([], Var pvoids) 
         UnarySum r  -> to_con r
         Sum{}
           -> do let psums_tc     = repr_psums_tc ss
                 let [psums_con]  = tyConDataCons psums_tc
                 (vars, exprs)   <- mapAndUnzipM to_con (repr_cons ss)
-                sel             <- newLocalVar (fsLit "sel") (repr_sel_ty ss) -- BROKEN: should be vector
+                sel             <- newLocalVar (fsLit "sels") (repr_sels_ty ss)
                 return ( sel : concat vars
                        , wrapFamInstBody psums_tc (repr_con_tys ss)
                          $ mkConApp psums_con 
                          $ map Type (repr_con_tys ss) ++ (Var sel : exprs))        
 
     to_prod ss
-     = case ss of       -- BROKEN: should be pvoids
-        EmptyProd    -> builtin pvoidVar >>= \pvoid -> return ([], Var pvoid)
+     = case ss of
+        EmptyProd    -> builtin pvoidsVar >>= \pvoids -> return ([], Var pvoids)
         UnaryProd r
          -> do  pty  <- mkPDatasType (compOrigType r)
                 var  <- newLocalVar (fsLit "x") pty
@@ -501,7 +501,7 @@ buildFromArrPReprs vect_tc prepr_tc _ pdatas_tc r
         Sum {}
          -> do  let psums_tc    =  repr_psums_tc ss
                 let [psums_con] =  tyConDataCons psums_tc
-                sel             <- newLocalVar (fsLit "sel") (repr_sel_ty ss)
+                sel             <- newLocalVar (fsLit "sels") (repr_sels_ty ss)
                 ptys            <- mapM mkPDatasType (repr_con_tys ss)
                 vars            <- newLocalVars (fsLit "xs") ptys
                 (res', args)    <- fold from_con res_ty res (map Var vars) (repr_cons ss)
diff --git a/compiler/vectorise/Vectorise/Generic/PData.hs b/compiler/vectorise/Vectorise/Generic/PData.hs
index f10afff972..3587452951 100644
--- a/compiler/vectorise/Vectorise/Generic/PData.hs
+++ b/compiler/vectorise/Vectorise/Generic/PData.hs
@@ -57,7 +57,7 @@ buildPDataDataCon :: Name -> TyCon -> TyCon -> SumRepr -> VM DataCon
 buildPDataDataCon orig_name vect_tc repr_tc repr
  = do let tvs   = tyConTyVars vect_tc
       dc_name   <- mkLocalisedName mkPDataDataConOcc orig_name
-      comp_tys  <- mkSumTys mkPDataType repr
+      comp_tys  <- mkSumTys repr_sel_ty mkPDataType repr
 
       liftDs $ buildDataCon dc_name
                             False                  -- not infix
@@ -106,7 +106,7 @@ buildPDatasDataCon orig_name vect_tc repr_tc repr
  = do let tvs   = tyConTyVars vect_tc
       dc_name        <- mkLocalisedName mkPDatasDataConOcc orig_name
 
-      comp_tys  <- mkSumTys mkPDatasType repr
+      comp_tys  <- mkSumTys repr_sels_ty mkPDatasType repr
 
       liftDs $ buildDataCon dc_name
                             False                  -- not infix
@@ -124,18 +124,18 @@ buildPDatasDataCon orig_name vect_tc repr_tc repr
 -- Utils ----------------------------------------------------------------------
 -- | Flatten a SumRepr into a list of data constructor types.
 mkSumTys 
-        :: (Type -> VM Type)
+        :: (SumRepr -> Type)
+        -> (Type -> VM Type)
         -> SumRepr
         -> VM [Type]
 
-mkSumTys mkTc repr
+mkSumTys repr_selX_ty mkTc repr
  = sum_tys repr
  where
     sum_tys EmptySum      = return []
     sum_tys (UnarySum r)  = con_tys r
-    sum_tys (Sum { repr_sel_ty = sel_ty
-                 , repr_cons   = cons })
-      = liftM (sel_ty :) (concatMapM con_tys cons)
+    sum_tys d@(Sum { repr_cons   = cons })
+      = liftM (repr_selX_ty d :) (concatMapM con_tys cons)
 
     con_tys (ConRepr _ r)  = prod_tys r