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