Adapt PArray instance generation to new scheme
[ghc-hetmet.git] / compiler / vectorise / VectUtils.hs
index 5d03521..e71d2a6 100644 (file)
@@ -5,13 +5,13 @@ module VectUtils (
   splitClosureTy,
 
   TyConRepr(..), mkTyConRepr,
-  mkToPRepr, mkToArrPRepr, mkFromPRepr, mkFromArrPRepr,
+  mkToArrPRepr, mkFromArrPRepr,
   mkPADictType, mkPArrayType, mkPReprType,
 
   parrayCoerce, parrayReprTyCon, parrayReprDataCon, mkVScrut,
-  prDictOfType, prCoerce,
+  prDFunOfTyCon, prCoerce,
   paDictArgType, paDictOfType, paDFunType,
-  paMethod, lengthPA, replicatePA, emptyPA, liftPA,
+  paMethod, mkPR, lengthPA, replicatePA, emptyPA, liftPA,
   polyAbstract, polyApply, polyVApply,
   hoistBinding, hoistExpr, hoistPolyVExpr, takeHoisted,
   buildClosure, buildClosures,
@@ -37,10 +37,12 @@ import MkId               ( unwrapFamInstScrut )
 import Name               ( Name )
 import PrelNames
 import TysWiredIn
+import TysPrim            ( intPrimTy )
 import BasicTypes         ( Boxity(..) )
 
 import Outputable
 import FastString
+import Maybes             ( orElse )
 
 import Data.List             ( zipWith4 )
 import Control.Monad         ( liftM, liftM2, zipWithM_ )
@@ -94,9 +96,6 @@ splitFixedTyConApp tc ty
 
   | otherwise = pprPanic "splitFixedTyConApp" (ppr tc <+> ppr ty)
 
-splitEmbedTy :: Type -> Type
-splitEmbedTy = splitUnTy "splitEmbedTy" embedTyConName
-
 splitClosureTy :: Type -> (Type, Type)
 splitClosureTy = splitBinTy "splitClosureTy" closureTyConName
 
@@ -129,111 +128,73 @@ mkBuiltinTyConApps1 get_tc dft tys
     mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
 
 data TyConRepr = TyConRepr {
-                   repr_tyvars      :: [TyVar]
-                 , repr_tys         :: [[Type]]
-
-                 , repr_embed_tys   :: [[Type]]
-                 , repr_prod_tycons :: [Maybe TyCon]
-                 , repr_prod_tys    :: [Type]
-                 , repr_sum_tycon   :: Maybe TyCon
-                 , repr_type        :: Type
+                   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]
+                 , repr_prod_tys       :: [Type]
+                 , repr_sum_tycon      :: Maybe TyCon
+                 , repr_sum_data_cons  :: [DataCon]
+                 , repr_type           :: Type
                  }
 
 mkTyConRepr :: TyCon -> VM TyConRepr
 mkTyConRepr vect_tc
   = do
-      embed_tys <- mapM (mapM mkEmbedType) rep_tys
-      prod_tycons <- mapM (mk_tycon prodTyCon) rep_tys
-      sum_tycon   <- mk_tycon sumTyCon rep_tys
-
-      let prod_tys = zipWith mk_tc_app_maybe prod_tycons embed_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
+      arr_repr_tys <- mapM (mapM mkPArrayType . arr_repr_elem_tys) rep_tys
 
       return $ TyConRepr {
-                 repr_tyvars      = tyvars
-               , repr_tys         = rep_tys
-
-               , repr_embed_tys   = embed_tys
-               , repr_prod_tycons = prod_tycons
-               , repr_prod_tys    = prod_tys
-               , repr_sum_tycon   = sum_tycon
-               , repr_type        = mk_tc_app_maybe sum_tycon prod_tys
+                 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
+               , repr_prod_tys       = prod_tys
+               , repr_sum_tycon      = sum_tycon
+               , repr_sum_data_cons  = fmap tyConDataCons sum_tycon `orElse` []
+               , repr_type           = mk_tc_app_maybe sum_tycon prod_tys
                }
   where
     tyvars = tyConTyVars 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
       where n = length tys
 
