Find the correct array type for primitive tycons

compiler/vectorise/VectBuiltIn.hs

@@ -2,7 +2,7 @@ module VectBuiltIn (
   Builtins(..), sumTyCon, prodTyCon,
   initBuiltins, initBuiltinTyCons, initBuiltinPAs, initBuiltinPRs,
 
-  primMethod
+  primMethod, primPArray
 ) where
 
 #include "HsVersions.h"
@@ -20,6 +20,7 @@ import NameEnv
 import OccName
 
 import TypeRep         ( funTyCon )
+import Type            ( Type )
 import TysPrim
 import TysWiredIn      ( unitTyCon, tupleTyCon, intTyConName )
 import PrelNames
@@ -203,6 +204,14 @@ primMethod tycon method
 
   | otherwise = return Nothing
 
+primPArray :: TyCon -> DsM (Maybe TyCon)
+primPArray tycon
+  | Just suffix <- lookupNameEnv prim_ty_cons (tyConName tycon)
+  = liftM Just
+  $ dsLookupTyCon =<< lookupOrig nDP_PRIM (mkOccName tcName $ "PArray" ++ suffix)
+
+  | otherwise = return Nothing
+
 prim_ty_cons = mkNameEnv [mk_prim intPrimTyCon]
   where
     mk_prim tycon = (tyConName tycon, '_' : getOccString tycon)

compiler/vectorise/VectMonad.hs

@@ -24,7 +24,7 @@ module VectMonad (
   lookupDataCon, defDataCon,
   lookupTyConPA, defTyConPA, defTyConPAs,
   lookupTyConPR,
-  lookupPrimMethod,
+  lookupPrimMethod, lookupPrimPArray,
   lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
 
   {-lookupInst,-} lookupFamInst
@@ -355,8 +355,11 @@ defDataCon :: DataCon -> DataCon -> VM ()
 defDataCon dc dc' = updGEnv $ \env ->
   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 
+lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
+lookupPrimPArray = liftDs . primPArray
+
 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
-lookupPrimMethod tycon method = liftDs $ primMethod tycon method
+lookupPrimMethod tycon = liftDs . primMethod tycon
 
 lookupTyConPA :: TyCon -> VM (Maybe Var)
 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)

compiler/vectorise/VectUtils.hs

@@ -100,6 +100,14 @@ splitClosureTy = splitBinTy "splitClosureTy" closureTyConName
 splitPArrayTy :: Type -> Type
 splitPArrayTy = splitUnTy "splitPArrayTy" parrayTyConName
 
+splitPrimTyCon :: Type -> Maybe TyCon
+splitPrimTyCon ty
+  | Just (tycon, []) <- splitTyConApp_maybe ty
+  , isPrimTyCon tycon
+  = Just tycon
+
+  | otherwise = Nothing
+
 mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type
 mkBuiltinTyConApp get_tc tys
   = do
@@ -138,6 +146,12 @@ mkPADictType :: Type -> VM Type
 mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
 
 mkPArrayType :: Type -> VM Type
+mkPArrayType ty
+  | Just tycon <- splitPrimTyCon ty
+  = do
+      arr <- traceMaybeV "mkPArrayType" (ppr tycon)
+           $ lookupPrimPArray tycon
+      return $ mkTyConApp arr []
 mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty]
 
 mkBuiltinCo :: (Builtins -> TyCon) -> VM Coercion
@@ -229,8 +243,7 @@ pa_empty     = (emptyPAVar,     "emptyPA")
 
 paMethod :: PAMethod -> Type -> VM CoreExpr
 paMethod (method, name) ty
-  | Just (tycon, []) <- splitTyConApp_maybe ty
-  , isPrimTyCon tycon
+  | Just tycon <- splitPrimTyCon ty
   = do
       fn <- traceMaybeV "paMethod" (ppr tycon <+> text name)
           $ lookupPrimMethod tycon name