vectorise/Vectorise/Utils/Hoisting.hs

    module Vectorise.Utils.Hoisting
      ( Inline(..)
      , addInlineArity
      , inlineMe

      , hoistBinding
      , hoistExpr
      , hoistVExpr
      , hoistPolyVExpr
      , takeHoisted
      )
    where

    import Vectorise.Monad
    import Vectorise.Env
    import Vectorise.Vect
    import Vectorise.Utils.Poly

    import CoreSyn
    import CoreUtils
    import CoreUnfold
    import Type
    import Id
    import BasicTypes  (Arity)
    import FastString
    import Control.Monad
    import Control.Applicative
    import Prelude -- avoid redundant import warning due to AMP

    -- Inline ---------------------------------------------------------------------

    -- |Records whether we should inline a particular binding.
    --
    data Inline
            = Inline Arity
            | DontInline

    -- |Add to the arity contained within an `Inline`, if any.
    --
    addInlineArity :: Inline -> Int -> Inline
    addInlineArity (Inline m) n = Inline (m+n)
    addInlineArity DontInline _ = DontInline

    -- |Says to always inline a binding.
    --
    inlineMe :: Inline
    inlineMe = Inline 0

    -- Hoisting --------------------------------------------------------------------

    hoistBinding :: Var -> CoreExpr -> VM ()
    hoistBinding v e = updGEnv $ \env ->
      env { global_bindings = (v,e) : global_bindings env }

    hoistExpr :: FastString -> CoreExpr -> Inline -> VM Var
    hoistExpr fs expr inl
      = do
          var <- mk_inline `liftM` newLocalVar fs (exprType expr)
          hoistBinding var expr
          return var
      where
        mk_inline var = case inl of
                          Inline arity -> var `setIdUnfolding`
                                          mkInlineUnfoldingWithArity arity expr
                          DontInline   -> var

    hoistVExpr :: VExpr -> Inline -> VM VVar
    hoistVExpr (ve, le) inl
      = do
          fs <- getBindName
          vv <- hoistExpr ('v' `consFS` fs) ve inl
          lv <- hoistExpr ('l' `consFS` fs) le (addInlineArity inl 1)
          return (vv, lv)

    -- |Hoist a polymorphic vectorised expression into a new top-level binding (representing a closure
    -- function).
    --
    -- The hoisted expression is parameterised by (1) a set of type variables and (2) a set of value
    -- variables that are passed as conventional type and value arguments.  The latter is implicitly
    -- extended by the set of 'PA' dictionaries required for the type variables.
    --
    hoistPolyVExpr :: [TyVar] -> [Var] -> Inline -> VM VExpr -> VM VExpr
    hoistPolyVExpr tvs vars inline p
      = do { inline' <- addInlineArity inline . (+ length vars) <$> polyArity tvs
           ; expr <- closedV . polyAbstract tvs $ \args ->
                       mapVect (mkLams $ tvs ++ args ++ vars) <$> p
           ; fn   <- hoistVExpr expr inline'
           ; let varArgs = varsToCoreExprs vars
           ; mapVect (\e -> e `mkApps` varArgs) <$> polyVApply (vVar fn) (mkTyVarTys tvs)
           }

    takeHoisted :: VM [(Var, CoreExpr)]
    takeHoisted
      = do
          env <- readGEnv id
          setGEnv $ env { global_bindings = [] }
          return $ global_bindings env