Incomplete support for boxing during vectorisation

4 files changed, 57 insertions, 3 deletions

diff --git a/compiler/vectorise/VectBuiltIn.hs b/compiler/vectorise/VectBuiltIn.hs
index 19df7ccc19..3975a19b2b 100644
--- a/compiler/vectorise/VectBuiltIn.hs
+++ b/compiler/vectorise/VectBuiltIn.hs
@@ -8,6 +8,7 @@
 module VectBuiltIn (
   Builtins(..), sumTyCon, prodTyCon, combinePAVar,
   initBuiltins, initBuiltinTyCons, initBuiltinPAs, initBuiltinPRs,
+  initBuiltinBoxedTyCons,
 
   primMethod, primPArray
 ) where
@@ -29,7 +30,7 @@ import OccName
 import TypeRep         ( funTyCon )
 import Type            ( Type )
 import TysPrim
-import TysWiredIn      ( unitTyCon, tupleTyCon, intTyConName )
+import TysWiredIn      ( unitTyCon, tupleTyCon, intTyCon, intTyConName )
 import Module
 import BasicTypes      ( Boxity(..) )
 
@@ -238,6 +239,13 @@ builtinPRs bi =
     mk_prod n = (tyConName $ prodTyCon n bi, nDP_REPR,
                  mkFastString ("dPR_" ++ show n))
 
+initBuiltinBoxedTyCons :: Builtins -> DsM [(Name, TyCon)]
+initBuiltinBoxedTyCons = return . builtinBoxedTyCons
+
+builtinBoxedTyCons :: Builtins -> [(Name, TyCon)]
+builtinBoxedTyCons bi =
+  [(tyConName intPrimTyCon, intTyCon)]
+
 externalVar :: Module -> FastString -> DsM Var
 externalVar mod fs
   = dsLookupGlobalId =<< lookupOrig mod (mkVarOccFS fs)
diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
index d91a60eb16..27f90f650c 100644
--- a/compiler/vectorise/VectMonad.hs
+++ b/compiler/vectorise/VectMonad.hs
@@ -31,6 +31,7 @@ module VectMonad (
   lookupDataCon, defDataCon,
   lookupTyConPA, defTyConPA, defTyConPAs,
   lookupTyConPR,
+  lookupBoxedTyCon,
   lookupPrimMethod, lookupPrimPArray,
   lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
 
@@ -102,6 +103,9 @@ data GlobalEnv = GlobalEnv {
                   -- Mapping from TyCons to their PR dfuns
                 , global_pr_funs :: NameEnv Var
 
+                  -- Mapping from unboxed TyCons to their boxed versions
+                , global_boxed_tycons :: NameEnv TyCon
+
                 -- External package inst-env & home-package inst-env for class
                 -- instances
                 --
@@ -142,6 +146,7 @@ initGlobalEnv info instEnvs famInstEnvs
     , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
     , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
     , global_pr_funs       = emptyNameEnv
+    , global_boxed_tycons  = emptyNameEnv
     , global_inst_env      = instEnvs
     , global_fam_inst_env  = famInstEnvs
     , global_bindings      = []
@@ -165,6 +170,10 @@ setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 setPRFunsEnv ps genv
   = genv { global_pr_funs = mkNameEnv ps }
 
+setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
+setBoxedTyConsEnv ps genv
+  = genv { global_boxed_tycons = mkNameEnv ps }
+
 emptyLocalEnv = LocalEnv {
                    local_vars     = emptyVarEnv
                  , local_tyvars   = []
@@ -389,6 +398,10 @@ lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 lookupTyConPR :: TyCon -> VM (Maybe Var)
 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
 
+lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
+lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
+                                                       (tyConName tc)
+
 defLocalTyVar :: TyVar -> VM ()
 defLocalTyVar tv = updLEnv $ \env ->
   env { local_tyvars   = tv : local_tyvars env
@@ -475,6 +488,7 @@ initV hsc_env guts info p
         let builtin_tycons = initBuiltinTyCons builtins
         builtin_pas    <- initBuiltinPAs builtins
         builtin_prs    <- initBuiltinPRs builtins
+        builtin_boxed  <- initBuiltinBoxedTyCons builtins
 
         eps <- ioToIOEnv $ hscEPS hsc_env
         let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
@@ -483,6 +497,7 @@ initV hsc_env guts info p
         let genv = extendTyConsEnv builtin_tycons
                  . extendPAFunsEnv builtin_pas
                  . setPRFunsEnv    builtin_prs
+                 . setBoxedTyConsEnv builtin_boxed
                  $ initGlobalEnv info instEnvs famInstEnvs
 
         r <- runVM p builtins genv emptyLocalEnv
diff --git a/compiler/vectorise/VectType.hs b/compiler/vectorise/VectType.hs
index 912eacfad7..c7046d4ba5 100644
--- a/compiler/vectorise/VectType.hs
+++ b/compiler/vectorise/VectType.hs
@@ -71,7 +71,7 @@ vectType (TyVarTy tv) = return $ TyVarTy tv
 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
-                                             (mapM vectType [ty1,ty2])
+                                             (mapM vectAndBoxType [ty1,ty2])
 vectType ty@(ForAllTy _ _)
   = do
       mdicts   <- mapM paDictArgType tyvars
@@ -82,6 +82,23 @@ vectType ty@(ForAllTy _ _)
 
 vectType ty = pprPanic "vectType:" (ppr ty)
 
+vectAndBoxType :: Type -> VM Type
+vectAndBoxType ty = vectType ty >>= boxType
+
+-- ----------------------------------------------------------------------------
+-- Boxing
+
+boxType :: Type -> VM Type
+boxType ty
+  | Just (tycon, []) <- splitTyConApp_maybe ty
+  , isUnLiftedTyCon tycon
+  = do
+      r <- lookupBoxedTyCon tycon
+      case r of
+        Just tycon' -> return $ mkTyConApp tycon' []
+        Nothing     -> return ty
+boxType ty = return ty
+
 -- ----------------------------------------------------------------------------
 -- Type definitions
 
@@ -285,7 +302,8 @@ boxedProductRepr tys
       tycon <- builtin (prodTyCon arity)
       let [data_con] = tyConDataCons tycon
 
-      (arr_tycon, _) <- parrayReprTyCon $ mkTyConApp tycon tys
+      tys' <- mapM boxType tys
+      (arr_tycon, _) <- parrayReprTyCon $ mkTyConApp tycon tys'
       let [arr_data_con] = tyConDataCons arr_tycon
 
       return $ ProdRepr {
diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 3e6143c398..a540b4d10c 100644
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -346,6 +346,19 @@ takeHoisted
       setGEnv $ env { global_bindings = [] }
       return $ global_bindings env
 
+boxExpr :: Type -> VExpr -> VM VExpr
+boxExpr ty (vexpr, lexpr)
+  | Just (tycon, []) <- splitTyConApp_maybe ty
+  , isUnLiftedTyCon tycon
+  = do
+      r <- lookupBoxedTyCon tycon
+      case r of
+        Just tycon' -> let [dc] = tyConDataCons tycon'
+                       in
+                       return (mkConApp dc [vexpr], lexpr)
+        Nothing     -> return (vexpr, lexpr)
+
+
 mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr
 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
   = do