compiler/vectorise/Vectorise/Builtins.hs

   1
   2 -- | Builtin types and functions used by the vectoriser.
   3 --   The source program uses functions from GHC.PArr, which the vectoriser rewrites
   4 --   to use equivalent vectorised versions in the DPH backend packages.
   5 --
   6 --   The `Builtins` structure holds the name of all the things in the DPH packages
   7 --   we will need. We can get specific things using the selectors, which print a
   8 --   civilized panic message if the specified thing cannot be found.
   9 --
  10 module Vectorise.Builtins (
  11         -- * Builtins
  12         Builtins(..),
  13         indexBuiltin,
  14
  15         -- * Wrapped selectors
  16         selTy,
  17         selReplicate,
  18         selPick,
  19         selTags,
  20         selElements,
  21         sumTyCon,
  22         prodTyCon,
  23         prodDataCon,
  24         combinePDVar,
  25         scalarZip,
  26         closureCtrFun,
  27
  28         -- * Initialisation
  29         initBuiltins, initBuiltinVars, initBuiltinTyCons, initBuiltinDataCons,
  30         initBuiltinPAs, initBuiltinPRs,
  31         initBuiltinBoxedTyCons, initBuiltinScalars,
  32
  33         -- * Lookup
  34         primMethod,
  35         primPArray
  36 ) where
  37 import Vectorise.Builtins.Base
  38 import Vectorise.Builtins.Modules
  39 import Vectorise.Builtins.Initialise
  40
  41 import TysPrim
  42 import IfaceEnv
  43 import TyCon
  44 import DsMonad
  45 import NameEnv
  46 import Name
  47 import Var
  48 import Control.Monad
  49
  50
  51 -- | Lookup a method function given its name and instance type.
  52 primMethod :: TyCon -> String -> Builtins -> DsM (Maybe Var)
  53 primMethod  tycon method (Builtins { dphModules = mods })
  54   | Just suffix <- lookupNameEnv prim_ty_cons (tyConName tycon)
  55   = liftM Just
  56   $ dsLookupGlobalId =<< lookupOrig (dph_Unboxed mods)
  57                                     (mkVarOcc $ method ++ suffix)
  58
  59   | otherwise = return Nothing
  60
  61 -- | Lookup the representation type we use for PArrays that contain a given element type.
  62 primPArray :: TyCon -> Builtins -> DsM (Maybe TyCon)
  63 primPArray tycon (Builtins { dphModules = mods })
  64   | Just suffix <- lookupNameEnv prim_ty_cons (tyConName tycon)
  65   = liftM Just
  66   $ dsLookupTyCon =<< lookupOrig (dph_Unboxed mods)
  67                                  (mkTcOcc $ "PArray" ++ suffix)
  68
  69   | otherwise = return Nothing
  70
  71 prim_ty_cons :: NameEnv String
  72 prim_ty_cons = mkNameEnv [mk_prim intPrimTyCon]
  73   where
  74     mk_prim tycon = (tyConName tycon, '_' : getOccString tycon)
  75