compiler/vectorise/Vectorise/Type/Env.hs

   1 {-# OPTIONS_GHC -XNoMonoLocalBinds -fno-warn-missing-signatures #-}
   2 -- Roman likes local bindings
   3 -- If this module lives on I'd like to get rid of this flag in due course
   4
   5 module Vectorise.Type.Env (
   6         vectTypeEnv,
   7 )
   8 where
   9 import Vectorise.Env
  10 import Vectorise.Vect
  11 import Vectorise.Monad
  12 import Vectorise.Builtins
  13 import Vectorise.Type.TyConDecl
  14 import Vectorise.Type.Classify
  15 import Vectorise.Type.PADict
  16 import Vectorise.Type.PData
  17 import Vectorise.Type.PRepr
  18 import Vectorise.Type.Repr
  19 import Vectorise.Utils
  20
  21 import HscTypes
  22 import CoreSyn
  23 import CoreUtils
  24 import CoreUnfold
  25 import DataCon
  26 import TyCon
  27 import Type
  28 import FamInstEnv
  29 import OccName
  30 import Id
  31 import MkId
  32 import Var
  33 import NameEnv
  34
  35 import Unique
  36 import UniqFM
  37 import Util
  38 import Outputable
  39 import FastString
  40 import MonadUtils
  41 import Control.Monad
  42 import Data.List
  43
  44 debug           = False
  45 dtrace s x      = if debug then pprTrace "VectType" s x else x
  46
  47
  48 -- | Vectorise a type environment.
  49 --   The type environment contains all the type things defined in a module.
  50 vectTypeEnv
  51         :: TypeEnv
  52         -> VM ( TypeEnv                 -- Vectorised type environment.
  53               , [FamInst]               -- New type family instances.
  54               , [(Var, CoreExpr)])      -- New top level bindings.
  55
  56 vectTypeEnv env
  57  = dtrace (ppr env)
  58  $ do
  59       cs <- readGEnv $ mk_map . global_tycons
  60
  61       -- Split the list of TyCons into the ones we have to vectorise vs the
  62       -- ones we can pass through unchanged. We also pass through algebraic
  63       -- types that use non Haskell98 features, as we don't handle those.
  64       let (conv_tcs, keep_tcs) = classifyTyCons cs groups
  65           keep_dcs             = concatMap tyConDataCons keep_tcs
  66
  67       zipWithM_ defTyCon   keep_tcs keep_tcs
  68       zipWithM_ defDataCon keep_dcs keep_dcs
  69
  70       new_tcs <- vectTyConDecls conv_tcs
  71
  72       let orig_tcs = keep_tcs ++ conv_tcs
  73
  74       -- We don't need to make new representation types for dictionary
  75       -- constructors. The constructors are always fully applied, and we don't
  76       -- need to lift them to arrays as a dictionary of a particular type
  77       -- always has the same value.
  78       let vect_tcs = filter (not . isClassTyCon)
  79                    $ keep_tcs ++ new_tcs
  80
  81       (_, binds, inst_tcs) <- fixV $ \ ~(dfuns', _, _) ->
  82         do
  83           defTyConPAs (zipLazy vect_tcs dfuns')
  84           reprs     <- mapM tyConRepr vect_tcs
  85           repr_tcs  <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
  86           pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
  87
  88           dfuns     <- sequence
  89                     $  zipWith5 buildTyConBindings
  90                                orig_tcs
  91                                vect_tcs
  92                                repr_tcs
  93                                pdata_tcs
  94                                reprs
  95
  96           binds     <- takeHoisted
  97           return (dfuns, binds, repr_tcs ++ pdata_tcs)
  98
  99       let all_new_tcs = new_tcs ++ inst_tcs
 100
 101       let new_env = extendTypeEnvList env
 102                        (map ATyCon all_new_tcs
 103                         ++ [ADataCon dc | tc <- all_new_tcs
 104                                         , dc <- tyConDataCons tc])
 105
 106       return (new_env, map mkLocalFamInst inst_tcs, binds)
 107   where
 108     tycons = typeEnvTyCons env
 109     groups = tyConGroups tycons
 110
 111     mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
 112
 113
 114
 115 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> SumRepr -> VM Var
 116 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc repr
 117  = do vectDataConWorkers orig_tc vect_tc pdata_tc
 118       buildPADict vect_tc prepr_tc pdata_tc repr
 119
 120
 121 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
 122 vectDataConWorkers orig_tc vect_tc arr_tc
 123  = do bs <- sequence
 124           . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
 125           $ zipWith4 mk_data_con (tyConDataCons vect_tc)
 126                                  rep_tys
 127                                  (inits rep_tys)
 128                                  (tail $ tails rep_tys)
 129       mapM_ (uncurry hoistBinding) bs
 130  where
 131     tyvars   = tyConTyVars vect_tc
 132     var_tys  = mkTyVarTys tyvars
 133     ty_args  = map Type var_tys
 134     res_ty   = mkTyConApp vect_tc var_tys
 135
 136     cons     = tyConDataCons vect_tc
 137     arity    = length cons
 138     [arr_dc] = tyConDataCons arr_tc
 139
 140     rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc
 141
 142
 143     mk_data_con con tys pre post
 144       = liftM2 (,) (vect_data_con con)
 145                    (lift_data_con tys pre post (mkDataConTag con))
 146
 147     sel_replicate len tag
 148       | arity > 1 = do
 149                       rep <- builtin (selReplicate arity)
 150                       return [rep `mkApps` [len, tag]]
 151
 152       | otherwise = return []
 153
 154     vect_data_con con = return $ mkConApp con ty_args
 155     lift_data_con tys pre_tys post_tys tag
 156       = do
 157           len  <- builtin liftingContext
 158           args <- mapM (newLocalVar (fsLit "xs"))
 159                   =<< mapM mkPDataType tys
 160
 161           sel  <- sel_replicate (Var len) tag
 162
 163           pre   <- mapM emptyPD (concat pre_tys)
 164           post  <- mapM emptyPD (concat post_tys)
 165
 166           return . mkLams (len : args)
 167                  . wrapFamInstBody arr_tc var_tys
 168                  . mkConApp arr_dc
 169                  $ ty_args ++ sel ++ pre ++ map Var args ++ post
 170
 171     def_worker data_con arg_tys mk_body
 172       = do
 173           arity <- polyArity tyvars
 174           body <- closedV
 175                 . inBind orig_worker
 176                 . polyAbstract tyvars $ \args ->
 177                   liftM (mkLams (tyvars ++ args) . vectorised)
 178                 $ buildClosures tyvars [] arg_tys res_ty mk_body
 179
 180           raw_worker <- cloneId mkVectOcc orig_worker (exprType body)
 181           let vect_worker = raw_worker `setIdUnfolding`
 182                               mkInlineRule body (Just arity)
 183           defGlobalVar orig_worker vect_worker
 184           return (vect_worker, body)
 185       where
 186         orig_worker = dataConWorkId data_con
 187