+    mk_single_datacon tc | [dc] <- tyConDataCons tc = dc
+
     mk_tc_app_maybe Nothing   []   = unitTy
     mk_tc_app_maybe Nothing   [ty] = ty
     mk_tc_app_maybe (Just tc) tys  = mkTyConApp tc tys
 
-{-
-mkPRepr :: [[Type]] -> VM Type
-mkPRepr tys
-  = do
-      embed_tc <- builtin embedTyCon
-      sum_tcs  <- builtins sumTyCon
-      prod_tcs <- builtins prodTyCon
-
-      let mk_sum []   = unitTy
-          mk_sum [ty] = ty
-          mk_sum tys  = mkTyConApp (sum_tcs $ length tys) tys
-
-          mk_prod []   = unitTy
-          mk_prod [ty] = ty
-          mk_prod tys  = mkTyConApp (prod_tcs $ length tys) tys
-
-          mk_embed ty = mkTyConApp embed_tc [ty]
-
-      return . mk_sum
-             . map (mk_prod . map mk_embed)
-             $ tys
--}
-
-mkToPRepr :: [[CoreExpr]] -> VM ([CoreExpr], Type)
-mkToPRepr ess
-  = do
-      embed_tc <- builtin embedTyCon
-      embed_dc <- builtin embedDataCon
-      sum_tcs  <- builtins sumTyCon
-      prod_tcs <- builtins prodTyCon
-
-      let mk_sum [] = ([Var unitDataConId], unitTy)
-          mk_sum [(expr, ty)] = ([expr], ty)
-          mk_sum es = (zipWith mk_alt (tyConDataCons sum_tc) exprs,
-                       mkTyConApp sum_tc tys)
-            where
-              (exprs, tys)   = unzip es
-              sum_tc         = sum_tcs (length es)
-              mk_alt dc expr = mkConApp dc (map Type tys ++ [expr])
-
-          mk_prod [] = (Var unitDataConId, unitTy)
-          mk_prod [(expr, ty)] = (expr, ty)
-          mk_prod es = (mkConApp prod_dc (map Type tys ++ exprs),
-                        mkTyConApp prod_tc tys)
-            where
-              (exprs, tys) = unzip es
-              prod_tc      = prod_tcs (length es)
-              [prod_dc]    = tyConDataCons prod_tc
-
-          mk_embed expr = (mkConApp embed_dc [Type ty, expr],
-                           mkTyConApp embed_tc [ty])
-            where ty = exprType expr
-
-      return . mk_sum $ map (mk_prod . map mk_embed) ess
+    arr_repr_elem_tys []  = [unitTy]
+    arr_repr_elem_tys tys = tys
 
 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
@@ -246,75 +207,25 @@ mkToArrPRepr len sel ess
             where
               (exprs, tys) = unzip es
 
-          mk_prod [(expr, ty)] = return (expr, ty)
-          mk_prod es
+          mk_prod [expr] = return (expr, splitPArrayTy (exprType expr))
+          mk_prod exprs
             = do
-                prod_tc <- builtin . prodTyCon $ length es
+                prod_tc <- builtin . prodTyCon $ length exprs
                 (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)
+              tys = map (splitPArrayTy . exprType) exprs
 
