VectType.hs 20 KB
Newer Older
1
module VectType ( vectTyCon, vectType, vectTypeEnv,
2
                   PAInstance, buildPADict )
3
4
5
6
7
8
where

#include "HsVersions.h"

import VectMonad
import VectUtils
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
9
import VectCore
10

11
import HscTypes          ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
12
import CoreSyn
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
13
import CoreUtils
14
import BuildTyCl
15
import DataCon
16
17
18
import TyCon
import Type
import TypeRep
19
import Coercion
20
import FamInstEnv        ( FamInst, mkLocalFamInst )
21
import InstEnv           ( Instance, mkLocalInstance, instanceDFunId )
22
23
import OccName
import MkId
24
import BasicTypes        ( StrictnessMark(..), OverlapFlag(..), boolToRecFlag )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
25
import Var               ( Var )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
26
import Id                ( mkWildId )
27
import Name              ( Name, getOccName )
28
import NameEnv
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
29
import TysWiredIn        ( unitTy, intTy, intDataCon, unitDataConId )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
30
import TysPrim           ( intPrimTy )
31

32
import Unique
33
34
35
36
import UniqFM
import UniqSet
import Digraph           ( SCC(..), stronglyConnComp )

37
38
import Outputable

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
39
import Control.Monad  ( liftM, liftM2, zipWithM, zipWithM_ )
40
import Data.List      ( inits, tails, zipWith4, zipWith5 )
41

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
42
43
44
-- ----------------------------------------------------------------------------
-- Types

45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
vectTyCon :: TyCon -> VM TyCon
vectTyCon tc
  | isFunTyCon tc        = builtin closureTyCon
  | isBoxedTupleTyCon tc = return tc
  | isUnLiftedTyCon tc   = return tc
  | otherwise = do
                  r <- lookupTyCon tc
                  case r of
                    Just tc' -> return tc'

                    -- FIXME: just for now
                    Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc

vectType :: Type -> VM Type
vectType ty | Just ty' <- coreView ty = vectType ty'
vectType (TyVarTy tv) = return $ TyVarTy tv
vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
                                             (mapM vectType [ty1,ty2])
vectType ty@(ForAllTy _ _)
  = do
      mdicts   <- mapM paDictArgType tyvars
      mono_ty' <- vectType mono_ty
      return $ tyvars `mkForAllTys` ([dict | Just dict <- mdicts] `mkFunTys` mono_ty')
  where
    (tyvars, mono_ty) = splitForAllTys ty

vectType ty = pprPanic "vectType:" (ppr ty)

75
76
77
78
79
-- ----------------------------------------------------------------------------
-- Type definitions

type TyConGroup = ([TyCon], UniqSet TyCon)

80
data PAInstance = PAInstance {
81
                    painstDFun      :: Var
82
                  , painstOrigTyCon :: TyCon
83
84
85
86
                  , painstVectTyCon :: TyCon
                  , painstArrTyCon  :: TyCon
                  }

87
vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [(Var, CoreExpr)])
88
89
90
91
92
93
94
vectTypeEnv env
  = do
      cs <- readGEnv $ mk_map . global_tycons
      let (conv_tcs, keep_tcs) = classifyTyCons cs groups
          keep_dcs             = concatMap tyConDataCons keep_tcs
      zipWithM_ defTyCon   keep_tcs keep_tcs
      zipWithM_ defDataCon keep_dcs keep_dcs
95
96
97
98
99
      new_tcs <- vectTyConDecls conv_tcs

      let orig_tcs = keep_tcs ++ conv_tcs
          vect_tcs  = keep_tcs ++ new_tcs

100
      repr_tcs <- zipWithM buildPReprTyCon   orig_tcs vect_tcs
101
      parr_tcs <- zipWithM buildPArrayTyCon orig_tcs vect_tcs
102
103
      dfuns    <- mapM mkPADFun vect_tcs
      defTyConPAs (zip vect_tcs dfuns)
104
105
106
107
108
109
      binds    <- sequence (zipWith5 buildTyConBindings orig_tcs
                                                        vect_tcs
                                                        repr_tcs
                                                        parr_tcs
                                                        dfuns)

110
      let all_new_tcs = new_tcs ++ repr_tcs ++ parr_tcs
111
112

      let new_env = extendTypeEnvList env
113
114
                       (map ATyCon all_new_tcs
                        ++ [ADataCon dc | tc <- all_new_tcs
115
116
                                        , dc <- tyConDataCons tc])

117
      return (new_env, map mkLocalFamInst (repr_tcs ++ parr_tcs), concat binds)
