VectType.hs 32.6 KB
Newer Older
1
{-# OPTIONS -w #-}
2
3
4
-- The above warning supression flag is a temporary kludge.
-- While working on this module you are encouraged to remove it and fix
-- any warnings in the module. See
Ian Lynagh's avatar
Ian Lynagh committed
5
--     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#Warnings
6
7
-- for details

8
module VectType ( vectTyCon, vectType, vectTypeEnv,
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
9
10
                  mkRepr, arrShapeTys, arrShapeVars, arrSelector,
                  PAInstance, buildPADict )
11
12
13
14
15
16
where

#include "HsVersions.h"

import VectMonad
import VectUtils
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
17
import VectCore
18

19
import HscTypes          ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
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 BuildTyCl
23
import DataCon
24
25
26
import TyCon
import Type
import TypeRep
27
import Coercion
28
import FamInstEnv        ( FamInst, mkLocalFamInst )
29
import InstEnv           ( Instance, mkLocalInstance, instanceDFunId )
30
31
import OccName
import MkId
32
import BasicTypes        ( StrictnessMark(..), OverlapFlag(..), boolToRecFlag )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
33
import Var               ( Var )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
34
import Id                ( mkWildId )
35
import Name              ( Name, getOccName )
36
import NameEnv
37
import TysWiredIn
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
38
import TysPrim           ( intPrimTy )
39

40
import Unique
41
42
import UniqFM
import UniqSet
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
43
import Util              ( singleton )
44
45
import Digraph           ( SCC(..), stronglyConnComp )

46
47
import Outputable

48
import Control.Monad  ( liftM, liftM2, zipWithM, zipWithM_, mapAndUnzipM )
49
import Data.List      ( inits, tails, zipWith4, zipWith5 )
50

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
51
52
53
-- ----------------------------------------------------------------------------
-- Types

54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
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)

84
85
86
87
88
-- ----------------------------------------------------------------------------
-- Type definitions

type TyConGroup = ([TyCon], UniqSet TyCon)

89
data PAInstance = PAInstance {
90
                    painstDFun      :: Var
91
                  , painstOrigTyCon :: TyCon
92
93
94
95
                  , painstVectTyCon :: TyCon
                  , painstArrTyCon  :: TyCon
                  }

96
vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [(Var, CoreExpr)])
97
98
99
100
101
102
103
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
104
105
106
107
108
      new_tcs <- vectTyConDecls conv_tcs

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

109
      repr_tcs <- zipWithM buildPReprTyCon   orig_tcs vect_tcs
110
      parr_tcs <- zipWithM buildPArrayTyCon orig_tcs vect_tcs
111
112
      dfuns    <- mapM mkPADFun vect_tcs
      defTyConPAs (zip vect_tcs dfuns)
113
114
115
116
117
118
      binds    <- sequence (zipWith5 buildTyConBindings orig_tcs
                                                        vect_tcs
                                                        repr_tcs
                                                        parr_tcs
                                                        dfuns)

119
      let all_new_tcs = new_tcs ++ repr_tcs ++ parr_tcs
120
121

      let new_env = extendTypeEnvList env
122
123
                       (map ATyCon all_new_tcs
                        ++ [ADataCon dc | tc <- all_new_tcs
124
125
                                        , dc <- tyConDataCons tc])

126
      return (new_env, map mkLocalFamInst (repr_tcs ++ parr_tcs), concat binds)
127
128
129
130
131
132
133
134
135
136
137
  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


138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
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)

153
154
155
156
157
158
159
160
      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
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
  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
186
187
188
189
190
191
192
193
194
195
196

      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'
197
198
199
  where
    name        = dataConName dc
    univ_tvs    = dataConUnivTyVars dc
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
200
    rep_arg_tys = dataConRepArgTys dc
201
202
    tycon       = dataConTyCon dc

203
204
205
206
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
207
208
209
210
buildPReprTyCon :: TyCon -> TyCon -> VM TyCon
buildPReprTyCon orig_tc vect_tc
  = do
      name     <- cloneName mkPReprTyConOcc (tyConName orig_tc)
