VectUtils.hs 14 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 VectUtils (
9
  collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
10
  collectAnnValBinders,
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
11
12
13
  dataConTagZ, mkDataConTag, mkDataConTagLit,

  newLocalVVar,
14

15
  mkBuiltinCo,
16
  mkPADictType, mkPArrayType, mkPReprType,
17

18
19
  parrayReprTyCon, parrayReprDataCon, mkVScrut,
  prDFunOfTyCon,
20
  paDictArgType, paDictOfType, paDFunType,
21
  paMethod, mkPR, lengthPA, replicatePA, emptyPA, packPA, combinePA, liftPA,
22
  polyAbstract, polyApply, polyVApply,
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
23
  hoistBinding, hoistExpr, hoistPolyVExpr, takeHoisted,
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
24
25
  buildClosure, buildClosures,
  mkClosureApp
26
27
28
29
) where

#include "HsVersions.h"

30
import VectCore
31
32
import VectMonad

33
import DsUtils
34
import CoreSyn
35
import CoreUtils
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
36
import Coercion
37
38
import Type
import TypeRep
39
import TyCon
40
import DataCon
41
import Var
42
43
import Id                 ( mkWildId )
import MkId               ( unwrapFamInstScrut )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
44
import Name               ( Name )
45
import PrelNames
46
import TysWiredIn
47
import TysPrim            ( intPrimTy )
48
import BasicTypes         ( Boxity(..) )
49
import Literal            ( Literal, mkMachInt )
50

51
import Outputable
52
import FastString
53

54
import Data.List             ( zipWith4 )
55
import Control.Monad         ( liftM, liftM2, zipWithM_ )
56

57
58
59
60
61
62
63
64
65
66
67
68
collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
collectAnnTypeArgs expr = go expr []
  where
    go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
    go e                             tys = (e, tys)

collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
collectAnnTypeBinders expr = go [] expr
  where
    go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
    go bs e                           = (reverse bs, e)

69
70
71
72
73
74
collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
collectAnnValBinders expr = go [] expr
  where
    go bs (_, AnnLam b e) | isId b = go (b:bs) e
    go bs e                        = (reverse bs, e)

75
76
77
78
isAnnTypeArg :: AnnExpr b ann -> Bool
isAnnTypeArg (_, AnnType t) = True
isAnnTypeArg _              = False

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
79
80
81
dataConTagZ :: DataCon -> Int
dataConTagZ con = dataConTag con - fIRST_TAG

82
mkDataConTagLit :: DataCon -> Literal
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
83
mkDataConTagLit = mkMachInt . toInteger . dataConTagZ
84

85
mkDataConTag :: DataCon -> CoreExpr
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
86
mkDataConTag = mkIntLitInt . dataConTagZ
87

88
89
90
91
92
93
94
95
splitPrimTyCon :: Type -> Maybe TyCon
splitPrimTyCon ty
  | Just (tycon, []) <- splitTyConApp_maybe ty
  , isPrimTyCon tycon
  = Just tycon

  | otherwise = Nothing

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
96
97
mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type
mkBuiltinTyConApp get_tc tys
98
  = do
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
99
100
      tc <- builtin get_tc
      return $ mkTyConApp tc tys
101

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
102
103
mkBuiltinTyConApps :: (Builtins -> TyCon) -> [Type] -> Type -> VM Type
mkBuiltinTyConApps get_tc tys ty
104
  = do
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
105
106
      tc <- builtin get_tc
      return $ foldr (mk tc) ty tys
107
  where
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
108
    mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
109

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
110
111
112
mkBuiltinTyConApps1 :: (Builtins -> TyCon) -> Type -> [Type] -> VM Type
mkBuiltinTyConApps1 get_tc dft [] = return dft
mkBuiltinTyConApps1 get_tc dft tys
113
  = do
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
114
115
116
117
118
119
120
121
122
123
124
125
126
      tc <- builtin get_tc
      case tys of
        [] -> pprPanic "mkBuiltinTyConApps1" (ppr tc)
        _  -> return $ foldr1 (mk tc) tys
  where
    mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]

mkClosureType :: Type -> Type -> VM Type
mkClosureType arg_ty res_ty = mkBuiltinTyConApp closureTyCon [arg_ty, res_ty]

mkClosureTypes :: [Type] -> Type -> VM Type
mkClosureTypes = mkBuiltinTyConApps closureTyCon

