Remove dead code
[ghc-hetmet.git] / compiler / vectorise / VectType.hs
1 module VectType ( vectTyCon, vectType, vectTypeEnv,
2                    PAInstance, buildPADict )
3 where
4
5 #include "HsVersions.h"
6
7 import VectMonad
8 import VectUtils
9 import VectCore
10
11 import HscTypes          ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
12 import CoreSyn
13 import CoreUtils
14 import BuildTyCl
15 import DataCon
16 import TyCon
17 import Type
18 import TypeRep
19 import Coercion
20 import FamInstEnv        ( FamInst, mkLocalFamInst )
21 import InstEnv           ( Instance, mkLocalInstance, instanceDFunId )
22 import OccName
23 import MkId
24 import BasicTypes        ( StrictnessMark(..), OverlapFlag(..), boolToRecFlag )
25 import Var               ( Var )
26 import Id                ( mkWildId )
27 import Name              ( Name, getOccName )
28 import NameEnv
29 import TysWiredIn        ( unitTy, intTy, intDataCon )
30 import TysPrim           ( intPrimTy )
31
32 import Unique
33 import UniqFM
34 import UniqSet
35 import Digraph           ( SCC(..), stronglyConnComp )
36
37 import Outputable
38
39 import Control.Monad  ( liftM, liftM2, zipWithM, zipWithM_ )
40 import Data.List      ( inits, tails, zipWith4, zipWith5 )
41
42 -- ----------------------------------------------------------------------------
43 -- Types
44
45 vectTyCon :: TyCon -> VM TyCon
46 vectTyCon tc
47   | isFunTyCon tc        = builtin closureTyCon
48   | isBoxedTupleTyCon tc = return tc
49   | isUnLiftedTyCon tc   = return tc
50   | otherwise = do
51                   r <- lookupTyCon tc
52                   case r of
53                     Just tc' -> return tc'
54
55                     -- FIXME: just for now
56                     Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc
57
58 vectType :: Type -> VM Type
59 vectType ty | Just ty' <- coreView ty = vectType ty'
60 vectType (TyVarTy tv) = return $ TyVarTy tv
61 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
62 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
63 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
64                                              (mapM vectType [ty1,ty2])
65 vectType ty@(ForAllTy _ _)
66   = do
67       mdicts   <- mapM paDictArgType tyvars
68       mono_ty' <- vectType mono_ty
69       return $ tyvars `mkForAllTys` ([dict | Just dict <- mdicts] `mkFunTys` mono_ty')
70   where
71     (tyvars, mono_ty) = splitForAllTys ty
72
73 vectType ty = pprPanic "vectType:" (ppr ty)
74
75 -- ----------------------------------------------------------------------------
76 -- Type definitions
77
78 type TyConGroup = ([TyCon], UniqSet TyCon)
79
80 data PAInstance = PAInstance {
81                     painstDFun      :: Var
82                   , painstOrigTyCon :: TyCon
83                   , painstVectTyCon :: TyCon
84                   , painstArrTyCon  :: TyCon
85                   }
86
87 vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [(Var, CoreExpr)])
88 vectTypeEnv env
89   = do
90       cs <- readGEnv $ mk_map . global_tycons
91       let (conv_tcs, keep_tcs) = classifyTyCons cs groups
92           keep_dcs             = concatMap tyConDataCons keep_tcs
93       zipWithM_ defTyCon   keep_tcs keep_tcs
94       zipWithM_ defDataCon keep_dcs keep_dcs
95       new_tcs <- vectTyConDecls conv_tcs
96
97       let orig_tcs = keep_tcs ++ conv_tcs
98           vect_tcs  = keep_tcs ++ new_tcs
99
100       repr_tcs <- zipWithM buildPReprTyCon   orig_tcs vect_tcs
101       parr_tcs <- zipWithM buildPArrayTyCon orig_tcs vect_tcs
102       dfuns    <- mapM mkPADFun vect_tcs
103       defTyConPAs (zip vect_tcs dfuns)
104       binds    <- sequence (zipWith5 buildTyConBindings orig_tcs
105                                                         vect_tcs
106                                                         repr_tcs
107                                                         parr_tcs
108                                                         dfuns)
109
110       let all_new_tcs = new_tcs ++ repr_tcs ++ parr_tcs
111
112       let new_env = extendTypeEnvList env
113                        (map ATyCon all_new_tcs
114                         ++ [ADataCon dc | tc <- all_new_tcs
115                                         , dc <- tyConDataCons tc])
116
117       return (new_env, map mkLocalFamInst (repr_tcs ++ parr_tcs), concat binds)
118   where
119     tycons = typeEnvTyCons env
120     groups = tyConGroups tycons
121
122     mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
123
124     keep_tc tc = let dcs = tyConDataCons tc
125                  in
126                  defTyCon tc tc >> zipWithM_ defDataCon dcs dcs
127
128
129 vectTyConDecls :: [TyCon] -> VM [TyCon]
130 vectTyConDecls tcs = fixV $ \tcs' ->
131   do
132     mapM_ (uncurry defTyCon) (lazy_zip tcs tcs')
133     mapM vectTyConDecl tcs
134   where
135     lazy_zip [] _ = []
136     lazy_zip (x:xs) ~(y:ys) = (x,y) : lazy_zip xs ys
137
138 vectTyConDecl :: TyCon -> VM TyCon
139 vectTyConDecl tc
140   = do
141       name' <- cloneName mkVectTyConOcc name
142       rhs'  <- vectAlgTyConRhs (algTyConRhs tc)
143
144       liftDs $ buildAlgTyCon name'
145                              tyvars
146                              []           -- no stupid theta
147                              rhs'
148                              rec_flag     -- FIXME: is this ok?
149                              False        -- FIXME: no generics
150                              False        -- not GADT syntax
151                              Nothing      -- not a family instance
152   where
153     name   = tyConName tc
154     tyvars = tyConTyVars tc
155     rec_flag = boolToRecFlag (isRecursiveTyCon tc)
156
157 vectAlgTyConRhs :: AlgTyConRhs -> VM AlgTyConRhs
158 vectAlgTyConRhs (DataTyCon { data_cons = data_cons
159                            , is_enum   = is_enum
160                            })
161   = do
162       data_cons' <- mapM vectDataCon data_cons
163       zipWithM_ defDataCon data_cons data_cons'
164       return $ DataTyCon { data_cons = data_cons'
165                          , is_enum   = is_enum
166                          }
167
168 vectDataCon :: DataCon -> VM DataCon
169 vectDataCon dc
170   | not . null $ dataConExTyVars dc = pprPanic "vectDataCon: existentials" (ppr dc)
171   | not . null $ dataConEqSpec   dc = pprPanic "vectDataCon: eq spec" (ppr dc)
172   | otherwise
173   = do
174       name'    <- cloneName mkVectDataConOcc name
175       tycon'   <- vectTyCon tycon
176       arg_tys  <- mapM vectType rep_arg_tys
177
178       liftDs $ buildDataCon name'
179                             False           -- not infix
180                             (map (const NotMarkedStrict) arg_tys)
181                             []              -- no labelled fields
182                             univ_tvs
183                             []              -- no existential tvs for now
184                             []              -- no eq spec for now
185                             []              -- no context
186                             arg_tys
187                             tycon'
188   where
189     name        = dataConName dc
190     univ_tvs    = dataConUnivTyVars dc
191     rep_arg_tys = dataConRepArgTys dc
192     tycon       = dataConTyCon dc
193
194 mk_fam_inst :: TyCon -> TyCon -> (TyCon, [Type])
195 mk_fam_inst fam_tc arg_tc
196   = (fam_tc, [mkTyConApp arg_tc . mkTyVarTys $ tyConTyVars arg_tc])
197
198 buildPReprTyCon :: TyCon -> TyCon -> VM TyCon
199 buildPReprTyCon orig_tc vect_tc
200   = do
201       name     <- cloneName mkPReprTyConOcc (tyConName orig_tc)
202       rhs_ty   <- buildPReprType vect_tc
203       prepr_tc <- builtin preprTyCon
204       liftDs $ buildSynTyCon name
205                              tyvars
206                              (SynonymTyCon rhs_ty)
207                              (Just $ mk_fam_inst prepr_tc vect_tc)
208   where
209     tyvars = tyConTyVars vect_tc
210
211 buildPReprType :: TyCon -> VM Type
212 buildPReprType = mkPRepr . map dataConRepArgTys . tyConDataCons
213
214 buildToPRepr :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
215 buildToPRepr _ vect_tc prepr_tc _
216   = do
217       arg <- newLocalVar FSLIT("x") arg_ty
218       bndrss <- mapM (mapM (newLocalVar FSLIT("x"))) rep_tys
219       (alt_bodies, res_ty) <- mkToPRepr $ map (map Var) bndrss
220
221       return . Lam arg
222              . wrapFamInstBody prepr_tc var_tys
223              . Case (Var arg) (mkWildId arg_ty) res_ty
224              $ zipWith3 mk_alt data_cons bndrss alt_bodies
225   where
226     var_tys   = mkTyVarTys $ tyConTyVars vect_tc
227     arg_ty    = mkTyConApp vect_tc var_tys
228     data_cons = tyConDataCons vect_tc
229     rep_tys   = map dataConRepArgTys data_cons
230
231     mk_alt data_con bndrs body = (DataAlt data_con, bndrs, body)
232
233 buildToArrPRepr :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
234 buildToArrPRepr _ vect_tc prepr_tc arr_tc
235   = do
236       arg_ty  <- mkPArrayType el_ty
237       rep_tys <- mapM (mapM mkPArrayType) rep_el_tys
238
239       arg     <- newLocalVar FSLIT("xs") arg_ty
240       bndrss  <- mapM (mapM (newLocalVar FSLIT("ys"))) rep_tys
241       len     <- newLocalVar FSLIT("len") intPrimTy
242       sel     <- newLocalVar FSLIT("sel") =<< mkPArrayType intTy
243
244       let add_sel xs | has_selector = sel : xs
245                      | otherwise    = xs
246
247           all_bndrs = len : add_sel (concat bndrss)
248
249       res      <- parrayCoerce prepr_tc var_tys
250                 =<< mkToArrPRepr (Var len) (Var sel) (map (map Var) bndrss)
251       res_ty   <- mkPArrayType =<< mkPReprType el_ty
252
253       return . Lam arg
254              $ Case (unwrapFamInstScrut arr_tc var_tys (Var arg))
255                     (mkWildId (mkTyConApp arr_tc var_tys))
256                     res_ty
257                     [(DataAlt arr_dc, all_bndrs, res)]
258   where
259     var_tys    = mkTyVarTys $ tyConTyVars vect_tc
260     el_ty      = mkTyConApp vect_tc var_tys
261     data_cons  = tyConDataCons vect_tc
262     rep_el_tys = map dataConRepArgTys data_cons
263
264     [arr_dc]   = tyConDataCons arr_tc
265
266     has_selector | [_] <- data_cons = False
267                  | otherwise        = True
268
269
270 buildFromPRepr :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
271 buildFromPRepr _ vect_tc prepr_tc _
272   = do
273       arg_ty <- mkPReprType res_ty
274       arg <- newLocalVar FSLIT("x") arg_ty
275       alts <- mapM mk_alt data_cons
276       body <- mkFromPRepr (unwrapFamInstScrut prepr_tc var_tys (Var arg))
277                           res_ty alts
278       return $ Lam arg body
279   where
280     var_tys   = mkTyVarTys $ tyConTyVars vect_tc
281     res_ty    = mkTyConApp vect_tc var_tys
282     data_cons = tyConDataCons vect_tc
283
284     mk_alt dc = do
285                   bndrs <- mapM (newLocalVar FSLIT("x")) $ dataConRepArgTys dc
286                   return (bndrs, mkConApp dc (map Type var_tys ++ map Var bndrs))
287
288 buildFromArrPRepr :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
289 buildFromArrPRepr _ vect_tc prepr_tc arr_tc
290   = mkFromArrPRepr undefined undefined undefined undefined undefined undefined
291
292 buildPRDict :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
293 buildPRDict _ vect_tc prepr_tc _
294   = prCoerce prepr_tc var_tys
295   =<< prDictOfType (mkTyConApp prepr_tc var_tys)
296   where
297     var_tys = mkTyVarTys $ tyConTyVars vect_tc
298
299 buildPArrayTyCon :: TyCon -> TyCon -> VM TyCon
300 buildPArrayTyCon orig_tc vect_tc = fixV $ \repr_tc ->
301   do
302     name'  <- cloneName mkPArrayTyConOcc orig_name
303     rhs    <- buildPArrayTyConRhs orig_name vect_tc repr_tc
304     parray <- builtin parrayTyCon
305
306     liftDs $ buildAlgTyCon name'
307                            tyvars
308                            []          -- no stupid theta
309                            rhs
310                            rec_flag    -- FIXME: is this ok?
311                            False       -- FIXME: no generics
312                            False       -- not GADT syntax
313                            (Just $ mk_fam_inst parray vect_tc)
314   where
315     orig_name = tyConName orig_tc
316     tyvars = tyConTyVars vect_tc
317     rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
318     
319
320 buildPArrayTyConRhs :: Name -> TyCon -> TyCon -> VM AlgTyConRhs
321 buildPArrayTyConRhs orig_name vect_tc repr_tc
322   = do
323       data_con <- buildPArrayDataCon orig_name vect_tc repr_tc
324       return $ DataTyCon { data_cons = [data_con], is_enum = False }
325
326 buildPArrayDataCon :: Name -> TyCon -> TyCon -> VM DataCon
327 buildPArrayDataCon orig_name vect_tc repr_tc
328   = do
329       dc_name  <- cloneName mkPArrayDataConOcc orig_name
330       shape    <- tyConShape vect_tc
331       repr_tys <- mapM mkPArrayType types
332
333       liftDs $ buildDataCon dc_name
334                             False                  -- not infix
335                             (shapeStrictness shape ++ map (const NotMarkedStrict) repr_tys)
336                             []                     -- no field labels
337                             (tyConTyVars vect_tc)
338                             []                     -- no existentials
339                             []                     -- no eq spec
340                             []                     -- no context
341                             (shapeReprTys shape ++ repr_tys)
342                             repr_tc
343   where
344     types = [ty | dc <- tyConDataCons vect_tc
345                 , ty <- dataConRepArgTys dc]
346
347 mkPADFun :: TyCon -> VM Var
348 mkPADFun vect_tc
349   = newExportedVar (mkPADFunOcc $ getOccName vect_tc) =<< paDFunType vect_tc
350
351 data Shape = Shape {
352                shapeReprTys    :: [Type]
353              , shapeStrictness :: [StrictnessMark]
354              , shapeLength     :: [CoreExpr] -> VM CoreExpr
355              , shapeReplicate  :: CoreExpr -> CoreExpr -> VM [CoreExpr]
356              }
357
358 tyConShape :: TyCon -> VM Shape
359 tyConShape vect_tc
360   | isProductTyCon vect_tc
361   = return $ Shape {
362                 shapeReprTys    = [intPrimTy]
363               , shapeStrictness = [NotMarkedStrict]
364               , shapeLength     = \[len] -> return len
365               , shapeReplicate  = \len _ -> return [len]
366               }
367
368   | otherwise
369   = do
370       repr_ty <- mkPArrayType intTy   -- FIXME: we want to unbox this
371       return $ Shape {
372                  shapeReprTys    = [repr_ty]
373                , shapeStrictness = [MarkedStrict]
374                , shapeLength     = \[sel] -> lengthPA sel
375                , shapeReplicate  = \len n -> do
376                                                e <- replicatePA len n
377                                                return [e]
378                }
379
380 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> Var
381                    -> VM [(Var, CoreExpr)]
382 buildTyConBindings orig_tc vect_tc prepr_tc arr_tc dfun
383   = do
384       shape <- tyConShape vect_tc
385       sequence_ (zipWith4 (vectDataConWorker shape vect_tc arr_tc arr_dc)
386                           orig_dcs
387                           vect_dcs
388                           (inits repr_tys)
389                           (tails repr_tys))
390       dict <- buildPADict shape vect_tc prepr_tc arr_tc dfun
391       binds <- takeHoisted
392       return $ (dfun, dict) : binds
393   where
394     orig_dcs = tyConDataCons orig_tc
395     vect_dcs = tyConDataCons vect_tc
396     [arr_dc] = tyConDataCons arr_tc
397
398     repr_tys = map dataConRepArgTys vect_dcs
399
400 vectDataConWorker :: Shape -> TyCon -> TyCon -> DataCon
401                   -> DataCon -> DataCon -> [[Type]] -> [[Type]]
402                   -> VM ()
403 vectDataConWorker shape vect_tc arr_tc arr_dc orig_dc vect_dc pre (dc_tys : post)
404   = do
405       clo <- closedV
406            . inBind orig_worker
407            . polyAbstract tvs $ \abstract ->
408              liftM (abstract . vectorised)
409            $ buildClosures tvs [] dc_tys res_ty (liftM2 (,) mk_vect mk_lift)
410
411       worker <- cloneId mkVectOcc orig_worker (exprType clo)
412       hoistBinding worker clo
413       defGlobalVar orig_worker worker
414       return ()
415   where
416     tvs     = tyConTyVars vect_tc
417     arg_tys = mkTyVarTys tvs
418     res_ty  = mkTyConApp vect_tc arg_tys
419
420     orig_worker = dataConWorkId orig_dc
421
422     mk_vect = return . mkConApp vect_dc $ map Type arg_tys
423     mk_lift = do
424                 len     <- newLocalVar FSLIT("n") intPrimTy
425                 arr_tys <- mapM mkPArrayType dc_tys
426                 args    <- mapM (newLocalVar FSLIT("xs")) arr_tys
427                 shapes  <- shapeReplicate shape
428                                           (Var len)
429                                           (mkDataConTag vect_dc)
430                 
431                 empty_pre  <- mapM emptyPA (concat pre)
432                 empty_post <- mapM emptyPA (concat post)
433
434                 return . mkLams (len : args)
435                        . wrapFamInstBody arr_tc arg_tys
436                        . mkConApp arr_dc
437                        $ map Type arg_tys ++ shapes
438                                           ++ empty_pre
439                                           ++ map Var args
440                                           ++ empty_post
441
442 buildPADict :: Shape -> TyCon -> TyCon -> TyCon -> Var -> VM CoreExpr
443 buildPADict shape vect_tc prepr_tc arr_tc dfun
444   = polyAbstract tvs $ \abstract ->
445     do
446       meth_binds <- mapM (mk_method shape) paMethods
447       let meth_exprs = map (Var . fst) meth_binds
448
449       pa_dc <- builtin paDataCon
450       let dict = mkConApp pa_dc (Type (mkTyConApp vect_tc arg_tys) : meth_exprs)
451           body = Let (Rec meth_binds) dict
452       return . mkInlineMe $ abstract body
453   where
454     tvs = tyConTyVars arr_tc
455     arg_tys = mkTyVarTys tvs
456
457     mk_method shape (name, build)
458       = localV
459       $ do
460           body <- build shape vect_tc prepr_tc arr_tc
461           var  <- newLocalVar name (exprType body)
462           return (var, mkInlineMe body)
463           
464 paMethods = [(FSLIT("toPRepr"),      buildToPRepr),
465              (FSLIT("fromPRepr"),    buildFromPRepr),
466              (FSLIT("toArrPRepr"),   buildToArrPRepr),
467              (FSLIT("fromArrPRepr"), buildFromArrPRepr),
468              (FSLIT("dictPRepr"),    buildPRDict)]
469
470 -- | Split the given tycons into two sets depending on whether they have to be
471 -- converted (first list) or not (second list). The first argument contains
472 -- information about the conversion status of external tycons:
473 -- 
474 --   * tycons which have converted versions are mapped to True
475 --   * tycons which are not changed by vectorisation are mapped to False
476 --   * tycons which can't be converted are not elements of the map
477 --
478 classifyTyCons :: UniqFM Bool -> [TyConGroup] -> ([TyCon], [TyCon])
479 classifyTyCons = classify [] []
480   where
481     classify conv keep cs [] = (conv, keep)
482     classify conv keep cs ((tcs, ds) : rs)
483       | can_convert && must_convert
484         = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
485       | can_convert
486         = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
487       | otherwise
488         = classify conv keep cs rs
489       where
490         refs = ds `delListFromUniqSet` tcs
491
492         can_convert  = isNullUFM (refs `minusUFM` cs) && all convertable tcs
493         must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
494
495         convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
496     
497 -- | Compute mutually recursive groups of tycons in topological order
498 --
499 tyConGroups :: [TyCon] -> [TyConGroup]
500 tyConGroups tcs = map mk_grp (stronglyConnComp edges)
501   where
502     edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
503                                 , let ds = tyConsOfTyCon tc]
504
505     mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
506     mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
507       where
508         (tcs, dss) = unzip els
509
510 tyConsOfTyCon :: TyCon -> UniqSet TyCon
511 tyConsOfTyCon 
512   = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
513
514 tyConsOfType :: Type -> UniqSet TyCon
515 tyConsOfType ty
516   | Just ty' <- coreView ty    = tyConsOfType ty'
517 tyConsOfType (TyVarTy v)       = emptyUniqSet
518 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
519   where
520     extend | isUnLiftedTyCon tc
521            || isTupleTyCon   tc = id
522
523            | otherwise          = (`addOneToUniqSet` tc)
524
525 tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
526 tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
527                                  `addOneToUniqSet` funTyCon
528 tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty
529 tyConsOfType other             = pprPanic "ClosureConv.tyConsOfType" $ ppr other
530
531 tyConsOfTypes :: [Type] -> UniqSet TyCon
532 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
533