Add code for looking up PA methods of primitive TyCons
[ghc-hetmet.git] / compiler / vectorise / VectUtils.hs
index e71d2a6..3c9d921 100644 (file)
@@ -4,12 +4,11 @@ module VectUtils (
   mkDataConTag,
   splitClosureTy,
 
-  TyConRepr(..), mkTyConRepr,
-  mkToArrPRepr, mkFromArrPRepr,
+  mkBuiltinCo,
   mkPADictType, mkPArrayType, mkPReprType,
 
-  parrayCoerce, parrayReprTyCon, parrayReprDataCon, mkVScrut,
-  prDFunOfTyCon, prCoerce,
+  parrayReprTyCon, parrayReprDataCon, mkVScrut,
+  prDFunOfTyCon,
   paDictArgType, paDictOfType, paDFunType,
   paMethod, mkPR, lengthPA, replicatePA, emptyPA, liftPA,
   polyAbstract, polyApply, polyVApply,
@@ -42,7 +41,6 @@ import BasicTypes         ( Boxity(..) )
 
 import Outputable
 import FastString
-import Maybes             ( orElse )
 
 import Data.List             ( zipWith4 )
 import Control.Monad         ( liftM, liftM2, zipWithM_ )
@@ -127,105 +125,6 @@ mkBuiltinTyConApps1 get_tc dft tys
   where
     mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
 
-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]
-                 , 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
-      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
-               , 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
-
-    arr_repr_elem_tys []  = [unitTy]
-    arr_repr_elem_tys tys = tys
-
-mkToArrPRepr :: CoreExpr -> CoreExpr -> [[CoreExpr]] -> VM CoreExpr
-mkToArrPRepr len sel ess
-  = do
-      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] = return (expr, splitPArrayTy (exprType expr))
-          mk_prod exprs
-            = do
-                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
-              tys = map (splitPArrayTy . exprType) exprs
-
-      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)
-
 mkClosureType :: Type -> Type -> VM Type
 mkClosureType arg_ty res_ty = mkBuiltinTyConApp closureTyCon [arg_ty, res_ty]
 
@@ -241,16 +140,11 @@ 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
+mkBuiltinCo :: (Builtins -> TyCon) -> VM Coercion
+mkBuiltinCo get_tc
   = do
-      parray <- builtin parrayTyCon
-
-      let co = mkAppCoercion (mkTyConApp parray [])
-                             (mkSymCoercion (mkTyConApp arg_co args))
-
-      return $ mkCoerce co expr
+      tc <- builtin get_tc
+      return $ mkTyConApp tc []
 
 parrayReprTyCon :: Type -> VM (TyCon, [Type])
 parrayReprTyCon ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
@@ -272,17 +166,6 @@ 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
-  | Just arg_co <- tyConFamilyCoercion_maybe repr_tc
-  = do
-      pr_tc <- builtin prTyCon
-
-      let co = mkAppCoercion (mkTyConApp pr_tc [])
-                             (mkSymCoercion (mkTyConApp arg_co args))
-
-      return $ mkCoerce co expr
-
 paDictArgType :: TyVar -> VM (Maybe Type)
 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
   where
@@ -338,27 +221,45 @@ paDFunApply dfun tys
       dicts <- mapM paDictOfType tys
       return $ mkApps (mkTyApps dfun tys) dicts
 
-paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr
-paMethod method ty
+type PAMethod = (Builtins -> Var, String)
+
+pa_length    = (lengthPAVar,    "lengthPA")
+pa_replicate = (replicatePAVar, "replicatePA")
+pa_empty     = (emptyPAVar,     "emptyPA")
+
+paMethod :: PAMethod -> Type -> VM CoreExpr
+paMethod (method, name) ty
+  | Just (tycon, []) <- splitTyConApp_maybe ty
+  , isPrimTyCon tycon
+  = do
+      fn <- traceMaybeV "paMethod" (ppr tycon <+> text name)
+          $ lookupPrimMethod tycon name
+      return (Var fn)
+
+paMethod (method, name) ty
   = do
       fn   <- builtin method
       dict <- paDictOfType ty
       return $ mkApps (Var fn) [Type ty, dict]
 
 mkPR :: Type -> VM CoreExpr
-mkPR = paMethod mkPRVar
+mkPR ty
+  = do
+      fn   <- builtin mkPRVar
+      dict <- paDictOfType ty
+      return $ mkApps (Var fn) [Type ty, dict]
 
 lengthPA :: CoreExpr -> VM CoreExpr
-lengthPA x = liftM (`App` x) (paMethod lengthPAVar ty)
+lengthPA x = liftM (`App` x) (paMethod pa_length ty)
   where
     ty = splitPArrayTy (exprType x)
 
 replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
 replicatePA len x = liftM (`mkApps` [len,x])
-                          (paMethod replicatePAVar (exprType x))
+                          (paMethod pa_replicate (exprType x))
 
 emptyPA :: Type -> VM CoreExpr
-emptyPA = paMethod emptyPAVar
+emptyPA = paMethod pa_empty
 
 liftPA :: CoreExpr -> VM CoreExpr
 liftPA x