211
      rhs_ty   <- buildPReprType vect_tc
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
212
213
214
215
216
217
218
219
      prepr_tc <- builtin preprTyCon
      liftDs $ buildSynTyCon name
                             tyvars
                             (SynonymTyCon rhs_ty)
                             (Just $ mk_fam_inst prepr_tc vect_tc)
  where
    tyvars = tyConTyVars vect_tc

220

221
222
223
224
225
226
227
data Repr = ProdRepr {
              prod_components   :: [Type]
            , prod_tycon        :: TyCon
            , prod_data_con     :: DataCon
            , prod_arr_tycon    :: TyCon
            , prod_arr_data_con :: DataCon
            }
228

229
230
231
232
233
234
235
          | SumRepr {
              sum_components    :: [Repr]
            , sum_tycon         :: TyCon
            , sum_arr_tycon     :: TyCon
            , sum_arr_data_con  :: DataCon
            }

236
237
          | IdRepr Type

238
239
240
241
242
          | VoidRepr {
              void_tycon        :: TyCon
            , void_bottom       :: CoreExpr
            }

243
244
245
246
247
248
249
          | EnumRepr {
              enum_tycon        :: TyCon
            , enum_data_con     :: DataCon
            , enum_arr_tycon    :: TyCon
            , enum_arr_data_con :: DataCon
            }

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
250
251
252
253
254
255
256
257
258
259
voidRepr :: VM Repr
voidRepr
  = do
      tycon <- builtin voidTyCon
      var   <- builtin voidVar
      return $ VoidRepr {
                 void_tycon  = tycon
               , void_bottom = Var var
               }

260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
enumRepr :: VM Repr
enumRepr
  = do
      (arr_tycon, _) <- parrayReprTyCon intTy
      let [arr_data_con] = tyConDataCons arr_tycon

      return $ EnumRepr {
                 enum_tycon       = tycon
               , enum_data_con     = data_con
               , enum_arr_tycon    = arr_tycon
               , enum_arr_data_con = arr_data_con
               }
  where
    tycon = intTyCon
    data_con = intDataCon

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
276
277
278
279
280
281
282
unboxedProductRepr :: [Type] -> VM Repr
unboxedProductRepr []   = voidRepr
unboxedProductRepr [ty] = return $ IdRepr ty
unboxedProductRepr tys  = boxedProductRepr tys

boxedProductRepr :: [Type] -> VM Repr
boxedProductRepr tys
283
284
285
  = do
      tycon <- builtin (prodTyCon arity)
      let [data_con] = tyConDataCons tycon
286

287
      (arr_tycon, _) <- parrayReprTyCon $ mkTyConApp tycon tys
288
289
      let [arr_data_con] = tyConDataCons arr_tycon

290
      return $ ProdRepr {
291
292
293
294
295
                 prod_components   = tys
               , prod_tycon        = tycon
               , prod_data_con     = data_con
               , prod_arr_tycon    = arr_tycon
               , prod_arr_data_con = arr_data_con
296
               }
297
298
  where
    arity = length tys
299

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
300
301
302
303
sumRepr :: [Repr] -> VM Repr
sumRepr []     = voidRepr
sumRepr [repr] = boxRepr repr
sumRepr reprs
304
  = do
305
306
307
308
309
310
      tycon <- builtin (sumTyCon arity)
      (arr_tycon, _) <- parrayReprTyCon
                      . mkTyConApp tycon
                      $ map reprType reprs

      let [arr_data_con] = tyConDataCons arr_tycon
311
312

      return $ SumRepr {
313
314
315
316
                 sum_components   = reprs
               , sum_tycon        = tycon
               , sum_arr_tycon    = arr_tycon
               , sum_arr_data_con = arr_data_con
317
318
               }
  where
319
    arity = length reprs
320

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
321
322
323
324
splitSumRepr :: Repr -> [Repr]
splitSumRepr (SumRepr { sum_components = reprs }) = reprs
splitSumRepr repr                                 = [repr]

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
325
326
327
328
329
boxRepr :: Repr -> VM Repr
boxRepr (VoidRepr {}) = boxedProductRepr []
boxRepr (IdRepr ty)   = boxedProductRepr [ty]
boxRepr repr          = return repr

