Env.hs 6.29 KB
Newer Older
Ian Lynagh's avatar
Ian Lynagh committed
1
{-# OPTIONS_GHC -XNoMonoLocalBinds #-}
2

3
module Vectorise.Type.Env ( 
4
	vectTypeEnv,
5 6
) where
  
7
import Vectorise.Env
8
import Vectorise.Vect
9 10
import Vectorise.Monad
import Vectorise.Builtins
11
import Vectorise.Type.TyConDecl
12
import Vectorise.Type.Classify
13
import Vectorise.Type.PADict
14 15 16
import Vectorise.Type.PData
import Vectorise.Type.PRepr
import Vectorise.Type.Repr
17
import Vectorise.Utils
18

19
import HscTypes
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
20
import CoreSyn
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
21
import CoreUtils
22
import CoreUnfold
23
import DataCon
24 25
import TyCon
import Type
26
import FamInstEnv
27
import OccName
28
import Id
29
import MkId
30
import NameEnv
31

32
import Unique
33
import UniqFM
34
import Util
35
import Outputable
36
import FastString
37 38
import MonadUtils
import Control.Monad
39 40
import Data.List

41

42 43
-- | Vectorise a type environment.
--   The type environment contains all the type things defined in a module.
44 45 46 47 48
--
vectTypeEnv :: TypeEnv
            -> VM ( TypeEnv             -- Vectorised type environment.
                  , [FamInst]           -- New type family instances.
                  , [(Var, CoreExpr)])  -- New top level bindings.
49
vectTypeEnv env
50 51 52
  = do
      traceVt "** vectTypeEnv" $ ppr env
      
53
      cs <- readGEnv $ mk_map . global_tycons
54 55 56 57

      -- Split the list of TyCons into the ones we have to vectorise vs the
      -- ones we can pass through unchanged. We also pass through algebraic 
      -- types that use non Haskell98 features, as we don't handle those.
58 59 60
      let tycons               = typeEnvTyCons env
          groups               = tyConGroups tycons

61
      let (conv_tcs, keep_tcs) = classifyTyCons cs groups
62
          orig_tcs             = keep_tcs ++ conv_tcs
63
          keep_dcs             = concatMap tyConDataCons keep_tcs
64

65
      -- Just use the unvectorised versions of these constructors in vectorised code.
66 67
      zipWithM_ defTyCon   keep_tcs keep_tcs
      zipWithM_ defDataCon keep_dcs keep_dcs
68

69 70
      -- Vectorise all the declarations.
      new_tcs      <- vectTyConDecls conv_tcs
71 72 73 74 75

      -- We don't need to make new representation types for dictionary
      -- constructors. The constructors are always fully applied, and we don't 
      -- need to lift them to arrays as a dictionary of a particular type
      -- always has the same value.
76 77
      let vect_tcs  = filter (not . isClassTyCon) 
                    $ keep_tcs ++ new_tcs
78

79 80 81 82 83 84 85
      reprs <- mapM tyConRepr vect_tcs
      repr_tcs  <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
      pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
      updGEnv $ extendFamEnv
              $ map mkLocalFamInst
              $ repr_tcs ++ pdata_tcs

86 87 88
      -- Create PRepr and PData instances for the vectorised types.
      -- We get back the binds for the instance functions, 
      -- and some new type constructors for the representation types.
89 90 91
      (_, binds, inst_tcs) <- fixV $ \ ~(dfuns', _, _) ->
        do
          defTyConPAs (zipLazy vect_tcs dfuns')
92 93 94 95 96 97 98 99 100 101 102
          reprs     <- mapM tyConRepr vect_tcs

          dfuns     <- sequence 
                    $  zipWith5 buildTyConBindings
                               orig_tcs
                               vect_tcs
                               repr_tcs
                               pdata_tcs
                               reprs

          binds     <- takeHoisted
103 104
          return (dfuns, binds, repr_tcs ++ pdata_tcs)

105 106
      -- The new type constructors are the vectorised versions of the originals, 
      -- plus the new type constructors that we use for the representations.
107
      let all_new_tcs = new_tcs ++ inst_tcs
108

109 110 111 112
      let new_env     =  extendTypeEnvList env
                      $  map ATyCon all_new_tcs
                      ++ [ADataCon dc | tc <- all_new_tcs
                                      , dc <- tyConDataCons tc]
113

114
      return (new_env, map mkLocalFamInst inst_tcs, binds)
115

116
   where
117 118
    mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]

119
buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> SumRepr -> VM Var
120
buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc repr
121
 = do vectDataConWorkers orig_tc vect_tc pdata_tc
122
      buildPADict vect_tc prepr_tc pdata_tc repr
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
123

124 125
vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
vectDataConWorkers orig_tc vect_tc arr_tc
126
 = do bs <- sequence
127 128 129
          . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
          $ zipWith4 mk_data_con (tyConDataCons vect_tc)
                                 rep_tys
130 131
                                 (inits rep_tys)
                                 (tail $ tails rep_tys)
132
      mapM_ (uncurry hoistBinding) bs
133
 where
134 135 136 137 138
    tyvars   = tyConTyVars vect_tc
    var_tys  = mkTyVarTys tyvars
    ty_args  = map Type var_tys
    res_ty   = mkTyConApp vect_tc var_tys

139 140 141 142
    cons     = tyConDataCons vect_tc
    arity    = length cons
    [arr_dc] = tyConDataCons arr_tc

143 144 145 146 147
    rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc


    mk_data_con con tys pre post
      = liftM2 (,) (vect_data_con con)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
148
                   (lift_data_con tys pre post (mkDataConTag con))
149

150 151 152 153 154 155 156
    sel_replicate len tag
      | arity > 1 = do
                      rep <- builtin (selReplicate arity)
                      return [rep `mkApps` [len, tag]]

      | otherwise = return []

157
    vect_data_con con = return $ mkConApp con ty_args
158
    lift_data_con tys pre_tys post_tys tag
159 160
      = do
          len  <- builtin liftingContext
Ian Lynagh's avatar
Ian Lynagh committed
161
          args <- mapM (newLocalVar (fsLit "xs"))
162
                  =<< mapM mkPDataType tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
163

164
          sel  <- sel_replicate (Var len) tag
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
165

166 167
          pre   <- mapM emptyPD (concat pre_tys)
          post  <- mapM emptyPD (concat post_tys)
168 169 170 171

          return . mkLams (len : args)
                 . wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
172
                 $ ty_args ++ sel ++ pre ++ map Var args ++ post
173 174 175

    def_worker data_con arg_tys mk_body
      = do
176
          arity <- polyArity tyvars
177 178
          body <- closedV
                . inBind orig_worker
179 180
                . polyAbstract tyvars $ \args ->
                  liftM (mkLams (tyvars ++ args) . vectorised)
181 182
                $ buildClosures tyvars [] arg_tys res_ty mk_body

183 184
          raw_worker <- cloneId mkVectOcc orig_worker (exprType body)
          let vect_worker = raw_worker `setIdUnfolding`
185
                              mkInlineUnfolding (Just arity) body
186 187 188 189 190
          defGlobalVar orig_worker vect_worker
          return (vect_worker, body)
      where
        orig_worker = dataConWorkId data_con