127
128
129
mkPReprType :: Type -> VM Type
mkPReprType ty = mkBuiltinTyConApp preprTyCon [ty]

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
130
131
mkPADictType :: Type -> VM Type
mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
132
133

mkPArrayType :: Type -> VM Type
134
135
136
137
138
139
mkPArrayType ty
  | Just tycon <- splitPrimTyCon ty
  = do
      arr <- traceMaybeV "mkPArrayType" (ppr tycon)
           $ lookupPrimPArray tycon
      return $ mkTyConApp arr []
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
140
mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty]
141

142
143
mkBuiltinCo :: (Builtins -> TyCon) -> VM Coercion
mkBuiltinCo get_tc
144
  = do
145
146
      tc <- builtin get_tc
      return $ mkTyConApp tc []
147

148
149
150
151
152
153
154
155
156
157
parrayReprTyCon :: Type -> VM (TyCon, [Type])
parrayReprTyCon ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])

parrayReprDataCon :: Type -> VM (DataCon, [Type])
parrayReprDataCon ty
  = do
      (tc, arg_tys) <- parrayReprTyCon ty
      let [dc] = tyConDataCons tc
      return (dc, arg_tys)

158
159
160
161
162
163
mkVScrut :: VExpr -> VM (VExpr, TyCon, [Type])
mkVScrut (ve, le)
  = do
      (tc, arg_tys) <- parrayReprTyCon (exprType ve)
      return ((ve, unwrapFamInstScrut tc arg_tys le), tc, arg_tys)

164
165
166
prDFunOfTyCon :: TyCon -> VM CoreExpr
prDFunOfTyCon tycon
  = liftM Var (traceMaybeV "prDictOfTyCon" (ppr tycon) (lookupTyConPR tycon))
167

168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
paDictArgType :: TyVar -> VM (Maybe Type)
paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
  where
    go ty k | Just k' <- kindView k = go ty k'
    go ty (FunTy k1 k2)
      = do
          tv   <- newTyVar FSLIT("a") k1
          mty1 <- go (TyVarTy tv) k1
          case mty1 of
            Just ty1 -> do
                          mty2 <- go (AppTy ty (TyVarTy tv)) k2
                          return $ fmap (ForAllTy tv . FunTy ty1) mty2
            Nothing  -> go ty k2

    go ty k
      | isLiftedTypeKind k
184
      = liftM Just (mkPADictType ty)
185
186
187

    go ty k = return Nothing

188
189
190
191
192
193
194
195
196
197
198
199
200
201
paDictOfType :: Type -> VM CoreExpr
paDictOfType ty = paDictOfTyApp ty_fn ty_args
  where
    (ty_fn, ty_args) = splitAppTys ty

paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
paDictOfTyApp ty_fn ty_args
  | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
paDictOfTyApp (TyVarTy tv) ty_args
  = do
      dfun <- maybeV (lookupTyVarPA tv)
      paDFunApply dfun ty_args
paDictOfTyApp (TyConApp tc _) ty_args
  = do
202
      dfun <- traceMaybeV "paDictOfTyApp" (ppr tc) (lookupTyConPA tc)
203
      paDFunApply (Var dfun) ty_args
204
205
paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)

206
207
208
209
210
211
212
213
214
215
216
paDFunType :: TyCon -> VM Type
paDFunType tc
  = do
      margs <- mapM paDictArgType tvs
      res   <- mkPADictType (mkTyConApp tc arg_tys)
      return . mkForAllTys tvs
             $ mkFunTys [arg | Just arg <- margs] res
  where
    tvs = tyConTyVars tc
    arg_tys = mkTyVarTys tvs

217
218
219
220
221
222
paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
paDFunApply dfun tys
  = do
      dicts <- mapM paDictOfType tys
      return $ mkApps (mkTyApps dfun tys) dicts

223
224
225
226
227
type PAMethod = (Builtins -> Var, String)

pa_length    = (lengthPAVar,    "lengthPA")
pa_replicate = (replicatePAVar, "replicatePA")
pa_empty     = (emptyPAVar,     "emptyPA")
228
pa_pack      = (packPAVar,      "packPA")
229
230
231

paMethod :: PAMethod -> Type -> VM CoreExpr
paMethod (method, name) ty
232
  | Just tycon <- splitPrimTyCon ty
233
234
235
236
237
238
  = do
      fn <- traceMaybeV "paMethod" (ppr tycon <+> text name)
          $ lookupPrimMethod tycon name
      return (Var fn)

