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