330
331
332
333
334
reprType :: Repr -> Type
reprType (ProdRepr { prod_tycon = tycon, prod_components = tys })
  = mkTyConApp tycon tys
reprType (SumRepr { sum_tycon = tycon, sum_components = reprs })
  = mkTyConApp tycon (map reprType reprs)
335
reprType (IdRepr ty) = ty
336
reprType (VoidRepr { void_tycon = tycon }) = mkTyConApp tycon []
337
reprType (EnumRepr { enum_tycon = tycon }) = mkTyConApp tycon []
338

339
340
arrReprType :: Repr -> VM Type
arrReprType = mkPArrayType . reprType
341

342
343
arrShapeTys :: Repr -> VM [Type]
arrShapeTys (SumRepr  {})
344
  = do
345
346
      int_arr <- builtin parrayIntPrimTyCon
      return [intPrimTy, mkTyConApp int_arr [], mkTyConApp int_arr []]
347
348
arrShapeTys (ProdRepr {}) = return [intPrimTy]
arrShapeTys (IdRepr _)    = return []
349
arrShapeTys (VoidRepr {}) = return [intPrimTy]
350
arrShapeTys (EnumRepr {}) = return [intPrimTy]
351

352
353
arrShapeVars :: Repr -> VM [Var]
arrShapeVars repr = mapM (newLocalVar FSLIT("sh")) =<< arrShapeTys repr
354

355
356
replicateShape :: Repr -> CoreExpr -> CoreExpr -> VM [CoreExpr]
replicateShape (ProdRepr {}) len _ = return [len]
357
358
359
360
361
replicateShape (SumRepr {})  len tag
  = do
      rep <- builtin replicatePAIntPrimVar
      up  <- builtin upToPAIntPrimVar
      return [len, Var rep `mkApps` [len, tag], Var up `App` len]
362
replicateShape (IdRepr _) _ _ = return []
363
replicateShape (VoidRepr {}) len _ = return [len]
364
replicateShape (EnumRepr {}) len _ = return [len]
365

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
366
367
368
369
arrSelector :: Repr -> [a] -> a
arrSelector (SumRepr {}) [_, sel, _] = sel
arrSelector _ _ = panic "arrSelector"

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
370
371
372
373
374
375
376
377
378
379
380
emptyArrRepr :: Repr -> VM [CoreExpr]
emptyArrRepr (SumRepr { sum_components = prods })
  = liftM concat $ mapM emptyArrRepr prods
emptyArrRepr (ProdRepr { prod_components = [] })
  = return [Var unitDataConId]
emptyArrRepr (ProdRepr { prod_components = tys })
  = mapM emptyPA tys
emptyArrRepr (IdRepr ty)
  = liftM singleton $ emptyPA ty
emptyArrRepr (VoidRepr { void_tycon = tycon })
  = liftM singleton $ emptyPA (mkTyConApp tycon [])
381
382
emptyArrRepr (EnumRepr { enum_tycon = tycon })
  = liftM singleton $ emptyPA (mkTyConApp tycon [])
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
383
384
385
386
387
388
389
390
391
392
393
394

arrReprTys :: Repr -> VM [Type]
arrReprTys (SumRepr { sum_components = reprs })
  = liftM concat $ mapM arrReprTys reprs
arrReprTys (ProdRepr { prod_components = [] })
  = return [unitTy]
arrReprTys (ProdRepr { prod_components = tys })
  = mapM mkPArrayType tys
arrReprTys (IdRepr ty)
  = liftM singleton $ mkPArrayType ty
arrReprTys (VoidRepr { void_tycon = tycon })
  = liftM singleton $ mkPArrayType (mkTyConApp tycon [])
395
396
arrReprTys (EnumRepr {})
  = liftM singleton $ mkPArrayType intPrimTy
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
397
398
399
400
401

arrReprTys' :: Repr -> VM [[Type]]
arrReprTys' (SumRepr { sum_components = reprs })
  = mapM arrReprTys reprs
arrReprTys' repr = liftM singleton $ arrReprTys repr
402
403
404