paMethod (method, name) ty
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
239
240
241
242
243
  = do
      fn   <- builtin method
      dict <- paDictOfType ty
      return $ mkApps (Var fn) [Type ty, dict]

244
mkPR :: Type -> VM CoreExpr
245
246
247
248
249
mkPR ty
  = do
      fn   <- builtin mkPRVar
      dict <- paDictOfType ty
      return $ mkApps (Var fn) [Type ty, dict]
250

251
252
lengthPA :: Type -> CoreExpr -> VM CoreExpr
lengthPA ty x = liftM (`App` x) (paMethod pa_length ty)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
253
254
255

replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
replicatePA len x = liftM (`mkApps` [len,x])
256
                          (paMethod pa_replicate (exprType x))
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
257

258
emptyPA :: Type -> VM CoreExpr
259
emptyPA = paMethod pa_empty
260

261
262
263
264
packPA :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr
packPA ty xs len sel = liftM (`mkApps` [len, sel])
                             (paMethod pa_pack ty)

265
266
267
268
269
270
271
272
combinePA :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> [CoreExpr]
          -> VM CoreExpr
combinePA ty len sel is xs
  = liftM (`mkApps` (len : sel : is : xs))
          (paMethod (combinePAVar n, "combine" ++ show n ++ "PA") ty)
  where
    n = length xs

273
274
275
276
277
278
liftPA :: CoreExpr -> VM CoreExpr
liftPA x
  = do
      lc <- builtin liftingContext
      replicatePA (Var lc) x

279
280
281
282
283
284
285
286
newLocalVVar :: FastString -> Type -> VM VVar
newLocalVVar fs vty
  = do
      lty <- mkPArrayType vty
      vv  <- newLocalVar fs vty
      lv  <- newLocalVar fs lty
      return (vv,lv)

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
287
288
289
290
polyAbstract :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
polyAbstract tvs p
  = localV
  $ do
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
291
292
293
294
295
296
297
298
299
300
301
302
      mdicts <- mapM mk_dict_var tvs
      zipWithM_ (\tv -> maybe (defLocalTyVar tv) (defLocalTyVarWithPA tv . Var)) tvs mdicts
      p (mk_lams mdicts)
  where
    mk_dict_var tv = do
                       r <- paDictArgType tv
                       case r of
                         Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty)
                         Nothing -> return Nothing

    mk_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts])

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
303
304
polyApply :: CoreExpr -> [Type] -> VM CoreExpr
polyApply expr tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
305
306
307
308
  = do
      dicts <- mapM paDictOfType tys
      return $ expr `mkTyApps` tys `mkApps` dicts

309
310
311
312
313
314
polyVApply :: VExpr -> [Type] -> VM VExpr
polyVApply expr tys
  = do
      dicts <- mapM paDictOfType tys
      return $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
315
316
317
318
hoistBinding :: Var -> CoreExpr -> VM ()
hoistBinding v e = updGEnv $ \env ->
  env { global_bindings = (v,e) : global_bindings env }

319
320
321
322
hoistExpr :: FastString -> CoreExpr -> VM Var
hoistExpr fs expr
  = do
      var <- newLocalVar fs (exprType expr)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
323
      hoistBinding var expr
324
325
      return var

326
327
hoistVExpr :: VExpr -> VM VVar
hoistVExpr (ve, le)
328
  = do
329
      fs <- getBindName
330
331
332
      vv <- hoistExpr ('v' `consFS` fs) ve
      lv <- hoistExpr ('l' `consFS` fs) le
      return (vv, lv)
333

334
335
hoistPolyVExpr :: [TyVar] -> VM VExpr -> VM VExpr
hoistPolyVExpr tvs p
336
  = do
337
338
      expr <- closedV . polyAbstract tvs $ \abstract ->
              liftM (mapVect abstract) p
339
      fn   <- hoistVExpr expr
340
      polyVApply (vVar fn) (mkTyVarTys tvs)
341

342
343
344
345
346
347
348
takeHoisted :: VM [(Var, CoreExpr)]
takeHoisted
  = do
      env <- readGEnv id
      setGEnv $ env { global_bindings = [] }
      return $ global_bindings env

349
350
mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr
mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
351
  = do
352
353
354
355
356
      dict <- paDictOfType env_ty
      mkv  <- builtin mkClosureVar
      mkl  <- builtin mkClosurePVar
      return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
              Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
