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