arrReprVars :: Repr -> VM [[Var]]
arrReprVars repr
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
405
  = mapM (mapM (newLocalVar FSLIT("rs"))) =<< arrReprTys' repr
406
407
408

mkRepr :: TyCon -> VM Repr
mkRepr vect_tc
409
410
411
  | [tys] <- rep_tys = boxedProductRepr tys
  | all null rep_tys = enumRepr
  | otherwise        = sumRepr =<< mapM unboxedProductRepr rep_tys
412
413
  where
    rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
414
415
416
417
418
419

buildPReprType :: TyCon -> VM Type
buildPReprType = liftM reprType . mkRepr

buildToPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
buildToPRepr repr vect_tc prepr_tc _
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
420
  = do
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
421
      arg    <- newLocalVar FSLIT("x") arg_ty
422
      result <- to_repr repr (Var arg)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
423

424
425
426
      return . Lam arg
             . wrapFamInstBody prepr_tc var_tys
             $ result
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
427
  where
428
429
430
    var_tys = mkTyVarTys $ tyConTyVars vect_tc
    arg_ty  = mkTyConApp vect_tc var_tys
    res_ty  = reprType repr
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
431

432
433
    cons    = tyConDataCons vect_tc
    [con]   = cons
434

435
436
437
    to_repr (SumRepr { sum_components = prods
                     , sum_tycon      = tycon })
            expr
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
438
      = do
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
439
          (vars, bodies) <- mapAndUnzipM to_unboxed prods
440
441
442
443
444
445
446
447
          return . Case expr (mkWildId (exprType expr)) res_ty
                 $ zipWith4 mk_alt cons vars (tyConDataCons tycon) bodies
      where
        mk_alt con vars sum_con body
          = (DataAlt con, vars, mkConApp sum_con (ty_args ++ [body]))

        ty_args = map (Type . reprType) prods

448
449
450
451
452
453
    to_repr (EnumRepr { enum_data_con = data_con }) expr
      = return . Case expr (mkWildId (exprType expr)) res_ty
               $ map mk_alt cons
      where
        mk_alt con = (DataAlt con, [], mkConApp data_con [mkDataConTag con])

454
    to_repr prod expr
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
455
      = do
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
456
          (vars, body) <- to_unboxed prod
457
458
          return $ Case expr (mkWildId (exprType expr)) res_ty
                   [(DataAlt con, vars, body)]
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
459

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
460
461
    to_unboxed (ProdRepr { prod_components = tys
                         , prod_data_con   = data_con })
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
462
      = do
463
464
          vars <- mapM (newLocalVar FSLIT("r")) tys
          return (vars, mkConApp data_con (map Type tys ++ map Var vars))
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
465

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
466
    to_unboxed (IdRepr ty)
467
468
469
470
      = do
          var <- newLocalVar FSLIT("y") ty
          return ([var], Var var)

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
471
    to_unboxed (VoidRepr { void_bottom = bottom })
472
473
474
      = return ([], bottom)


475
476
477
478
479
buildFromPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
buildFromPRepr repr vect_tc prepr_tc _
  = do
      arg_ty <- mkPReprType res_ty
      arg    <- newLocalVar FSLIT("x") arg_ty
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
480

481
482
483
484
485
486
      liftM (Lam arg)
           . from_repr repr
           $ unwrapFamInstScrut prepr_tc var_tys (Var arg)
  where
    var_tys = mkTyVarTys $ tyConTyVars vect_tc
    res_ty  = mkTyConApp vect_tc var_tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
487

488
489
    cons    = map (`mkConApp` map Type var_tys) (tyConDataCons vect_tc)
    [con]   = cons
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
490

491
492
493
494
495
    from_repr repr@(SumRepr { sum_components = prods
                            , sum_tycon      = tycon })
              expr
      = do
          vars   <- mapM (newLocalVar FSLIT("x")) (map reprType prods)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
496
          bodies <- sequence . zipWith3 from_unboxed prods cons
