Conversions to/from generic array representation (not finished yet)

2 files changed, 101 insertions, 7 deletions

diff --git a/compiler/vectorise/VectType.hs b/compiler/vectorise/VectType.hs
index a7c463b462..aa0eae2f10 100644
--- a/compiler/vectorise/VectType.hs
+++ b/compiler/vectorise/VectType.hs
@@ -230,6 +230,43 @@ buildToPRepr _ vect_tc prepr_tc _
 
     mk_alt data_con bndrs body = (DataAlt data_con, bndrs, body)
 
+buildToArrPRepr :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
+buildToArrPRepr _ vect_tc prepr_tc arr_tc
+  = do
+      arg_ty  <- mkPArrayType el_ty
+      rep_tys <- mapM (mapM mkPArrayType) rep_el_tys
+
+      arg     <- newLocalVar FSLIT("xs") arg_ty
+      bndrss  <- mapM (mapM (newLocalVar FSLIT("ys"))) rep_tys
+      len     <- newLocalVar FSLIT("len") intPrimTy
+      sel     <- newLocalVar FSLIT("sel") =<< mkPArrayType intTy
+
+      let add_sel xs | has_selector = sel : xs
+                     | otherwise    = xs
+
+          all_bndrs = len : add_sel (concat bndrss)
+
+      res      <- parrayCoerce prepr_tc var_tys
+                =<< mkToArrPRepr (Var len) (Var sel) (map (map Var) bndrss)
+      res_ty   <- mkPArrayType =<< mkPReprType el_ty
+
+      return . Lam arg
+             $ Case (unwrapFamInstScrut arr_tc var_tys (Var arg))
+                    (mkWildId (mkTyConApp arr_tc var_tys))
+                    res_ty
+                    [(DataAlt arr_dc, all_bndrs, res)]
+  where
+    var_tys    = mkTyVarTys $ tyConTyVars vect_tc
+    el_ty      = mkTyConApp vect_tc var_tys
+    data_cons  = tyConDataCons vect_tc
+    rep_el_tys = map dataConRepArgTys data_cons
+
+    [arr_dc]   = tyConDataCons arr_tc
+
+    has_selector | [_] <- data_cons = False
+                 | otherwise        = True
+
+
 buildFromPRepr :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
 buildFromPRepr _ vect_tc prepr_tc _
   = do
@@ -248,6 +285,10 @@ buildFromPRepr _ vect_tc prepr_tc _
                   bndrs <- mapM (newLocalVar FSLIT("x")) $ dataConRepArgTys dc
                   return (bndrs, mkConApp dc (map Type var_tys ++ map Var bndrs))
 
+buildFromArrPRepr :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
+buildFromArrPRepr _ vect_tc prepr_tc arr_tc
+  = mkFromArrPRepr undefined undefined undefined undefined undefined undefined
+
 buildPRDict :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
 buildPRDict _ vect_tc prepr_tc _
   = prCoerce prepr_tc var_tys
@@ -420,11 +461,11 @@ buildPADict shape vect_tc prepr_tc arr_tc dfun
           var  <- newLocalVar name (exprType body)
           return (var, mkInlineMe body)
           
-paMethods = [(FSLIT("lengthPA"),    buildLengthPA),
-             (FSLIT("replicatePA"), buildReplicatePA),
-             (FSLIT("toPRepr"),     buildToPRepr),
-             (FSLIT("fromPRepr"),   buildFromPRepr),
-             (FSLIT("dictPRepr"),   buildPRDict)]
+paMethods = [(FSLIT("toPRepr"),      buildToPRepr),
+             (FSLIT("fromPRepr"),    buildFromPRepr),
+             (FSLIT("toArrPRepr"),   buildToArrPRepr),
+             (FSLIT("fromArrPRepr"), buildFromArrPRepr),
+             (FSLIT("dictPRepr"),    buildPRDict)]
 
 buildLengthPA :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
 buildLengthPA shape vect_tc _ arr_tc
diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 0f101bd3cb..acbbe45b9a 100644
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -3,9 +3,9 @@ module VectUtils (
   collectAnnValBinders,
   mkDataConTag,
   splitClosureTy,
-  mkPRepr, mkToPRepr, mkFromPRepr,
+  mkPRepr, mkToPRepr, mkToArrPRepr, mkFromPRepr, mkFromArrPRepr,
   mkPADictType, mkPArrayType, mkPReprType,
-  parrayReprTyCon, parrayReprDataCon, mkVScrut,
+  parrayCoerce, parrayReprTyCon, parrayReprDataCon, mkVScrut,
   prDictOfType, prCoerce,
   paDictArgType, paDictOfType, paDFunType,
   paMethod, lengthPA, replicatePA, emptyPA, liftPA,
@@ -178,6 +178,43 @@ mkToPRepr ess
 
       return . mk_sum $ map (mk_prod . map mk_embed) ess
 
+mkToArrPRepr :: CoreExpr -> CoreExpr -> [[CoreExpr]] -> VM CoreExpr
+mkToArrPRepr len sel ess
+  = do
+      embed_tc <- builtin embedTyCon
+      (embed_rtc, _) <- parrayReprTyCon (mkTyConApp embed_tc [unitTy])
+      let [embed_rdc] = tyConDataCons embed_rtc
+
+      let mk_sum [(expr, ty)] = return (expr, ty)
+          mk_sum es
+            = do
+                sum_tc <- builtin . sumTyCon $ length es
+                (sum_rtc, _) <- parrayReprTyCon (mkTyConApp sum_tc tys)
+                let [sum_rdc] = tyConDataCons sum_rtc
+
+                return (mkConApp sum_rdc (map Type tys ++ (len : sel : exprs)),
+                        mkTyConApp sum_tc tys)
+            where
+              (exprs, tys) = unzip es
+
+          mk_prod [(expr, ty)] = return (expr, ty)
+          mk_prod es
+            = do
+                prod_tc <- builtin . prodTyCon $ length es
+                (prod_rtc, _) <- parrayReprTyCon (mkTyConApp prod_tc tys)
+                let [prod_rdc] = tyConDataCons prod_rtc
+
+                return (mkConApp prod_rdc (map Type tys ++ (len : exprs)),
+                        mkTyConApp prod_tc tys)
+            where
+              (exprs, tys) = unzip es
+
+          mk_embed expr = (mkConApp embed_rdc [Type ty, expr],
+                           mkTyConApp embed_tc [ty])
+            where ty = splitPArrayTy (exprType expr)
+
+      liftM fst (mk_sum =<< mapM (mk_prod . map mk_embed) ess)
+
 mkFromPRepr :: CoreExpr -> Type -> [([Var], CoreExpr)] -> VM CoreExpr
 mkFromPRepr scrut res_ty alts
   = do
@@ -221,6 +258,11 @@ mkFromPRepr scrut res_ty alts
 
       un_sum scrut (exprType scrut) alts
 
+mkFromArrPRepr :: CoreExpr -> Type -> Var -> Var -> [[Var]] -> CoreExpr
+               -> VM CoreExpr
+mkFromArrPRepr scrut res_ty len sel vars res
+  = return (Var unitDataConId)
+
 mkClosureType :: Type -> Type -> VM Type
 mkClosureType arg_ty res_ty = mkBuiltinTyConApp closureTyCon [arg_ty, res_ty]
 
@@ -236,6 +278,17 @@ mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
 mkPArrayType :: Type -> VM Type
 mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty]
 
+parrayCoerce :: TyCon -> [Type] -> CoreExpr -> VM CoreExpr
+parrayCoerce repr_tc args expr
+  | Just arg_co <- tyConFamilyCoercion_maybe repr_tc
+  = do
+      parray <- builtin parrayTyCon
+
+      let co = mkAppCoercion (mkTyConApp parray [])
+                             (mkSymCoercion (mkTyConApp arg_co args))
+
+      return $ mkCoerce co expr
+
 parrayReprTyCon :: Type -> VM (TyCon, [Type])
 parrayReprTyCon ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])