118
119
120
121
122
123
124
125
126
127
128
  where
    tycons = typeEnvTyCons env
    groups = tyConGroups tycons

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

    keep_tc tc = let dcs = tyConDataCons tc
                 in
                 defTyCon tc tc >> zipWithM_ defDataCon dcs dcs


129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
vectTyConDecls :: [TyCon] -> VM [TyCon]
vectTyConDecls tcs = fixV $ \tcs' ->
  do
    mapM_ (uncurry defTyCon) (lazy_zip tcs tcs')
    mapM vectTyConDecl tcs
  where
    lazy_zip [] _ = []
    lazy_zip (x:xs) ~(y:ys) = (x,y) : lazy_zip xs ys

vectTyConDecl :: TyCon -> VM TyCon
vectTyConDecl tc
  = do
      name' <- cloneName mkVectTyConOcc name
      rhs'  <- vectAlgTyConRhs (algTyConRhs tc)

144
145
146
147
148
149
150
151
      liftDs $ buildAlgTyCon name'
                             tyvars
                             []           -- no stupid theta
                             rhs'
                             rec_flag     -- FIXME: is this ok?
                             False        -- FIXME: no generics
                             False        -- not GADT syntax
                             Nothing      -- not a family instance
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
  where
    name   = tyConName tc
    tyvars = tyConTyVars tc
    rec_flag = boolToRecFlag (isRecursiveTyCon tc)

vectAlgTyConRhs :: AlgTyConRhs -> VM AlgTyConRhs
vectAlgTyConRhs (DataTyCon { data_cons = data_cons
                           , is_enum   = is_enum
                           })
  = do
      data_cons' <- mapM vectDataCon data_cons
      zipWithM_ defDataCon data_cons data_cons'
      return $ DataTyCon { data_cons = data_cons'
                         , is_enum   = is_enum
                         }

vectDataCon :: DataCon -> VM DataCon
vectDataCon dc
  | not . null $ dataConExTyVars dc = pprPanic "vectDataCon: existentials" (ppr dc)
  | not . null $ dataConEqSpec   dc = pprPanic "vectDataCon: eq spec" (ppr dc)
  | otherwise
  = do
      name'    <- cloneName mkVectDataConOcc name
      tycon'   <- vectTyCon tycon
      arg_tys  <- mapM vectType rep_arg_tys
177
178
179
180
181
182
183
184
185
186
187

      liftDs $ buildDataCon name'
                            False           -- not infix
                            (map (const NotMarkedStrict) arg_tys)
                            []              -- no labelled fields
                            univ_tvs
                            []              -- no existential tvs for now
                            []              -- no eq spec for now
                            []              -- no context
                            arg_tys
                            tycon'
188
189
190
  where
    name        = dataConName dc
    univ_tvs    = dataConUnivTyVars dc
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
191
    rep_arg_tys = dataConRepArgTys dc
192
193
    tycon       = dataConTyCon dc

194
195
196
197
mk_fam_inst :: TyCon -> TyCon -> (TyCon, [Type])
mk_fam_inst fam_tc arg_tc
  = (fam_tc, [mkTyConApp arg_tc . mkTyVarTys $ tyConTyVars arg_tc])

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
198
199
200
201
buildPReprTyCon :: TyCon -> TyCon -> VM TyCon
buildPReprTyCon orig_tc vect_tc
  = do
      name     <- cloneName mkPReprTyConOcc (tyConName orig_tc)
202
      rhs_ty   <- buildPReprType vect_tc
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
203
204
205
206
207
208
209
210
      prepr_tc <- builtin preprTyCon
      liftDs $ buildSynTyCon name
                             tyvars
                             (SynonymTyCon rhs_ty)
                             (Just $ mk_fam_inst prepr_tc vect_tc)
  where
    tyvars = tyConTyVars vect_tc

211
buildPReprType :: TyCon -> VM Type
212
buildPReprType = liftM repr_type . mkTyConRepr
213

214
buildToPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
215
216
217
218
219
220
221
222
buildToPRepr (TyConRepr {
                repr_tys         = repr_tys
              , repr_prod_tycons = prod_tycons
              , repr_prod_tys    = prod_tys
              , repr_sum_tycon   = repr_sum_tycon
              , repr_type        = repr_type
              })
              vect_tc prepr_tc _
223
  = do
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
224
225
      arg  <- newLocalVar FSLIT("x") arg_ty
      vars <- mapM (mapM (newLocalVar FSLIT("x"))) repr_tys
226
227
228

      return . Lam arg
             . wrapFamInstBody prepr_tc var_tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
229
230
231
             . Case (Var arg) (mkWildId arg_ty) repr_type
             . mk_alts data_cons vars
             . zipWith3 mk_prod prod_tycons repr_tys $ map (map Var) vars
232
233
234
235
  where
    var_tys   = mkTyVarTys $ tyConTyVars vect_tc
    arg_ty    = mkTyConApp vect_tc var_tys
    data_cons = tyConDataCons vect_tc
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
236

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
    Just sum_tycon = repr_sum_tycon
    sum_datacons   = tyConDataCons sum_tycon

    mk_alts _    _      []     = [(DEFAULT, [], Var unitDataConId)]
    mk_alts [dc] [vars] [expr] = [(DataAlt dc, vars, expr)]
    mk_alts dcs  vars   exprs  = zipWith4 mk_alt dcs vars sum_datacons exprs 

    mk_alt dc vars sum_dc expr = (DataAlt dc, vars,
                                  mkConApp sum_dc (map Type prod_tys ++ [expr]))

    mk_prod _         _   []     = Var unitDataConId
    mk_prod _         _   [expr] = expr
    mk_prod (Just tc) tys exprs  = mkConApp dc (map Type tys ++ exprs)
      where
        [dc] = tyConDataCons tc
252

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
buildFromPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
buildFromPRepr _ vect_tc prepr_tc _
  = do
      arg_ty <- mkPReprType res_ty
      arg <- newLocalVar FSLIT("x") arg_ty
      alts <- mapM mk_alt data_cons
      body <- mkFromPRepr (unwrapFamInstScrut prepr_tc var_tys (Var arg))
                          res_ty alts
      return $ Lam arg body
  where
    var_tys   = mkTyVarTys $ tyConTyVars vect_tc
    res_ty    = mkTyConApp vect_tc var_tys
    data_cons = tyConDataCons vect_tc

    mk_alt dc = do
                  bndrs <- mapM (newLocalVar FSLIT("x")) $ dataConRepArgTys dc
                  return (bndrs, mkConApp dc (map Type var_tys ++ map Var bndrs))

271
buildToArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
buildToArrPRepr _ vect_tc prepr_tc arr_tc
  = do
      arg_ty  <- mkPArrayType el_ty
      rep_tys <- mapM (mapM mkPArrayType) rep_el_tys

      arg     <- newLocalVar FSLIT("xs") arg_ty
      bndrss  <- mapM (mapM (newLocalVar FSLIT("ys"))) rep_tys
      len     <- newLocalVar FSLIT("len") intPrimTy
      sel     <- newLocalVar FSLIT("sel") =<< mkPArrayType intTy

      let add_sel xs | has_selector = sel : xs
                     | otherwise    = xs

          all_bndrs = len : add_sel (concat bndrss)

      res      <- parrayCoerce prepr_tc var_tys
                =<< mkToArrPRepr (Var len) (Var sel) (map (map Var) bndrss)
      res_ty   <- mkPArrayType =<< mkPReprType el_ty

      return . Lam arg
             $ Case (unwrapFamInstScrut arr_tc var_tys (Var arg))
                    (mkWildId (mkTyConApp arr_tc var_tys))
                    res_ty
                    [(DataAlt arr_dc, all_bndrs, res)]
  where
    var_tys    = mkTyVarTys $ tyConTyVars vect_tc
    el_ty      = mkTyConApp vect_tc var_tys
    data_cons  = tyConDataCons vect_tc
    rep_el_tys = map dataConRepArgTys data_cons

    [arr_dc]   = tyConDataCons arr_tc

    has_selector | [_] <- data_cons = False
                 | otherwise        = True

307
buildFromArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
308
309
310
buildFromArrPRepr _ vect_tc prepr_tc arr_tc
  = mkFromArrPRepr undefined undefined undefined undefined undefined undefined

311
buildPRDict :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
312
313
314
315
316
317
buildPRDict _ vect_tc prepr_tc _
  = prCoerce prepr_tc var_tys
  =<< prDictOfType (mkTyConApp prepr_tc var_tys)
  where
    var_tys = mkTyVarTys $ tyConTyVars vect_tc

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
318
319
buildPArrayTyCon :: TyCon -> TyCon -> VM TyCon
buildPArrayTyCon orig_tc vect_tc = fixV $ \repr_tc ->
320
  do
321
322
    name'  <- cloneName mkPArrayTyConOcc orig_name
    rhs    <- buildPArrayTyConRhs orig_name vect_tc repr_tc
323
324
325
326
327
328
329
330
331
    parray <- builtin parrayTyCon

    liftDs $ buildAlgTyCon name'
                           tyvars
                           []          -- no stupid theta
                           rhs
                           rec_flag    -- FIXME: is this ok?
                           False       -- FIXME: no generics
                           False       -- not GADT syntax
332
                           (Just $ mk_fam_inst parray vect_tc)
333
  where
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
334
    orig_name = tyConName orig_tc
335
336
    tyvars = tyConTyVars vect_tc
    rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
337
338
    

339
340
buildPArrayTyConRhs :: Name -> TyCon -> TyCon -> VM AlgTyConRhs
buildPArrayTyConRhs orig_name vect_tc repr_tc
341
  = do
342
      data_con <- buildPArrayDataCon orig_name vect_tc repr_tc
343
      return $ DataTyCon { data_cons = [data_con], is_enum = False }
344

345
346
buildPArrayDataCon :: Name -> TyCon -> TyCon -> VM DataCon
buildPArrayDataCon orig_name vect_tc repr_tc
347
  = do
348
      dc_name  <- cloneName mkPArrayDataConOcc orig_name
349
      shape    <- tyConShape vect_tc
350
      repr_tys <- mapM mkPArrayType types
351
352
353
354
355
356
357
358
359
360
361

      liftDs $ buildDataCon dc_name
                            False                  -- not infix
                            (shapeStrictness shape ++ map (const NotMarkedStrict) repr_tys)
                            []                     -- no field labels
                            (tyConTyVars vect_tc)
                            []                     -- no existentials
                            []                     -- no eq spec
                            []                     -- no context
                            (shapeReprTys shape ++ repr_tys)
                            repr_tc
362
  where
363
    types = [ty | dc <- tyConDataCons vect_tc
364
365
                , ty <- dataConRepArgTys dc]

366
367
368
369
mkPADFun :: TyCon -> VM Var
mkPADFun vect_tc
  = newExportedVar (mkPADFunOcc $ getOccName vect_tc) =<< paDFunType vect_tc

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
data Shape = Shape {
               shapeReprTys    :: [Type]
             , shapeStrictness :: [StrictnessMark]
             , shapeLength     :: [CoreExpr] -> VM CoreExpr
             , shapeReplicate  :: CoreExpr -> CoreExpr -> VM [CoreExpr]
             }

tyConShape :: TyCon -> VM Shape
tyConShape vect_tc
  | isProductTyCon vect_tc
  = return $ Shape {
                shapeReprTys    = [intPrimTy]
              , shapeStrictness = [NotMarkedStrict]
              , shapeLength     = \[len] -> return len
              , shapeReplicate  = \len _ -> return [len]
              }

  | otherwise
  = do
      repr_ty <- mkPArrayType intTy   -- FIXME: we want to unbox this
      return $ Shape {
                 shapeReprTys    = [repr_ty]
               , shapeStrictness = [MarkedStrict]
               , shapeLength     = \[sel] -> lengthPA sel
               , shapeReplicate  = \len n -> do
                                               e <- replicatePA len n
                                               return [e]
               }
398

399
400
401
buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> Var
                   -> VM [(Var, CoreExpr)]
buildTyConBindings orig_tc vect_tc prepr_tc arr_tc dfun
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
402
403
  = do
      shape <- tyConShape vect_tc
404
      repr  <- mkTyConRepr vect_tc
405
406
407
      sequence_ (zipWith4 (vectDataConWorker shape vect_tc arr_tc arr_dc)
                          orig_dcs
                          vect_dcs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
408
409
                          (inits repr_tys)
                          (tails repr_tys))
410
      dict <- buildPADict repr vect_tc prepr_tc arr_tc dfun
411
412
      binds <- takeHoisted
      return $ (dfun, dict) : binds
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
413
414
415
416
417
418
419
420
  where
    orig_dcs = tyConDataCons orig_tc
    vect_dcs = tyConDataCons vect_tc
    [arr_dc] = tyConDataCons arr_tc

    repr_tys = map dataConRepArgTys vect_dcs

vectDataConWorker :: Shape -> TyCon -> TyCon -> DataCon
421
                  -> DataCon -> DataCon -> [[Type]] -> [[Type]]
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
422
                  -> VM ()
423
vectDataConWorker shape vect_tc arr_tc arr_dc orig_dc vect_dc pre (dc_tys : post)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
  = do
      clo <- closedV
           . inBind orig_worker
           . polyAbstract tvs $ \abstract ->
             liftM (abstract . vectorised)
           $ buildClosures tvs [] dc_tys res_ty (liftM2 (,) mk_vect mk_lift)

      worker <- cloneId mkVectOcc orig_worker (exprType clo)
      hoistBinding worker clo
      defGlobalVar orig_worker worker
      return ()
  where
    tvs     = tyConTyVars vect_tc
    arg_tys = mkTyVarTys tvs
    res_ty  = mkTyConApp vect_tc arg_tys

    orig_worker = dataConWorkId orig_dc

    mk_vect = return . mkConApp vect_dc $ map Type arg_tys
    mk_lift = do
                len     <- newLocalVar FSLIT("n") intPrimTy
                arr_tys <- mapM mkPArrayType dc_tys
                args    <- mapM (newLocalVar FSLIT("xs")) arr_tys
447
448
                shapes  <- shapeReplicate shape
                                          (Var len)
449
                                          (mkDataConTag vect_dc)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
450
451
452
453
454
455
456
457
458
459
460
461
                
                empty_pre  <- mapM emptyPA (concat pre)
                empty_post <- mapM emptyPA (concat post)

                return . mkLams (len : args)
                       . wrapFamInstBody arr_tc arg_tys
                       . mkConApp arr_dc
                       $ map Type arg_tys ++ shapes
                                          ++ empty_pre
                                          ++ map Var args
                                          ++ empty_post

462
463
buildPADict :: TyConRepr -> TyCon -> TyCon -> TyCon -> Var -> VM CoreExpr
buildPADict repr vect_tc prepr_tc arr_tc dfun
464
  = polyAbstract tvs $ \abstract ->
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
465
    do
466
      meth_binds <- mapM (mk_method repr) paMethods
467
      let meth_exprs = map (Var . fst) meth_binds
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
468

469
      pa_dc <- builtin paDataCon
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
470
      let dict = mkConApp pa_dc (Type (mkTyConApp vect_tc arg_tys) : meth_exprs)
471
          body = Let (Rec meth_binds) dict
472
      return . mkInlineMe $ abstract body
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
473
474
475
476
  where
    tvs = tyConTyVars arr_tc
    arg_tys = mkTyVarTys tvs

477
    mk_method repr (name, build)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
478
479
      = localV
      $ do
480
          body <- build repr vect_tc prepr_tc arr_tc
481
          var  <- newLocalVar name (exprType body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
482
483
          return (var, mkInlineMe body)
          
484
485
486
487
488
paMethods = [(FSLIT("toPRepr"),      buildToPRepr),
             (FSLIT("fromPRepr"),    buildFromPRepr),
             (FSLIT("toArrPRepr"),   buildToArrPRepr),
             (FSLIT("fromArrPRepr"), buildFromArrPRepr),
             (FSLIT("dictPRepr"),    buildPRDict)]
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
489

490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
-- | Split the given tycons into two sets depending on whether they have to be
-- converted (first list) or not (second list). The first argument contains
-- information about the conversion status of external tycons:
-- 
--   * tycons which have converted versions are mapped to True
--   * tycons which are not changed by vectorisation are mapped to False
--   * tycons which can't be converted are not elements of the map
--
classifyTyCons :: UniqFM Bool -> [TyConGroup] -> ([TyCon], [TyCon])
classifyTyCons = classify [] []
  where
    classify conv keep cs [] = (conv, keep)
    classify conv keep cs ((tcs, ds) : rs)
      | can_convert && must_convert
        = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
      | can_convert
        = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
      | otherwise
        = classify conv keep cs rs
      where
        refs = ds `delListFromUniqSet` tcs

        can_convert  = isNullUFM (refs `minusUFM` cs) && all convertable tcs
        must_convert = foldUFM (||) False (intersectUFM_C const cs refs)

        convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
    
-- | Compute mutually recursive groups of tycons in topological order
--
tyConGroups :: [TyCon] -> [TyConGroup]
tyConGroups tcs = map mk_grp (stronglyConnComp edges)
  where
    edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
                                , let ds = tyConsOfTyCon tc]

    mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
    mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
      where
        (tcs, dss) = unzip els

tyConsOfTyCon :: TyCon -> UniqSet TyCon
tyConsOfTyCon 
  = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons

tyConsOfType :: Type -> UniqSet TyCon
tyConsOfType ty
  | Just ty' <- coreView ty    = tyConsOfType ty'
tyConsOfType (TyVarTy v)       = emptyUniqSet
538
539
540
541
542
543
544
tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
  where
    extend | isUnLiftedTyCon tc
           || isTupleTyCon   tc = id

           | otherwise          = (`addOneToUniqSet` tc)

545
546
547
548
549
550
551
552
553
tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
                                 `addOneToUniqSet` funTyCon
tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty
tyConsOfType other             = pprPanic "ClosureConv.tyConsOfType" $ ppr other

tyConsOfTypes :: [Type] -> UniqSet TyCon
tyConsOfTypes = unionManyUniqSets . map tyConsOfType