1 {-# OPTIONS_GHC -fno-warn-missing-signatures #-}
2 {-# OPTIONS_GHC -XNoMonoLocalBinds #-}
3 -- Roman likes local bindings
4 -- If this module lives on I'd like to get rid of this flag in due course
6 module Vectorise.Type.Env (
12 import Vectorise.Monad
13 import Vectorise.Builtins
14 import Vectorise.Type.TyConDecl
15 import Vectorise.Type.Classify
16 import Vectorise.Type.PADict
17 import Vectorise.Type.PData
18 import Vectorise.Type.PRepr
19 import Vectorise.Type.Repr
20 import Vectorise.Utils
46 dtrace s x = if debug then pprTrace "VectType" s x else x
48 -- | Vectorise a type environment.
49 -- The type environment contains all the type things defined in a module.
52 -> VM ( TypeEnv -- Vectorised type environment.
53 , [FamInst] -- New type family instances.
54 , [(Var, CoreExpr)]) -- New top level bindings.
59 cs <- readGEnv $ mk_map . global_tycons
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 tycons = typeEnvTyCons env
65 groups = tyConGroups tycons
67 let (conv_tcs, keep_tcs) = classifyTyCons cs groups
68 orig_tcs = keep_tcs ++ conv_tcs
69 keep_dcs = concatMap tyConDataCons keep_tcs
71 -- Just use the unvectorised versions of these constructors in vectorised code.
72 zipWithM_ defTyCon keep_tcs keep_tcs
73 zipWithM_ defDataCon keep_dcs keep_dcs
75 -- Vectorise all the declarations.
76 new_tcs <- vectTyConDecls conv_tcs
78 -- We don't need to make new representation types for dictionary
79 -- constructors. The constructors are always fully applied, and we don't
80 -- need to lift them to arrays as a dictionary of a particular type
81 -- always has the same value.
82 let vect_tcs = filter (not . isClassTyCon)
85 reprs <- mapM tyConRepr vect_tcs
86 repr_tcs <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
87 pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
88 updGEnv $ extendFamEnv
90 $ repr_tcs ++ pdata_tcs
92 -- Create PRepr and PData instances for the vectorised types.
93 -- We get back the binds for the instance functions,
94 -- and some new type constructors for the representation types.
95 (_, binds, inst_tcs) <- fixV $ \ ~(dfuns', _, _) ->
97 defTyConPAs (zipLazy vect_tcs dfuns')
98 reprs <- mapM tyConRepr vect_tcs
101 $ zipWith5 buildTyConBindings
109 return (dfuns, binds, repr_tcs ++ pdata_tcs)
111 -- The new type constructors are the vectorised versions of the originals,
112 -- plus the new type constructors that we use for the representations.
113 let all_new_tcs = new_tcs ++ inst_tcs
115 let new_env = extendTypeEnvList env
116 $ map ATyCon all_new_tcs
117 ++ [ADataCon dc | tc <- all_new_tcs
118 , dc <- tyConDataCons tc]
120 return (new_env, map mkLocalFamInst inst_tcs, binds)
123 mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
127 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> SumRepr -> VM Var
128 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc repr
129 = do vectDataConWorkers orig_tc vect_tc pdata_tc
130 buildPADict vect_tc prepr_tc pdata_tc repr
133 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
134 vectDataConWorkers orig_tc vect_tc arr_tc
136 . zipWith3 def_worker (tyConDataCons orig_tc) rep_tys
137 $ zipWith4 mk_data_con (tyConDataCons vect_tc)
140 (tail $ tails rep_tys)
141 mapM_ (uncurry hoistBinding) bs
143 tyvars = tyConTyVars vect_tc
144 var_tys = mkTyVarTys tyvars
145 ty_args = map Type var_tys
146 res_ty = mkTyConApp vect_tc var_tys
148 cons = tyConDataCons vect_tc
150 [arr_dc] = tyConDataCons arr_tc
152 rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
155 mk_data_con con tys pre post
156 = liftM2 (,) (vect_data_con con)
157 (lift_data_con tys pre post (mkDataConTag con))
159 sel_replicate len tag
161 rep <- builtin (selReplicate arity)
162 return [rep `mkApps` [len, tag]]
164 | otherwise = return []
166 vect_data_con con = return $ mkConApp con ty_args
167 lift_data_con tys pre_tys post_tys tag
169 len <- builtin liftingContext
170 args <- mapM (newLocalVar (fsLit "xs"))
171 =<< mapM mkPDataType tys
173 sel <- sel_replicate (Var len) tag
175 pre <- mapM emptyPD (concat pre_tys)
176 post <- mapM emptyPD (concat post_tys)
178 return . mkLams (len : args)
179 . wrapFamInstBody arr_tc var_tys
181 $ ty_args ++ sel ++ pre ++ map Var args ++ post
183 def_worker data_con arg_tys mk_body
185 arity <- polyArity tyvars
188 . polyAbstract tyvars $ \args ->
189 liftM (mkLams (tyvars ++ args) . vectorised)
190 $ buildClosures tyvars [] arg_tys res_ty mk_body
192 raw_worker <- cloneId mkVectOcc orig_worker (exprType body)
193 let vect_worker = raw_worker `setIdUnfolding`
194 mkInlineUnfolding (Just arity) body
195 defGlobalVar orig_worker vect_worker
196 return (vect_worker, body)
198 orig_worker = dataConWorkId data_con