compiler/vectorise/Vectorise/Utils/Closure.hs

   1
   2 -- | Utils concerning closure construction and application.
   3 module Vectorise.Utils.Closure (
   4         mkClosure,
   5         mkClosureApp,
   6         buildClosure,
   7         buildClosures,
   8         buildEnv
   9 )
  10 where
  11 import Vectorise.Builtins
  12 import Vectorise.Vect
  13 import Vectorise.Monad
  14 import Vectorise.Utils.Base
  15 import Vectorise.Utils.PADict
  16 import Vectorise.Utils.Hoisting
  17
  18 import CoreSyn
  19 import Type
  20 import MkCore
  21 import CoreUtils
  22 import TyCon
  23 import DataCon
  24 import MkId
  25 import TysWiredIn
  26 import BasicTypes( Boxity(..) )
  27 import FastString
  28
  29
  30 -- | Make a closure.
  31 mkClosure
  32         :: Type         -- ^ Type of the argument.
  33         -> Type         -- ^ Type of the result.
  34         -> Type         -- ^ Type of the environment.
  35         -> VExpr        -- ^ The function to apply.
  36         -> VExpr        -- ^ The environment to use.
  37         -> VM VExpr
  38
  39 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
  40  = do dict <- paDictOfType env_ty
  41       mkv  <- builtin closureVar
  42       mkl  <- builtin liftedClosureVar
  43       return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
  44               Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
  45
  46
  47 -- | Make a closure application.
  48 mkClosureApp
  49         :: Type         -- ^ Type of the argument.
  50         -> Type         -- ^ Type of the result.
  51         -> VExpr        -- ^ Closure to apply.
  52         -> VExpr        -- ^ Argument to use.
  53         -> VM VExpr
  54
  55 mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg)
  56  = do vapply <- builtin applyVar
  57       lapply <- builtin liftedApplyVar
  58       lc     <- builtin liftingContext
  59       return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
  60               Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [Var lc, lclo, larg])
  61
  62
  63
  64 buildClosures
  65         :: [TyVar]
  66         -> [VVar]
  67         -> [Type]       -- ^ Type of the arguments.
  68         -> Type         -- ^ Type of result.
  69         -> VM VExpr
  70         -> VM VExpr
  71
  72 buildClosures _   _    [] _ mk_body
  73  = mk_body
  74
  75 buildClosures tvs vars [arg_ty] res_ty mk_body
  76  =  buildClosure tvs vars arg_ty res_ty mk_body
  77
  78 buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body
  79  = do res_ty' <- mkClosureTypes arg_tys res_ty
  80       arg     <- newLocalVVar (fsLit "x") arg_ty
  81       buildClosure tvs vars arg_ty res_ty'
  82         . hoistPolyVExpr tvs (Inline (length vars + 1))
  83         $ do
  84             lc     <- builtin liftingContext
  85             clo    <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body
  86             return $ vLams lc (vars ++ [arg]) clo
  87
  88
  89 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
  90 --   where
  91 --     f  = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
  92 --     f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
  93 --
  94 buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
  95 buildClosure tvs vars arg_ty res_ty mk_body
  96   = do
  97       (env_ty, env, bind) <- buildEnv vars
  98       env_bndr <- newLocalVVar (fsLit "env") env_ty
  99       arg_bndr <- newLocalVVar (fsLit "arg") arg_ty
 100
 101       fn <- hoistPolyVExpr tvs (Inline 2)
 102           $ do
 103               lc     <- builtin liftingContext
 104               body   <- mk_body
 105               return .  vLams lc [env_bndr, arg_bndr]
 106                      $  bind (vVar env_bndr)
 107                              (vVarApps lc body (vars ++ [arg_bndr]))
 108
 109       mkClosure arg_ty res_ty env_ty fn env
 110
 111
 112 -- Environments ---------------------------------------------------------------
 113 buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr)
 114 buildEnv []
 115  = do
 116       ty    <- voidType
 117       void  <- builtin voidVar
 118       pvoid <- builtin pvoidVar
 119       return (ty, vVar (void, pvoid), \_ body -> body)
 120
 121 buildEnv [v] = return (vVarType v, vVar v,
 122                     \env body -> vLet (vNonRec v env) body)
 123
 124 buildEnv vs
 125  = do (lenv_tc, lenv_tyargs) <- pdataReprTyCon ty
 126
 127       let venv_con   = tupleCon Boxed (length vs)
 128           [lenv_con] = tyConDataCons lenv_tc
 129
 130           venv       = mkCoreTup (map Var vvs)
 131           lenv       = Var (dataConWrapId lenv_con)
 132                        `mkTyApps` lenv_tyargs
 133                        `mkApps`   map Var lvs
 134
 135           vbind env body = mkWildCase env ty (exprType body)
 136                            [(DataAlt venv_con, vvs, body)]
 137
 138           lbind env body =
 139             let scrut = unwrapFamInstScrut lenv_tc lenv_tyargs env
 140             in
 141             mkWildCase scrut (exprType scrut) (exprType body)
 142               [(DataAlt lenv_con, lvs, body)]
 143
 144           bind (venv, lenv) (vbody, lbody) = (vbind venv vbody,
 145                                               lbind lenv lbody)
 146
 147       return (ty, (venv, lenv), bind)
 148   where
 149     (vvs, lvs) = unzip vs
 150     tys        = map vVarType vs
 151     ty         = mkBoxedTupleTy tys