497
498
499
500
501
                             $ map Var vars
          return . Case expr (mkWildId (reprType repr)) res_ty
                 $ zipWith3 sum_alt (tyConDataCons tycon) vars bodies
      where
        sum_alt data_con var body = (DataAlt data_con, [var], body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
502

503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
    from_repr repr@(EnumRepr { enum_data_con = data_con }) expr
      = do
          var <- newLocalVar FSLIT("n") intPrimTy

          let res = Case (Var var) (mkWildId intPrimTy) res_ty
                  $ (DEFAULT, [], error_expr)
                  : zipWith mk_alt (tyConDataCons vect_tc) cons

          return $ Case expr (mkWildId (reprType repr)) res_ty
                   [(DataAlt data_con, [var], res)]
      where
        mk_alt data_con con = (LitAlt (mkDataConTagLit data_con), [], con)

        error_expr = mkRuntimeErrorApp rUNTIME_ERROR_ID res_ty
                   . showSDoc
                   $ sep [text "Invalid NDP representation of", ppr vect_tc]

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
520
    from_repr repr expr = from_unboxed repr con expr
521

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
522
523
    from_unboxed prod@(ProdRepr { prod_components = tys
                                , prod_data_con   = data_con })
524
525
526
527
528
529
530
              con
              expr
      = do
          vars <- mapM (newLocalVar FSLIT("y")) tys
          return $ Case expr (mkWildId (reprType prod)) res_ty
                   [(DataAlt data_con, vars, con `mkVarApps` vars)]

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
531
    from_unboxed (IdRepr _) con expr
532
533
       = return $ con `App` expr

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
534
    from_unboxed (VoidRepr {}) con expr
535
536
       = return con

537
538
buildToArrPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
buildToArrPRepr repr vect_tc prepr_tc arr_tc
539
  = do
540
541
      arg_ty     <- mkPArrayType el_ty
      arg        <- newLocalVar FSLIT("xs") arg_ty
542
543
544
545
546

      res_ty     <- mkPArrayType (reprType repr)

      shape_vars <- arrShapeVars repr
      repr_vars  <- arrReprVars  repr
547
548
549

      parray_co  <- mkBuiltinCo parrayTyCon

550
551
552
553
554
555
      let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
          co           = mkAppCoercion parray_co
                       . mkSymCoercion
                       $ mkTyConApp repr_co var_tys

          scrut   = unwrapFamInstScrut arr_tc var_tys (Var arg)
556

557
      result <- to_repr shape_vars repr_vars repr
558
559
560

      return . Lam arg
             . mkCoerce co
561
562
             $ Case scrut (mkWildId (mkTyConApp arr_tc var_tys)) res_ty
               [(DataAlt arr_dc, shape_vars ++ concat repr_vars, result)]
563
564
565
  where
    var_tys = mkTyVarTys $ tyConTyVars vect_tc
    el_ty   = mkTyConApp vect_tc var_tys
566
567
568

    [arr_dc] = tyConDataCons arr_tc

569
570
571
572
573
574
    to_repr shape_vars@(len_var : _)
            repr_vars
            (SumRepr { sum_components   = prods
                     , sum_arr_tycon    = tycon
                     , sum_arr_data_con = data_con })
      = do
575
          exprs <- zipWithM to_prod repr_vars prods
576

577
578
579
580
581
          return . wrapFamInstBody tycon tys
                 . mkConApp data_con
                 $ map Type tys ++ map Var shape_vars ++ exprs
      where
        tys = map reprType prods
582

583
584
585
586
587
588
589
590
    to_repr [len_var]
            [repr_vars]
            (ProdRepr { prod_components   = tys
                      , prod_arr_tycon    = tycon
                      , prod_arr_data_con = data_con })
       = return . wrapFamInstBody tycon tys
                . mkConApp data_con
                $ map Type tys ++ map Var (len_var : repr_vars)
591

592
593
594
595
596
597
598
    to_repr [len_var]
            [[repr_var]]
            (EnumRepr { enum_arr_tycon    = tycon
                      , enum_arr_data_con = data_con })
      = return . wrapFamInstBody tycon []
               $ mkConApp data_con [Var len_var, Var repr_var]

599
    to_prod repr_vars@(r : _)
600
            (ProdRepr { prod_components   = tys@(ty : _)
601
602
                      , prod_arr_tycon    = tycon
                      , prod_arr_data_con = data_con })
603
      = do
604
          len <- lengthPA ty (Var r)
605
606
607
608
          return . wrapFamInstBody tycon tys
                 . mkConApp data_con
                 $ map Type tys ++ len : map Var repr_vars

609
610
611
    to_prod [var] (IdRepr ty)   = return (Var var)
    to_prod [var] (VoidRepr {}) = return (Var var)

612

613
614
buildFromArrPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
buildFromArrPRepr repr vect_tc prepr_tc arr_tc
615
  = do
616
      arg_ty     <- mkPArrayType =<< mkPReprType el_ty
617
618
      arg        <- newLocalVar FSLIT("xs") arg_ty

619
620
621
622
      res_ty     <- mkPArrayType el_ty

      shape_vars <- arrShapeVars repr
      repr_vars  <- arrReprVars  repr
623
624
625

      parray_co  <- mkBuiltinCo parrayTyCon

626
627
628
      let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
          co           = mkAppCoercion parray_co
                       $ mkTyConApp repr_co var_tys
629

630
          scrut  = mkCoerce co (Var arg)
631

632
633
634
          result = wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
                 $ map Type var_tys ++ map Var (shape_vars ++ concat repr_vars)
635

636
637
      liftM (Lam arg)
            (from_repr repr scrut shape_vars repr_vars res_ty result)
638
639
640
641
642
  where
    var_tys = mkTyVarTys $ tyConTyVars vect_tc
    el_ty   = mkTyConApp vect_tc var_tys

    [arr_dc] = tyConDataCons arr_tc
643

644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
    from_repr (SumRepr { sum_components   = prods
                       , sum_arr_tycon    = tycon
                       , sum_arr_data_con = data_con })
              expr
              shape_vars
              repr_vars
              res_ty
              body
      = do
          vars <- mapM (newLocalVar FSLIT("xs")) =<< mapM arrReprType prods
          result <- go prods repr_vars vars body

          let scrut = unwrapFamInstScrut tycon ty_args expr
          return . Case scrut (mkWildId scrut_ty) res_ty
                 $ [(DataAlt data_con, shape_vars ++ vars, result)]
      where
        ty_args  = map reprType prods
        scrut_ty = mkTyConApp tycon ty_args

        go [] [] [] body = return body
        go (prod : prods) (repr_vars : rss) (var : vars) body
          = do
              shape_vars <- mapM (newLocalVar FSLIT("s")) =<< arrShapeTys prod

              from_prod prod (Var var) shape_vars repr_vars res_ty
                =<< go prods rss vars body

    from_repr repr expr shape_vars [repr_vars] res_ty body
      = from_prod repr expr shape_vars repr_vars res_ty body

    from_prod prod@(ProdRepr { prod_components = tys
                             , prod_arr_tycon  = tycon
                             , prod_arr_data_con = data_con })
              expr
              shape_vars
              repr_vars
              res_ty
              body
      = do
          let scrut    = unwrapFamInstScrut tycon tys expr
              scrut_ty = mkTyConApp tycon tys
          ty <- arrReprType prod

          return $ Case scrut (mkWildId scrut_ty) res_ty
                   [(DataAlt data_con, shape_vars ++ repr_vars, body)]

690
691
692
693
694
695
696
697
698
699
700
701
702
    from_prod (EnumRepr { enum_arr_tycon = tycon
                        , enum_arr_data_con = data_con })
              expr
              [len_var]
              [repr_var]
              res_ty
              body
      = let scrut    = unwrapFamInstScrut tycon [] expr
            scrut_ty = mkTyConApp tycon []
        in
        return $ Case scrut (mkWildId scrut_ty) res_ty
                 [(DataAlt data_con, [len_var, repr_var], body)]

703
704
705
706
707
708
709
710
    from_prod (IdRepr ty)
              expr
              shape_vars
              [repr_var]
              res_ty
              body
      = return $ Let (NonRec repr_var expr) body

711
712
713
714
715
716
717
718
    from_prod (VoidRepr {})
              expr
              shape_vars
              [repr_var]
              res_ty
              body
      = return $ Let (NonRec repr_var expr) body

719
buildPRDictRepr :: Repr -> VM CoreExpr
720
721
buildPRDictRepr (VoidRepr { void_tycon = tycon })
  = prDFunOfTyCon tycon
722
buildPRDictRepr (IdRepr ty) = mkPR ty
723
buildPRDictRepr (ProdRepr {
724
725
                   prod_components = tys
                 , prod_tycon      = tycon
726
                 })
727
  = do
728
729
730
      prs  <- mapM mkPR tys
      dfun <- prDFunOfTyCon tycon
      return $ dfun `mkTyApps` tys `mkApps` prs
731

732
buildPRDictRepr (SumRepr {
733
734
                   sum_components = prods
                 , sum_tycon      = tycon })
735
  = do
736
737
738
      prs  <- mapM buildPRDictRepr prods
      dfun <- prDFunOfTyCon tycon
      return $ dfun `mkTyApps` map reprType prods `mkApps` prs
739

740
741
742
buildPRDictRepr (EnumRepr { enum_tycon = tycon })
  = prDFunOfTyCon tycon

743
buildPRDict :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
buildPRDict repr vect_tc prepr_tc _
  = do
      dict  <- buildPRDictRepr repr

      pr_co <- mkBuiltinCo prTyCon
      let co = mkAppCoercion pr_co
             . mkSymCoercion
             $ mkTyConApp arg_co var_tys

      return $ mkCoerce co dict
  where
    var_tys = mkTyVarTys $ tyConTyVars vect_tc

    Just arg_co = tyConFamilyCoercion_maybe prepr_tc

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
759
760
buildPArrayTyCon :: TyCon -> TyCon -> VM TyCon
buildPArrayTyCon orig_tc vect_tc = fixV $ \repr_tc ->
761
  do
762
763
    name'  <- cloneName mkPArrayTyConOcc orig_name
    rhs    <- buildPArrayTyConRhs orig_name vect_tc repr_tc
764
765
766
767
768
769
770
771
772
    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
773
                           (Just $ mk_fam_inst parray vect_tc)
774
  where
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
775
    orig_name = tyConName orig_tc
776
777
    tyvars = tyConTyVars vect_tc
    rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
778

779

780
781
buildPArrayTyConRhs :: Name -> TyCon -> TyCon -> VM AlgTyConRhs
buildPArrayTyConRhs orig_name vect_tc repr_tc
782
  = do
783
      data_con <- buildPArrayDataCon orig_name vect_tc repr_tc
784
      return $ DataTyCon { data_cons = [data_con], is_enum = False }
785

786
787
buildPArrayDataCon :: Name -> TyCon -> TyCon -> VM DataCon
buildPArrayDataCon orig_name vect_tc repr_tc
788
  = do
789
      dc_name  <- cloneName mkPArrayDataConOcc orig_name
790
      repr     <- mkRepr vect_tc
791

792
793
794
      shape_tys <- arrShapeTys repr
      repr_tys  <- arrReprTys  repr

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
795
      let tys = shape_tys ++ repr_tys
796
797
798

      liftDs $ buildDataCon dc_name
                            False                  -- not infix
799
                            (map (const NotMarkedStrict) tys)
800
801
802
803
804
                            []                     -- no field labels
                            (tyConTyVars vect_tc)
                            []                     -- no existentials
                            []                     -- no eq spec
                            []                     -- no context
805
                            tys
806
                            repr_tc
807

808
809
810
811
mkPADFun :: TyCon -> VM Var
mkPADFun vect_tc
  = newExportedVar (mkPADFunOcc $ getOccName vect_tc) =<< paDFunType vect_tc

812
813
814
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
815
  = do
816
      repr  <- mkRepr vect_tc
817
      vectDataConWorkers repr orig_tc vect_tc arr_tc
818
      dict <- buildPADict repr vect_tc prepr_tc arr_tc dfun
819
820
      binds <- takeHoisted
      return $ (dfun, dict) : binds
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
821
822
823
824
825
826
827
  where
    orig_dcs = tyConDataCons orig_tc
    vect_dcs = tyConDataCons vect_tc
    [arr_dc] = tyConDataCons arr_tc

    repr_tys = map dataConRepArgTys vect_dcs

828
829
830
831
832
833
834
835
vectDataConWorkers :: Repr -> TyCon -> TyCon -> TyCon
                   -> VM ()
vectDataConWorkers repr orig_tc vect_tc arr_tc
  = do
      bs <- sequence
          . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
          $ zipWith4 mk_data_con (tyConDataCons vect_tc)
                                 rep_tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
836
837
                                 (inits reprs)
                                 (tail $ tails reprs)
838
839
840
841
842
843
844
845
846
      mapM_ (uncurry hoistBinding) bs
  where
    tyvars   = tyConTyVars vect_tc
    var_tys  = mkTyVarTys tyvars
    ty_args  = map Type var_tys

    res_ty   = mkTyConApp vect_tc var_tys

    rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
847
    reprs    = splitSumRepr repr
848
849
850
851
852

    [arr_dc] = tyConDataCons arr_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
853
                   (lift_data_con tys pre post (mkDataConTag con))
854
855

    vect_data_con con = return $ mkConApp con ty_args
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
856
    lift_data_con tys pre_reprs post_reprs tag
857
858
859
860
      = do
          len  <- builtin liftingContext
          args <- mapM (newLocalVar FSLIT("xs"))
                  =<< mapM mkPArrayType tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
861

862
863
          shape <- replicateShape repr (Var len) tag
          repr  <- mk_arr_repr (Var len) (map Var args)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
864

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
865
866
          pre   <- liftM concat $ mapM emptyArrRepr pre_reprs
          post  <- liftM concat $ mapM emptyArrRepr post_reprs
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893

          return . mkLams (len : args)
                 . wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
                 $ ty_args ++ shape ++ pre ++ repr ++ post

    mk_arr_repr len []
      = do
          units <- replicatePA len (Var unitDataConId)
          return [units]

    mk_arr_repr len arrs = return arrs

    def_worker data_con arg_tys mk_body
      = do
          body <- closedV
                . inBind orig_worker
                . polyAbstract tyvars $ \abstract ->
                  liftM (abstract . vectorised)
                $ buildClosures tyvars [] arg_tys res_ty mk_body

          vect_worker <- cloneId mkVectOcc orig_worker (exprType body)
          defGlobalVar orig_worker vect_worker
          return (vect_worker, body)
      where
        orig_worker = dataConWorkId data_con

894
buildPADict :: Repr -> TyCon -> TyCon -> TyCon -> Var -> VM CoreExpr
895
buildPADict repr vect_tc prepr_tc arr_tc dfun
896
  = polyAbstract tvs $ \abstract ->
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
897
    do
898
      meth_binds <- mapM (mk_method repr) paMethods
899
      let meth_exprs = map (Var . fst) meth_binds
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
900

901
      pa_dc <- builtin paDataCon
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
902
      let dict = mkConApp pa_dc (Type (mkTyConApp vect_tc arg_tys) : meth_exprs)
903
          body = Let (Rec meth_binds) dict
904
      return . mkInlineMe $ abstract body
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
905
906
907
908
  where
    tvs = tyConTyVars arr_tc
    arg_tys = mkTyVarTys tvs

909
    mk_method repr (name, build)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
910
911
      = localV
      $ do
912
          body <- build repr vect_tc prepr_tc arr_tc
913
          var  <- newLocalVar name (exprType body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
914
          return (var, mkInlineMe body)
915

916
917
918
919
920
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
921

922
923
924
-- | 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:
925
--
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
--   * 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)
948

949
950
951
952
953
954
955
956
957
958
959
960
961
962
-- | 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
963
tyConsOfTyCon
964
965
966
967
968
969
  = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons

tyConsOfType :: Type -> UniqSet TyCon
tyConsOfType ty
  | Just ty' <- coreView ty    = tyConsOfType ty'
tyConsOfType (TyVarTy v)       = emptyUniqSet
970
971
972
973
974
975
976
tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
  where
    extend | isUnLiftedTyCon tc
           || isTupleTyCon   tc = id

           | otherwise          = (`addOneToUniqSet` tc)

977
978
979
980
981
982
983
984
985
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