-      liftM fst (mk_sum =<< mapM (mk_prod . map mk_embed) ess)
-
-mkFromPRepr :: CoreExpr -> Type -> [([Var], CoreExpr)] -> VM CoreExpr
-mkFromPRepr scrut res_ty alts
-  = do
-      embed_dc <- builtin embedDataCon
-      sum_tcs  <- builtins sumTyCon
-      prod_tcs <- builtins prodTyCon
-
-      let un_sum expr ty [(vars, res)] = un_prod expr ty vars res
-          un_sum expr ty bs
-            = do
-                ps     <- mapM (newLocalVar FSLIT("p")) tys
-                bodies <- sequence
-                        $ zipWith4 un_prod (map Var ps) tys vars rs
-                return . Case expr (mkWildId ty) res_ty
-                       $ zipWith3 mk_alt sum_dcs ps bodies
-            where
-              (vars, rs) = unzip bs
-              tys        = splitFixedTyConApp sum_tc ty
-              sum_tc     = sum_tcs $ length bs
-              sum_dcs    = tyConDataCons sum_tc
-
-              mk_alt dc p body = (DataAlt dc, [p], body)
-
-          un_prod expr ty []    r = return r
-          un_prod expr ty [var] r = return $ un_embed expr ty var r
-          un_prod expr ty vars  r
-            = do
-                xs <- mapM (newLocalVar FSLIT("x")) tys
-                let body = foldr (\(e,t,v) r -> un_embed e t v r) r
-                         $ zip3 (map Var xs) tys vars
-                return $ Case expr (mkWildId ty) res_ty
-                         [(DataAlt prod_dc, xs, body)]
-            where
-              tys       = splitFixedTyConApp prod_tc ty
-              prod_tc   = prod_tcs $ length vars
-              [prod_dc] = tyConDataCons prod_tc
-
-          un_embed expr ty var r
-            = Case expr (mkWildId ty) res_ty
-                [(DataAlt embed_dc, [var], r)]
-
-      un_sum scrut (exprType scrut) alts
+      liftM fst (mk_sum =<< mapM mk_prod ess)
 
 mkFromArrPRepr :: CoreExpr -> Type -> Var -> Var -> [[Var]] -> CoreExpr
                -> VM CoreExpr
 mkFromArrPRepr scrut res_ty len sel vars res
   = return (Var unitDataConId)
 
-mkEmbedType :: Type -> VM Type
-mkEmbedType ty = mkBuiltinTyConApp embedTyCon [ty]
-
 mkClosureType :: Type -> Type -> VM Type
 mkClosureType arg_ty res_ty = mkBuiltinTyConApp closureTyCon [arg_ty, res_ty]
 
@@ -357,35 +268,9 @@ mkVScrut (ve, le)
       (tc, arg_tys) <- parrayReprTyCon (exprType ve)
       return ((ve, unwrapFamInstScrut tc arg_tys le), tc, arg_tys)
 
-prDictOfType :: Type -> VM CoreExpr
-prDictOfType orig_ty
-  | Just (tycon, ty_args) <- splitTyConApp_maybe orig_ty
-  = do
-      dfun <- traceMaybeV "prDictOfType" (ppr tycon) (lookupTyConPR tycon)
-      prDFunApply (Var dfun) ty_args
-
-prDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
-prDFunApply dfun tys
-  = do
-      args <- mapM mkDFunArg arg_tys
-      return $ mkApps mono_dfun args
-  where
-    mono_dfun    = mkTyApps dfun tys
-    (arg_tys, _) = splitFunTys (exprType mono_dfun)
-
-mkDFunArg :: Type -> VM CoreExpr
-mkDFunArg ty
-  | Just (tycon, [arg]) <- splitTyConApp_maybe ty
-
-  = let name = tyConName tycon
-
-        get_dict | name == paTyConName = paDictOfType
-                 | name == prTyConName = prDictOfType
-                 | otherwise           = pprPanic "mkDFunArg" (ppr ty)
-
-    in get_dict arg
-
-mkDFunArg ty = pprPanic "mkDFunArg" (ppr ty)
+prDFunOfTyCon :: TyCon -> VM CoreExpr
+prDFunOfTyCon tycon
+  = liftM Var (traceMaybeV "prDictOfTyCon" (ppr tycon) (lookupTyConPR tycon))
 
 prCoerce :: TyCon -> [Type] -> CoreExpr -> VM CoreExpr
 prCoerce repr_tc args expr
@@ -460,6 +345,9 @@ paMethod method ty
       dict <- paDictOfType ty
       return $ mkApps (Var fn) [Type ty, dict]
 
+mkPR :: Type -> VM CoreExpr
+mkPR = paMethod mkPRVar
+
 lengthPA :: CoreExpr -> VM CoreExpr
 lengthPA x = liftM (`App` x) (paMethod lengthPAVar ty)
   where