Env.hs 5.75 KB
Newer Older
1 2 3
{-# OPTIONS_GHC -XNoMonoLocalBinds -fno-warn-missing-signatures #-}
-- Roman likes local bindings
-- If this module lives on I'd like to get rid of this flag in due course
4

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

21
import HscTypes
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
22
import CoreSyn
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
23
import CoreUtils
24
import CoreUnfold
25
import DataCon
26 27
import TyCon
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 41
import MonadUtils
import Control.Monad
42 43 44 45
import Data.List

debug		= False
dtrace s x	= if debug then pprTrace "VectType" s x else x
46

47

48 49
-- | Vectorise a type environment.
--   The type environment contains all the type things defined in a module.
50 51 52 53 54 55
vectTypeEnv 
	:: TypeEnv
	-> VM ( TypeEnv			-- Vectorised type environment.
	      , [FamInst]		-- New type family instances.
	      , [(Var, CoreExpr)])	-- New top level bindings.
	
56
vectTypeEnv env
57 58
 = dtrace (ppr env)
 $ do
59
      cs <- readGEnv $ mk_map . global_tycons
60 61 62 63

      -- 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.
64 65
      let (conv_tcs, keep_tcs) = classifyTyCons cs groups
          keep_dcs             = concatMap tyConDataCons keep_tcs
66

67 68
      zipWithM_ defTyCon   keep_tcs keep_tcs
      zipWithM_ defDataCon keep_dcs keep_dcs
69

70 71 72
      new_tcs <- vectTyConDecls conv_tcs

      let orig_tcs = keep_tcs ++ conv_tcs
73 74 75 76 77 78 79 80

      -- 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.
      let vect_tcs = filter (not . isClassTyCon) 
                   $ keep_tcs ++ new_tcs

81 82 83
      (_, binds, inst_tcs) <- fixV $ \ ~(dfuns', _, _) ->
        do
          defTyConPAs (zipLazy vect_tcs dfuns')
84
          reprs     <- mapM tyConRepr vect_tcs
85 86
          repr_tcs  <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
          pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
87 88 89 90 91 92 93 94 95 96

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

          binds     <- takeHoisted
97 98 99
          return (dfuns, binds, repr_tcs ++ pdata_tcs)

      let all_new_tcs = new_tcs ++ inst_tcs
100 101

      let new_env = extendTypeEnvList env
102 103
                       (map ATyCon all_new_tcs
                        ++ [ADataCon dc | tc <- all_new_tcs
104 105
                                        , dc <- tyConDataCons tc])

106
      return (new_env, map mkLocalFamInst inst_tcs, binds)
107 108 109 110 111 112 113
  where
    tycons = typeEnvTyCons env
    groups = tyConGroups tycons

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


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
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
119

120

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

136 137 138 139
    cons     = tyConDataCons vect_tc
    arity    = length cons
    [arr_dc] = tyConDataCons arr_tc

140 141 142 143 144
    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
145
                   (lift_data_con tys pre post (mkDataConTag con))
146

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

      | otherwise = return []

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

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

163 164
          pre   <- mapM emptyPD (concat pre_tys)
          post  <- mapM emptyPD (concat post_tys)
165 166 167 168

          return . mkLams (len : args)
                 . wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
169
                 $ ty_args ++ sel ++ pre ++ map Var args ++ post
170 171 172

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

180 181
          raw_worker <- cloneId mkVectOcc orig_worker (exprType body)
          let vect_worker = raw_worker `setIdUnfolding`
182
                              mkInlineRule body (Just arity)
183 184 185 186 187
          defGlobalVar orig_worker vect_worker
          return (vect_worker, body)
      where
        orig_worker = dataConWorkId data_con