Adapt PArray instance generation to new scheme

2 files changed, 27 insertions, 9 deletions

diff --git a/compiler/vectorise/VectType.hs b/compiler/vectorise/VectType.hs
index d47f391404..10c3bbffc9 100644
--- a/compiler/vectorise/VectType.hs
+++ b/compiler/vectorise/VectType.hs
@@ -398,22 +398,20 @@ buildPArrayDataCon :: Name -> TyCon -> TyCon -> VM DataCon
 buildPArrayDataCon orig_name vect_tc repr_tc
   = do
       dc_name  <- cloneName mkPArrayDataConOcc orig_name
-      shape    <- tyConShape vect_tc
-      repr_tys <- mapM mkPArrayType types
+      repr     <- mkTyConRepr vect_tc
+
+      let all_tys = arr_shape_tys repr ++ concat (arr_repr_tys repr)
 
       liftDs $ buildDataCon dc_name
                             False                  -- not infix
-                            (shapeStrictness shape ++ map (const NotMarkedStrict) repr_tys)
+                            (map (const NotMarkedStrict) all_tys)
                             []                     -- no field labels
                             (tyConTyVars vect_tc)
                             []                     -- no existentials
                             []                     -- no eq spec
                             []                     -- no context
-                            (shapeReprTys shape ++ repr_tys)
+                            all_tys
                             repr_tc
-  where
-    types = [ty | dc <- tyConDataCons vect_tc
-                , ty <- dataConRepArgTys dc]
 
 mkPADFun :: TyCon -> VM Var
 mkPADFun vect_tc
diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index a1f554db50..e71d2a6484 100644
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -37,6 +37,7 @@ import MkId               ( unwrapFamInstScrut )
 import Name               ( Name )
 import PrelNames
 import TysWiredIn
+import TysPrim            ( intPrimTy )
 import BasicTypes         ( Boxity(..) )
 
 import Outputable
@@ -129,6 +130,8 @@ mkBuiltinTyConApps1 get_tc dft tys
 data TyConRepr = TyConRepr {
                    repr_tyvars         :: [TyVar]
                  , repr_tys            :: [[Type]]
+                 , arr_shape_tys       :: [Type]
+                 , arr_repr_tys        :: [[Type]]
 
                  , repr_prod_tycons    :: [Maybe TyCon]
                  , repr_prod_data_cons :: [Maybe DataCon]
@@ -141,13 +144,17 @@ data TyConRepr = TyConRepr {
 mkTyConRepr :: TyCon -> VM TyConRepr
 mkTyConRepr vect_tc
   = do
-      prod_tycons <- mapM (mk_tycon prodTyCon) rep_tys
+      uarr <- builtin uarrTyCon
+      prod_tycons  <- mapM (mk_tycon prodTyCon) rep_tys
       let prod_tys = zipWith mk_tc_app_maybe prod_tycons rep_tys
-      sum_tycon   <- mk_tycon sumTyCon prod_tys
+      sum_tycon    <- mk_tycon sumTyCon prod_tys
+      arr_repr_tys <- mapM (mapM mkPArrayType . arr_repr_elem_tys) rep_tys
 
       return $ TyConRepr {
                  repr_tyvars         = tyvars
                , repr_tys            = rep_tys
+               , arr_shape_tys       = mk_shape uarr
+               , arr_repr_tys        = arr_repr_tys
 
                , repr_prod_tycons    = prod_tycons
                , repr_prod_data_cons = map (fmap mk_single_datacon) prod_tycons
@@ -161,6 +168,16 @@ mkTyConRepr vect_tc
     data_cons = tyConDataCons vect_tc
     rep_tys   = map dataConRepArgTys data_cons
 
+    is_product | [_] <- data_cons = True
+               | otherwise        = False
+
+    mk_shape uarr = intPrimTy : mk_sel uarr
+
+    mk_sel uarr | is_product = []
+                | otherwise  = [uarr_int, uarr_int]
+      where
+        uarr_int = mkTyConApp uarr [intTy]
+
     mk_tycon get_tc tys
       | n > 1     = builtin (Just . get_tc n)
       | otherwise = return Nothing
@@ -172,6 +189,9 @@ mkTyConRepr vect_tc
     mk_tc_app_maybe Nothing   [ty] = ty
     mk_tc_app_maybe (Just tc) tys  = mkTyConApp tc tys
 
+    arr_repr_elem_tys []  = [unitTy]
+    arr_repr_elem_tys tys = tys
+
 mkToArrPRepr :: CoreExpr -> CoreExpr -> [[CoreExpr]] -> VM CoreExpr
 mkToArrPRepr len sel ess
   = do