compiler/vectorise/VectUtils.hs

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