357

358
359
mkClosureApp :: Type -> Type -> VExpr -> VExpr -> VM VExpr
mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
360
361
362
363
364
365
  = do
      vapply <- builtin applyClosureVar
      lapply <- builtin applyClosurePVar
      return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
              Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [lclo, larg])

366
buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
367
368
buildClosures tvs vars [] res_ty mk_body
  = mk_body
369
370
371
buildClosures tvs vars [arg_ty] res_ty mk_body
  = buildClosure tvs vars arg_ty res_ty mk_body
buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body
372
373
374
  = do
      res_ty' <- mkClosureTypes arg_tys res_ty
      arg <- newLocalVVar FSLIT("x") arg_ty
375
      buildClosure tvs vars arg_ty res_ty'
376
        . hoistPolyVExpr tvs
377
        $ do
378
379
            lc <- builtin liftingContext
            clo <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body
380
381
            return $ vLams lc (vars ++ [arg]) clo

382
383
384
385
-- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
--   where
--     f  = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
--     f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
386
--
387
388
buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
buildClosure tvs vars arg_ty res_ty mk_body
389
  = do
390
      (env_ty, env, bind) <- buildEnv vars
391
392
      env_bndr <- newLocalVVar FSLIT("env") env_ty
      arg_bndr <- newLocalVVar FSLIT("arg") arg_ty
393

394
      fn <- hoistPolyVExpr tvs
395
          $ do
396
              lc    <- builtin liftingContext
397
398
              body  <- mk_body
              body' <- bind (vVar env_bndr)
399
                            (vVarApps lc body (vars ++ [arg_bndr]))
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
400
              return (vLamsWithoutLC [env_bndr, arg_bndr] body')
401

402
      mkClosure arg_ty res_ty env_ty fn env
403

404
405
buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VM VExpr)
buildEnv vvs
406
  = do
407
      lc <- builtin liftingContext
408
      let (ty, venv, vbind) = mkVectEnv tys vs
409
      (lenv, lbind) <- mkLiftEnv lc tys ls
410
      return (ty, (venv, lenv),
411
412
413
414
415
              \(venv,lenv) (vbody,lbody) ->
              do
                let vbody' = vbind venv vbody
                lbody' <- lbind lenv lbody
                return (vbody', lbody'))
416
  where
417
418
    (vs,ls) = unzip vvs
    tys     = map idType vs
419
420
421
422
423
424
425
426
427
428

mkVectEnv :: [Type] -> [Var] -> (Type, CoreExpr, CoreExpr -> CoreExpr -> CoreExpr)
mkVectEnv []   []  = (unitTy, Var unitDataConId, \env body -> body)
mkVectEnv [ty] [v] = (ty, Var v, \env body -> Let (NonRec v env) body)
mkVectEnv tys  vs  = (ty, mkCoreTup (map Var vs),
                        \env body -> Case env (mkWildId ty) (exprType body)
                                       [(DataAlt (tupleCon Boxed (length vs)), vs, body)])
  where
    ty = mkCoreTupTy tys

429
mkLiftEnv :: Var -> [Type] -> [Var] -> VM (CoreExpr, CoreExpr -> CoreExpr -> VM CoreExpr)
430
mkLiftEnv lc [ty] [v]
431
432
  = return (Var v, \env body ->
                   do
433
                     len <- lengthPA ty (Var v)
434
                     return . Let (NonRec v env)
435
                            $ Case len lc (exprType body) [(DEFAULT, [], body)])
436
437

-- NOTE: this transparently deals with empty environments
438
mkLiftEnv lc tys vs
439
  = do
440
      (env_tc, env_tyargs) <- parrayReprTyCon vty
441
442
443
      let [env_con] = tyConDataCons env_tc
          
          env = Var (dataConWrapId env_con)
444
                `mkTyApps`  env_tyargs
445
                `mkVarApps` (lc : vs)
446

447
448
          bind env body = let scrut = unwrapFamInstScrut env_tc env_tyargs env
                          in
449
450
                          return $ Case scrut (mkWildId (exprType scrut))
                                        (exprType body)
451
                                        [(DataAlt env_con, lc : bndrs, body)]
452
453
454
455
      return (env, bind)
  where
    vty = mkCoreTupTy tys

456
457
458
    bndrs | null vs   = [mkWildId unitTy]
          | otherwise = vs