compiler/vectorise/VectUtils.hs

   1 module VectUtils (
   2   collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
   3   collectAnnValBinders,
   4   mkDataConTag,
   5   splitClosureTy,
   6   mkPRepr, mkToPRepr, mkToArrPRepr, mkFromPRepr, mkFromArrPRepr,
   7   mkPADictType, mkPArrayType, mkPReprType,
   8   parrayCoerce, parrayReprTyCon, parrayReprDataCon, mkVScrut,
   9   prDictOfType, prCoerce,
  10   paDictArgType, paDictOfType, paDFunType,
  11   paMethod, lengthPA, replicatePA, emptyPA, liftPA,
  12   polyAbstract, polyApply, polyVApply,
  13   hoistBinding, hoistExpr, hoistPolyVExpr, takeHoisted,
  14   buildClosure, buildClosures,
  15   mkClosureApp
  16 ) where
  17
  18 #include "HsVersions.h"
  19
  20 import VectCore
  21 import VectMonad
  22
  23 import DsUtils
  24 import CoreSyn
  25 import CoreUtils
  26 import Coercion
  27 import Type
  28 import TypeRep
  29 import TyCon
  30 import DataCon            ( DataCon, dataConWrapId, dataConTag )
  31 import Var
  32 import Id                 ( mkWildId )
  33 import MkId               ( unwrapFamInstScrut )
  34 import Name               ( Name )
  35 import PrelNames
  36 import TysWiredIn
  37 import BasicTypes         ( Boxity(..) )
  38
  39 import Outputable
  40 import FastString
  41
  42 import Data.List             ( zipWith4 )
  43 import Control.Monad         ( liftM, liftM2, zipWithM_ )
  44
  45 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
  46 collectAnnTypeArgs expr = go expr []
  47   where
  48     go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
  49     go e                             tys = (e, tys)
  50
  51 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  52 collectAnnTypeBinders expr = go [] expr
  53   where
  54     go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
  55     go bs e                           = (reverse bs, e)
  56
  57 collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  58 collectAnnValBinders expr = go [] expr
  59   where
  60     go bs (_, AnnLam b e) | isId b = go (b:bs) e
  61     go bs e                        = (reverse bs, e)
  62
  63 isAnnTypeArg :: AnnExpr b ann -> Bool
  64 isAnnTypeArg (_, AnnType t) = True
  65 isAnnTypeArg _              = False
  66
  67 mkDataConTag :: DataCon -> CoreExpr
  68 mkDataConTag dc = mkConApp intDataCon [mkIntLitInt $ dataConTag dc]
  69
  70 splitUnTy :: String -> Name -> Type -> Type
  71 splitUnTy s name ty
  72   | Just (tc, [ty']) <- splitTyConApp_maybe ty
  73   , tyConName tc == name
  74   = ty'
  75
  76   | otherwise = pprPanic s (ppr ty)
  77
  78 splitBinTy :: String -> Name -> Type -> (Type, Type)
  79 splitBinTy s name ty
  80   | Just (tc, [ty1, ty2]) <- splitTyConApp_maybe ty
  81   , tyConName tc == name
  82   = (ty1, ty2)
  83
  84   | otherwise = pprPanic s (ppr ty)
  85
  86 splitFixedTyConApp :: TyCon -> Type -> [Type]
  87 splitFixedTyConApp tc ty
  88   | Just (tc', tys) <- splitTyConApp_maybe ty
  89   , tc == tc'
  90   = tys
  91
  92   | otherwise = pprPanic "splitFixedTyConApp" (ppr tc <+> ppr ty)
  93
  94 splitEmbedTy :: Type -> Type
  95 splitEmbedTy = splitUnTy "splitEmbedTy" embedTyConName
  96
  97 splitClosureTy :: Type -> (Type, Type)
  98 splitClosureTy = splitBinTy "splitClosureTy" closureTyConName
  99
 100 splitPArrayTy :: Type -> Type
 101 splitPArrayTy = splitUnTy "splitPArrayTy" parrayTyConName
 102
 103 mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type
 104 mkBuiltinTyConApp get_tc tys
 105   = do
 106       tc <- builtin get_tc
 107       return $ mkTyConApp tc tys
 108
 109 mkBuiltinTyConApps :: (Builtins -> TyCon) -> [Type] -> Type -> VM Type
 110 mkBuiltinTyConApps get_tc tys ty
 111   = do
 112       tc <- builtin get_tc
 113       return $ foldr (mk tc) ty tys
 114   where
 115     mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
 116
 117 mkBuiltinTyConApps1 :: (Builtins -> TyCon) -> Type -> [Type] -> VM Type
 118 mkBuiltinTyConApps1 get_tc dft [] = return dft
 119 mkBuiltinTyConApps1 get_tc dft tys
 120   = do
 121       tc <- builtin get_tc
 122       case tys of
 123         [] -> pprPanic "mkBuiltinTyConApps1" (ppr tc)
 124         _  -> return $ foldr1 (mk tc) tys
 125   where
 126     mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
 127
 128 mkPRepr :: [[Type]] -> VM Type
 129 mkPRepr tys
 130   = do
 131       embed_tc <- builtin embedTyCon
 132       sum_tcs  <- builtins sumTyCon
 133       prod_tcs <- builtins prodTyCon
 134
 135       let mk_sum []   = unitTy
 136           mk_sum [ty] = ty
 137           mk_sum tys  = mkTyConApp (sum_tcs $ length tys) tys
 138
 139           mk_prod []   = unitTy
 140           mk_prod [ty] = ty
 141           mk_prod tys  = mkTyConApp (prod_tcs $ length tys) tys
 142
 143           mk_embed ty = mkTyConApp embed_tc [ty]
 144
 145       return . mk_sum
 146              . map (mk_prod . map mk_embed)
 147              $ tys
 148
 149 mkToPRepr :: [[CoreExpr]] -> VM ([CoreExpr], Type)
 150 mkToPRepr ess
 151   = do
 152       embed_tc <- builtin embedTyCon
 153       embed_dc <- builtin embedDataCon
 154       sum_tcs  <- builtins sumTyCon
 155       prod_tcs <- builtins prodTyCon
 156
 157       let mk_sum [] = ([Var unitDataConId], unitTy)
 158           mk_sum [(expr, ty)] = ([expr], ty)
 159           mk_sum es = (zipWith mk_alt (tyConDataCons sum_tc) exprs,
 160                        mkTyConApp sum_tc tys)
 161             where
 162               (exprs, tys)   = unzip es
 163               sum_tc         = sum_tcs (length es)
 164               mk_alt dc expr = mkConApp dc (map Type tys ++ [expr])
 165
 166           mk_prod [] = (Var unitDataConId, unitTy)
 167           mk_prod [(expr, ty)] = (expr, ty)
 168           mk_prod es = (mkConApp prod_dc (map Type tys ++ exprs),
 169                         mkTyConApp prod_tc tys)
 170             where
 171               (exprs, tys) = unzip es
 172               prod_tc      = prod_tcs (length es)
 173               [prod_dc]    = tyConDataCons prod_tc
 174
 175           mk_embed expr = (mkConApp embed_dc [Type ty, expr],
 176                            mkTyConApp embed_tc [ty])
 177             where ty = exprType expr
 178
 179       return . mk_sum $ map (mk_prod . map mk_embed) ess
 180
 181 mkToArrPRepr :: CoreExpr -> CoreExpr -> [[CoreExpr]] -> VM CoreExpr
 182 mkToArrPRepr len sel ess
 183   = do
 184       embed_tc <- builtin embedTyCon
 185       (embed_rtc, _) <- parrayReprTyCon (mkTyConApp embed_tc [unitTy])
 186       let [embed_rdc] = tyConDataCons embed_rtc
 187
 188       let mk_sum [(expr, ty)] = return (expr, ty)
 189           mk_sum es
 190             = do
 191                 sum_tc <- builtin . sumTyCon $ length es
 192                 (sum_rtc, _) <- parrayReprTyCon (mkTyConApp sum_tc tys)
 193                 let [sum_rdc] = tyConDataCons sum_rtc
 194
 195                 return (mkConApp sum_rdc (map Type tys ++ (len : sel : exprs)),
 196                         mkTyConApp sum_tc tys)
 197             where
 198               (exprs, tys) = unzip es
 199
 200           mk_prod [(expr, ty)] = return (expr, ty)
 201           mk_prod es
 202             = do
 203                 prod_tc <- builtin . prodTyCon $ length es
 204                 (prod_rtc, _) <- parrayReprTyCon (mkTyConApp prod_tc tys)
 205                 let [prod_rdc] = tyConDataCons prod_rtc
 206
 207                 return (mkConApp prod_rdc (map Type tys ++ (len : exprs)),
 208                         mkTyConApp prod_tc tys)
 209             where
 210               (exprs, tys) = unzip es
 211
 212           mk_embed expr = (mkConApp embed_rdc [Type ty, expr],
 213                            mkTyConApp embed_tc [ty])
 214             where ty = splitPArrayTy (exprType expr)
 215
 216       liftM fst (mk_sum =<< mapM (mk_prod . map mk_embed) ess)
 217
 218 mkFromPRepr :: CoreExpr -> Type -> [([Var], CoreExpr)] -> VM CoreExpr
 219 mkFromPRepr scrut res_ty alts
 220   = do
 221       embed_dc <- builtin embedDataCon
 222       sum_tcs  <- builtins sumTyCon
 223       prod_tcs <- builtins prodTyCon
 224
 225       let un_sum expr ty [(vars, res)] = un_prod expr ty vars res
 226           un_sum expr ty bs
 227             = do
 228                 ps     <- mapM (newLocalVar FSLIT("p")) tys
 229                 bodies <- sequence
 230                         $ zipWith4 un_prod (map Var ps) tys vars rs
 231                 return . Case expr (mkWildId ty) res_ty
 232                        $ zipWith3 mk_alt sum_dcs ps bodies
 233             where
 234               (vars, rs) = unzip bs
 235               tys        = splitFixedTyConApp sum_tc ty
 236               sum_tc     = sum_tcs $ length bs
 237               sum_dcs    = tyConDataCons sum_tc
 238
 239               mk_alt dc p body = (DataAlt dc, [p], body)
 240
 241           un_prod expr ty []    r = return r
 242           un_prod expr ty [var] r = return $ un_embed expr ty var r
 243           un_prod expr ty vars  r
 244             = do
 245                 xs <- mapM (newLocalVar FSLIT("x")) tys
 246                 let body = foldr (\(e,t,v) r -> un_embed e t v r) r
 247                          $ zip3 (map Var xs) tys vars
 248                 return $ Case expr (mkWildId ty) res_ty
 249                          [(DataAlt prod_dc, xs, body)]
 250             where
 251               tys       = splitFixedTyConApp prod_tc ty
 252               prod_tc   = prod_tcs $ length vars
 253               [prod_dc] = tyConDataCons prod_tc
 254
 255           un_embed expr ty var r
 256             = Case expr (mkWildId ty) res_ty
 257                 [(DataAlt embed_dc, [var], r)]
 258
 259       un_sum scrut (exprType scrut) alts
 260
 261 mkFromArrPRepr :: CoreExpr -> Type -> Var -> Var -> [[Var]] -> CoreExpr
 262                -> VM CoreExpr
 263 mkFromArrPRepr scrut res_ty len sel vars res
 264   = return (Var unitDataConId)
 265
 266 mkClosureType :: Type -> Type -> VM Type
 267 mkClosureType arg_ty res_ty = mkBuiltinTyConApp closureTyCon [arg_ty, res_ty]
 268
 269 mkClosureTypes :: [Type] -> Type -> VM Type
 270 mkClosureTypes = mkBuiltinTyConApps closureTyCon
 271
 272 mkPReprType :: Type -> VM Type
 273 mkPReprType ty = mkBuiltinTyConApp preprTyCon [ty]
 274
 275 mkPADictType :: Type -> VM Type
 276 mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
 277
 278 mkPArrayType :: Type -> VM Type
 279 mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty]
 280
 281 parrayCoerce :: TyCon -> [Type] -> CoreExpr -> VM CoreExpr
 282 parrayCoerce repr_tc args expr
 283   | Just arg_co <- tyConFamilyCoercion_maybe repr_tc
 284   = do
 285       parray <- builtin parrayTyCon
 286
 287       let co = mkAppCoercion (mkTyConApp parray [])
 288                              (mkSymCoercion (mkTyConApp arg_co args))
 289
 290       return $ mkCoerce co expr
 291
 292 parrayReprTyCon :: Type -> VM (TyCon, [Type])
 293 parrayReprTyCon ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
 294
 295 parrayReprDataCon :: Type -> VM (DataCon, [Type])
 296 parrayReprDataCon ty
 297   = do
 298       (tc, arg_tys) <- parrayReprTyCon ty
 299       let [dc] = tyConDataCons tc
 300       return (dc, arg_tys)
 301
 302 mkVScrut :: VExpr -> VM (VExpr, TyCon, [Type])
 303 mkVScrut (ve, le)
 304   = do
 305       (tc, arg_tys) <- parrayReprTyCon (exprType ve)
 306       return ((ve, unwrapFamInstScrut tc arg_tys le), tc, arg_tys)
 307
 308 prDictOfType :: Type -> VM CoreExpr
 309 prDictOfType orig_ty
 310   | Just (tycon, ty_args) <- splitTyConApp_maybe orig_ty
 311   = do
 312       dfun <- traceMaybeV "prDictOfType" (ppr tycon) (lookupTyConPR tycon)
 313       prDFunApply (Var dfun) ty_args
 314
 315 prDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
 316 prDFunApply dfun tys
 317   = do
 318       args <- mapM mkDFunArg arg_tys
 319       return $ mkApps mono_dfun args
 320   where
 321     mono_dfun    = mkTyApps dfun tys
 322     (arg_tys, _) = splitFunTys (exprType mono_dfun)
 323
 324 mkDFunArg :: Type -> VM CoreExpr
 325 mkDFunArg ty
 326   | Just (tycon, [arg]) <- splitTyConApp_maybe ty
 327
 328   = let name = tyConName tycon
 329
 330         get_dict | name == paTyConName = paDictOfType
 331                  | name == prTyConName = prDictOfType
 332                  | otherwise           = pprPanic "mkDFunArg" (ppr ty)
 333
 334     in get_dict arg
 335
 336 mkDFunArg ty = pprPanic "mkDFunArg" (ppr ty)
 337
 338 prCoerce :: TyCon -> [Type] -> CoreExpr -> VM CoreExpr
 339 prCoerce repr_tc args expr
 340   | Just arg_co <- tyConFamilyCoercion_maybe repr_tc
 341   = do
 342       pr_tc <- builtin prTyCon
 343
 344       let co = mkAppCoercion (mkTyConApp pr_tc [])
 345                              (mkSymCoercion (mkTyConApp arg_co args))
 346
 347       return $ mkCoerce co expr
 348
 349 paDictArgType :: TyVar -> VM (Maybe Type)
 350 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
 351   where
 352     go ty k | Just k' <- kindView k = go ty k'
 353     go ty (FunTy k1 k2)
 354       = do
 355           tv   <- newTyVar FSLIT("a") k1
 356           mty1 <- go (TyVarTy tv) k1
 357           case mty1 of
 358             Just ty1 -> do
 359                           mty2 <- go (AppTy ty (TyVarTy tv)) k2
 360                           return $ fmap (ForAllTy tv . FunTy ty1) mty2
 361             Nothing  -> go ty k2
 362
 363     go ty k
 364       | isLiftedTypeKind k
 365       = liftM Just (mkPADictType ty)
 366
 367     go ty k = return Nothing
 368
 369 paDictOfType :: Type -> VM CoreExpr
 370 paDictOfType ty = paDictOfTyApp ty_fn ty_args
 371   where
 372     (ty_fn, ty_args) = splitAppTys ty
 373
 374 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
 375 paDictOfTyApp ty_fn ty_args
 376   | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
 377 paDictOfTyApp (TyVarTy tv) ty_args
 378   = do
 379       dfun <- maybeV (lookupTyVarPA tv)
 380       paDFunApply dfun ty_args
 381 paDictOfTyApp (TyConApp tc _) ty_args
 382   = do
 383       dfun <- traceMaybeV "paDictOfTyApp" (ppr tc) (lookupTyConPA tc)
 384       paDFunApply (Var dfun) ty_args
 385 paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
 386
 387 paDFunType :: TyCon -> VM Type
 388 paDFunType tc
 389   = do
 390       margs <- mapM paDictArgType tvs
 391       res   <- mkPADictType (mkTyConApp tc arg_tys)
 392       return . mkForAllTys tvs
 393              $ mkFunTys [arg | Just arg <- margs] res
 394   where
 395     tvs = tyConTyVars tc
 396     arg_tys = mkTyVarTys tvs
 397
 398 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
 399 paDFunApply dfun tys
 400   = do
 401       dicts <- mapM paDictOfType tys
 402       return $ mkApps (mkTyApps dfun tys) dicts
 403
 404 paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr
 405 paMethod method ty
 406   = do
 407       fn   <- builtin method
 408       dict <- paDictOfType ty
 409       return $ mkApps (Var fn) [Type ty, dict]
 410
 411 lengthPA :: CoreExpr -> VM CoreExpr
 412 lengthPA x = liftM (`App` x) (paMethod lengthPAVar ty)
 413   where
 414     ty = splitPArrayTy (exprType x)
 415
 416 replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
 417 replicatePA len x = liftM (`mkApps` [len,x])
 418                           (paMethod replicatePAVar (exprType x))
 419
 420 emptyPA :: Type -> VM CoreExpr
 421 emptyPA = paMethod emptyPAVar
 422
 423 liftPA :: CoreExpr -> VM CoreExpr
 424 liftPA x
 425   = do
 426       lc <- builtin liftingContext
 427       replicatePA (Var lc) x
 428
 429 newLocalVVar :: FastString -> Type -> VM VVar
 430 newLocalVVar fs vty
 431   = do
 432       lty <- mkPArrayType vty
 433       vv  <- newLocalVar fs vty
 434       lv  <- newLocalVar fs lty
 435       return (vv,lv)
 436
 437 polyAbstract :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
 438 polyAbstract tvs p
 439   = localV
 440   $ do
 441       mdicts <- mapM mk_dict_var tvs
 442       zipWithM_ (\tv -> maybe (defLocalTyVar tv) (defLocalTyVarWithPA tv . Var)) tvs mdicts
 443       p (mk_lams mdicts)
 444   where
 445     mk_dict_var tv = do
 446                        r <- paDictArgType tv
 447                        case r of
 448                          Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty)
 449                          Nothing -> return Nothing
 450
 451     mk_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts])
 452
 453 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
 454 polyApply expr tys
 455   = do
 456       dicts <- mapM paDictOfType tys
 457       return $ expr `mkTyApps` tys `mkApps` dicts
 458
 459 polyVApply :: VExpr -> [Type] -> VM VExpr
 460 polyVApply expr tys
 461   = do
 462       dicts <- mapM paDictOfType tys
 463       return $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr
 464
 465 hoistBinding :: Var -> CoreExpr -> VM ()
 466 hoistBinding v e = updGEnv $ \env ->
 467   env { global_bindings = (v,e) : global_bindings env }
 468
 469 hoistExpr :: FastString -> CoreExpr -> VM Var
 470 hoistExpr fs expr
 471   = do
 472       var <- newLocalVar fs (exprType expr)
 473       hoistBinding var expr
 474       return var
 475
 476 hoistVExpr :: VExpr -> VM VVar
 477 hoistVExpr (ve, le)
 478   = do
 479       fs <- getBindName
 480       vv <- hoistExpr ('v' `consFS` fs) ve
 481       lv <- hoistExpr ('l' `consFS` fs) le
 482       return (vv, lv)
 483
 484 hoistPolyVExpr :: [TyVar] -> VM VExpr -> VM VExpr
 485 hoistPolyVExpr tvs p
 486   = do
 487       expr <- closedV . polyAbstract tvs $ \abstract ->
 488               liftM (mapVect abstract) p
 489       fn   <- hoistVExpr expr
 490       polyVApply (vVar fn) (mkTyVarTys tvs)
 491
 492 takeHoisted :: VM [(Var, CoreExpr)]
 493 takeHoisted
 494   = do
 495       env <- readGEnv id
 496       setGEnv $ env { global_bindings = [] }
 497       return $ global_bindings env
 498
 499 mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr
 500 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
 501   = do
 502       dict <- paDictOfType env_ty
 503       mkv  <- builtin mkClosureVar
 504       mkl  <- builtin mkClosurePVar
 505       return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
 506               Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
 507
 508 mkClosureApp :: VExpr -> VExpr -> VM VExpr
 509 mkClosureApp (vclo, lclo) (varg, larg)
 510   = do
 511       vapply <- builtin applyClosureVar
 512       lapply <- builtin applyClosurePVar
 513       return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
 514               Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [lclo, larg])
 515   where
 516     (arg_ty, res_ty) = splitClosureTy (exprType vclo)
 517
 518 buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
 519 buildClosures tvs vars [] res_ty mk_body
 520   = mk_body
 521 buildClosures tvs vars [arg_ty] res_ty mk_body
 522   = buildClosure tvs vars arg_ty res_ty mk_body
 523 buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body
 524   = do
 525       res_ty' <- mkClosureTypes arg_tys res_ty
 526       arg <- newLocalVVar FSLIT("x") arg_ty
 527       buildClosure tvs vars arg_ty res_ty'
 528         . hoistPolyVExpr tvs
 529         $ do
 530             lc <- builtin liftingContext
 531             clo <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body
 532             return $ vLams lc (vars ++ [arg]) clo
 533
 534 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
 535 --   where
 536 --     f  = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
 537 --     f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
 538 --
 539 buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
 540 buildClosure tvs vars arg_ty res_ty mk_body
 541   = do
 542       (env_ty, env, bind) <- buildEnv vars
 543       env_bndr <- newLocalVVar FSLIT("env") env_ty
 544       arg_bndr <- newLocalVVar FSLIT("arg") arg_ty
 545
 546       fn <- hoistPolyVExpr tvs
 547           $ do
 548               lc    <- builtin liftingContext
 549               body  <- mk_body
 550               body' <- bind (vVar env_bndr)
 551                             (vVarApps lc body (vars ++ [arg_bndr]))
 552               return (vLamsWithoutLC [env_bndr, arg_bndr] body')
 553
 554       mkClosure arg_ty res_ty env_ty fn env
 555
 556 buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VM VExpr)
 557 buildEnv vvs
 558   = do
 559       lc <- builtin liftingContext
 560       let (ty, venv, vbind) = mkVectEnv tys vs
 561       (lenv, lbind) <- mkLiftEnv lc tys ls
 562       return (ty, (venv, lenv),
 563               \(venv,lenv) (vbody,lbody) ->
 564               do
 565                 let vbody' = vbind venv vbody
 566                 lbody' <- lbind lenv lbody
 567                 return (vbody', lbody'))
 568   where
 569     (vs,ls) = unzip vvs
 570     tys     = map idType vs
 571
 572 mkVectEnv :: [Type] -> [Var] -> (Type, CoreExpr, CoreExpr -> CoreExpr -> CoreExpr)
 573 mkVectEnv []   []  = (unitTy, Var unitDataConId, \env body -> body)
 574 mkVectEnv [ty] [v] = (ty, Var v, \env body -> Let (NonRec v env) body)
 575 mkVectEnv tys  vs  = (ty, mkCoreTup (map Var vs),
 576                         \env body -> Case env (mkWildId ty) (exprType body)
 577                                        [(DataAlt (tupleCon Boxed (length vs)), vs, body)])
 578   where
 579     ty = mkCoreTupTy tys
 580
 581 mkLiftEnv :: Var -> [Type] -> [Var] -> VM (CoreExpr, CoreExpr -> CoreExpr -> VM CoreExpr)
 582 mkLiftEnv lc [ty] [v]
 583   = return (Var v, \env body ->
 584                    do
 585                      len <- lengthPA (Var v)
 586                      return . Let (NonRec v env)
 587                             $ Case len lc (exprType body) [(DEFAULT, [], body)])
 588
 589 -- NOTE: this transparently deals with empty environments
 590 mkLiftEnv lc tys vs
 591   = do
 592       (env_tc, env_tyargs) <- parrayReprTyCon vty
 593       let [env_con] = tyConDataCons env_tc
 594
 595           env = Var (dataConWrapId env_con)
 596                 `mkTyApps`  env_tyargs
 597                 `mkVarApps` (lc : vs)
 598
 599           bind env body = let scrut = unwrapFamInstScrut env_tc env_tyargs env
 600                           in
 601                           return $ Case scrut (mkWildId (exprType scrut))
 602                                         (exprType body)
 603                                         [(DataAlt env_con, lc : bndrs, body)]
 604       return (env, bind)
 605   where
 606     vty = mkCoreTupTy tys
 607
 608     bndrs | null vs   = [mkWildId unitTy]
 609           | otherwise = vs
 610