path: root/src/lib/Inline.hs

-- Copyright 2022 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

module Inline (inlineBindings) where

import Builder
import Core
import Err
import CheapReduction
import IRVariants
import Name
import Subst
import Occurrence hiding (Var)
import PeepholeOptimize
import Types.Core
import Types.Primitives

-- === External API ===

inlineBindings :: (EnvReader m) => STopLam n -> m n (STopLam n)
inlineBindings lam = liftLamExpr lam \body -> liftInlineM $ buildBlock $ inlineExpr Stop body
{-# INLINE inlineBindings #-}
{-# SCC inlineBindings #-}

-- === Data Structure ===

data InlineExpr (r::IR) (o::S) where
  DoneEx :: SExpr o -> InlineExpr SimpIR o
  SuspEx :: SExpr i -> Subst InlineSubstVal i o -> InlineExpr SimpIR o

instance Show (InlineExpr r n) where
  show = \case
    DoneEx e -> "finished " ++ show e
    SuspEx e _ -> "unfinished " ++ show e

instance RenameE (InlineExpr r) where
  renameE (scope, subst) = \case
    DoneEx e -> DoneEx $ renameE (scope, subst) e
    SuspEx e s -> SuspEx e $ renameE (scope, subst) s

instance SinkableE (InlineExpr r) where
  sinkingProofE rename = \case
    DoneEx e   -> DoneEx $ sinkingProofE rename e
    SuspEx e s -> SuspEx e $ sinkingProofE rename s

type InlineSubstVal = SubstVal InlineExpr

newtype InlineM (i::S) (o::S) (a:: *) = InlineM
  { runInlineM :: SubstReaderT InlineSubstVal (BuilderM SimpIR) i o a }
  deriving ( Functor, Applicative, Monad, MonadFail, Fallible, ScopeReader
           , EnvExtender, EnvReader, SubstReader InlineSubstVal, (Builder SimpIR)
           , (ScopableBuilder SimpIR))

liftInlineM :: (EnvReader m) => InlineM n n a -> m n a
liftInlineM act = liftBuilder $ runSubstReaderT idSubst $ runInlineM act
{-# INLINE liftInlineM #-}

-- === Inliner ===

data SizePreservationInfo =
  -- Explicit noinline, or inlining doesn't preserve work
    NoInline
  -- Used once statically, ergo size-preserving to inline.  In Secrets, this
  -- corresponds to either OnceSafe, or OnceUnsafe with whnfOrBot == True
  | UsedOnce
  -- Used more than once statically, ergo potentially size-increasing to inline.
  -- In Secrets, this corresponds to either MultiSafe, or MultiUnsafe with
  -- whnfOrBot == True
  | UsedMulti
  deriving (Eq, Show)

inlineDeclsSubst :: Emits o => Nest SDecl i i' -> InlineM i o (Subst InlineSubstVal i' o)
inlineDeclsSubst = \case
  Empty -> getSubst
  Nest (Let b (DeclBinding ann expr)) rest -> do
    if preInlineUnconditionally ann then do
      s <- getSubst
      extendSubst (b @> SubstVal (SuspEx expr s)) $ inlineDeclsSubst rest
    else do
      expr' <- peepholeExpr <$> inlineExpr Stop expr
      -- If the inliner starts moving effectful expressions, it may become
      -- necessary to query the effects of the new expression here.
      let presInfo = resolveWorkConservation ann expr'
      -- See NoteSecretsSubtlety
      if presInfo == UsedOnce then do
        let substVal = case expr' of
             Atom (Stuck _ (Var name')) -> Rename $ atomVarName name'
             _ -> SubstVal (DoneEx expr')
        extendSubst (b @> substVal) $ inlineDeclsSubst rest
      else do
        -- expr' can't be Atom (Var x) here
        name' <- emitDecl (getNameHint b) (dropOccInfo ann) expr'
        extendSubst (b @> Rename (atomVarName name')) do
          -- See NoteConditionalInlining
          inlineDeclsSubst rest
  where
    dropOccInfo PlainLet = PlainLet
    dropOccInfo LinearLet = LinearLet
    dropOccInfo InlineLet = InlineLet
    dropOccInfo NoInlineLet = NoInlineLet
    dropOccInfo (OccInfoPure _) = PlainLet
    dropOccInfo (OccInfoImpure _) = PlainLet
    resolveWorkConservation PlainLet _ =
      NoInline  -- No occurrence info, assume the worst
    resolveWorkConservation LinearLet _ = NoInline
    resolveWorkConservation InlineLet _ = NoInline
    resolveWorkConservation NoInlineLet _ = NoInline
    -- Quick hack to always unconditionally inline renames, until we get
    -- a better story about measuring the sizes of atoms and expressions.
    resolveWorkConservation (OccInfoPure _) (Atom (Stuck _ (Var _))) = UsedOnce
    resolveWorkConservation (OccInfoPure (UsageInfo s (ixDepth, d))) expr
      | d <= One = case ixDepthExpr expr >= ixDepth of
        True -> if s <= One then UsedOnce else UsedMulti
        False -> NoInline
    resolveWorkConservation (OccInfoPure (UsageInfo s _)) (Atom _) =
      if s <= One then UsedOnce else UsedMulti
    -- TODO In Haskell, inlining expressions that are guaranteed to be bottom is
    -- also work-preserving, and profitable because it can avoid allocating
    -- thunks.  Do we care about that case here?  (It can also change the
    -- semantics in a strict language, from "error" to "no error" (or to
    -- "different error").)
    resolveWorkConservation (OccInfoPure _) _ = NoInline
    -- TODO Tagging impure expressions "noinline" here misses two potential
    -- opportunities.
    -- - The OccInfo annotation is from the target expression before inlining.
    --   It's conceivable that inlining will expose a peephole optimization that
    --   removes the effect, making the post-inlining expression pure.  It would
    --   be nice to be able to inline it here in this case, but (i) as of this
    --   writing, the inliner does not remove effects, and (ii) even if it did,
    --   we could recover the end state by running another inlining pass.
    -- - More interestingly, the inliner could reorder provably-independent
    --   effects, like `State` on two different heap parameters.  For that,
    --   however, we would need to know something about the locations into which
    --   an effectful binding could be moved, and whether doing so reorders
    --   effects that should not be reordered.  And because the inliner inlines
    --   multiple bindings in one pass, we may also need to be careful about any
    --   effectful bindings that are already in the substitution.
    resolveWorkConservation (OccInfoImpure _) _ = NoInline
    -- This function for detecting available indexing depth reports all decls as
    -- aborting indexing, even if the decl will be inlined away later.
    --
    -- For example, the binding
    --   xs = for i.
    --     ys = for j. <something>
    --     for k. ys.k
    -- really does have two available indexing levels, because the ys binding is
    -- inlinable into the result of the outer `for`; but the below function will
    -- not detect it, and thus we will decline to inline `xs` into a context
    -- where two levels of indexing are required.  However, we can recover the
    -- desired effect by running a second pass of occurrence analysis and
    -- inlining.
    --
    -- A more subtle case occurs with
    --   xs = for i.
    --     ys = for j. <something>
    --     view k. ys.k
    -- (note the `view`).  In this case, if we went ahead and inlined `xs` into
    -- a context that required two levels of indexing, we would temporarily
    -- duplicate work, but a second pass of occurence analysis and inlining
    -- would fix it by inlining `ys`.  However, if we decline to inline here,
    -- running inlining on the body of `for i.` will not inline into the `view`,
    -- because in general views are supposed to be duplicable all over.  Maybe
    -- the solution is to just not have view atoms in the IR by this point,
    -- since their main purpose is to force inlining in the simplifier, and if
    -- one just stuck like this it has become equivalent to a `for` anyway.
    ixDepthExpr :: Expr SimpIR n -> Int
    ixDepthExpr (PrimOp (Hof (TypedHof _ (For _ _ (UnaryLamExpr _ body))))) = 1 + ixDepthExpr body
    ixDepthExpr _ = 0

-- Should we decide to inline this binding wherever it appears, before we even
-- know the expression?  "Yes" only if we know it only occurs once, and in a
-- context where inlining it does not duplicate work.
preInlineUnconditionally :: LetAnn -> Bool
preInlineUnconditionally = \case
  PlainLet -> False  -- "Missing occurrence annotation"
  InlineLet   -> True
  NoInlineLet -> False
  LinearLet   -> False
  OccInfoPure (UsageInfo s (0, d)) | s <= One && d <= One -> True
  OccInfoPure _ -> False
  OccInfoImpure _ -> False

-- A context in which an E-kinded thing is to be reconstructed.  This amounts to
-- a defunctionalization of the interesting part of rebuilding an expression,
-- which supports now-local optimizations in `reconstruct`.  Read `reconstruct`
-- together with this data type: `reconstruct` is an interpreter for it.
--
-- Note: This pattern of inserting an object into a context can do local
-- optimizations that would otherwise be hidden from a peephole optimizer by
-- intervening bindings.  For example, table indexing only permits an Atom in
-- the array position, but `reconstruct` can check whether it's trying to insert
-- a `for` expression into that spot and perform the beta reduction immediately,
-- instead of emitting the binding.
data Context (from::E) (to::E) (o::S) where
  Stop :: Context e e o
  TabAppCtx :: SAtom i -> Subst InlineSubstVal i o
            -> Context SExpr e o -> Context SExpr e o
  CaseCtx :: [SAlt i] -> SType i -> EffectRow SimpIR i
          -> Subst InlineSubstVal i o
          -> Context SExpr e o -> Context SExpr e o
  EmitToAtomCtx :: Context SAtom e o -> Context SExpr e o
  EmitToNameCtx :: Context SAtomVar e o -> Context SAtom e o

class Inlinable (e1::E) where
  inline :: Emits o => Context e1 e2 o -> e1 i -> InlineM i o (e2 o)

  default inline :: (VisitGeneric e1 SimpIR, Emits o)
    => Context e1 e2 o -> e1 i -> InlineM i o (e2 o)
  inline ctx e = visitGeneric e >>= reconstruct ctx

instance NonAtomRenamer (InlineM i o) i o where renameN = substM
instance Emits o => Visitor (InlineM i o) SimpIR i o where
  visitType = inline Stop
  visitAtom = inline Stop
  visitLam  = inline Stop
  visitPi   = inline Stop

inlineExpr :: Emits o => Context SExpr e o -> SExpr i -> InlineM i o (e o)
inlineExpr ctx = \case
  Atom atom -> inlineAtom ctx atom
  TabApp _ tbl ix -> do
    s <- getSubst
    inlineAtom (TabAppCtx ix s ctx) tbl
  Case scrut alts (EffTy effs resultTy) -> do
    s <- getSubst
    inlineAtom (CaseCtx alts resultTy effs s ctx) scrut
  Block _ (Abs decls ans) -> do
    s <- inlineDeclsSubst decls
    withSubst s $ inlineExpr ctx ans
  expr -> visitGeneric expr >>= reconstruct ctx

inlineAtom :: Emits o => Context SExpr e o -> SAtom i -> InlineM i o (e o)
inlineAtom ctx = \case
  Stuck _ stuck -> inlineStuck ctx stuck
  Con con -> (toExpr <$> visitGeneric con) >>= reconstruct ctx

inlineStuck :: Emits o => Context SExpr e o -> SStuck i -> InlineM i o (e o)
inlineStuck ctx = \case
  Var name -> inlineName ctx name
  StuckProject i x -> do
    ans <- proj i =<< emit =<< inlineStuck Stop x
    reconstruct ctx $ Atom ans
  StuckTabApp _ _ -> error "not implemented"
  PtrVar t p -> do
    s <- mkStuck =<< (PtrVar t <$> substM p)
    reconstruct ctx (toExpr s)
  RepValAtom repVal -> do
    s <- mkStuck =<< (RepValAtom <$> visitGeneric repVal)
    reconstruct ctx (toExpr s)

inlineName :: Emits o => Context SExpr e o -> SAtomVar i -> InlineM i o (e o)
inlineName ctx name =
  lookupSubstM (atomVarName name) >>= \case
    Rename name' -> do
      -- This is the considerInline function from the Secrets paper; this
      -- is where we decide whether to inline a binding that isn't to be
      -- inlined unconditionally.
      -- For now, we just don't.  If we did, we would want to start with
      -- something like
      -- lookupEnv name' >>= \case
      --   (expr', presInfo) | inline presInfo expr' ctx -> inline
      --   no info -> do not inline (as now)
      v <- toAtomVar name'
      reconstruct ctx (toExpr v)
    SubstVal (DoneEx expr) -> dropSubst $ inlineExpr ctx expr
    SubstVal (SuspEx expr s') -> withSubst s' $ inlineExpr ctx expr

instance Inlinable SAtomVar where
  inline ctx a = inlineName (EmitToAtomCtx $ EmitToNameCtx ctx) a

instance Inlinable SAtom where
  inline ctx a = inlineAtom (EmitToAtomCtx ctx) a

instance Inlinable SType where
  inline ctx (TyCon ty) = (TyCon <$> visitGeneric ty) >>= reconstruct ctx

instance Inlinable SLam where
  inline ctx (LamExpr bs body) = do
    reconstruct ctx =<< withBinders bs \bs' -> do
      body' <- buildBlock $ inlineExpr Stop body
      return $ LamExpr bs' body'

withBinders
  :: Nest SBinder i i'
  -> (forall o'. DExt o o' => Nest SBinder o o' -> InlineM i' o' a)
  -> InlineM i o a
withBinders Empty cont = getDistinct >>= \Distinct -> cont Empty
withBinders (Nest (b:>ty) bs) cont = do
  ty' <- buildScopedAssumeNoDecls $ inline Stop ty
  withFreshBinder (getNameHint b) ty' \b' ->
    extendRenamer (b@>binderName b') do
      withBinders bs \bs' -> cont $ Nest b' bs'

instance Inlinable (PiType SimpIR) where
  inline ctx (PiType bs effTy)  =
    reconstruct ctx =<< withBinders bs \bs' -> do
      effTy' <- buildScopedAssumeNoDecls $ inline Stop effTy
      return $ PiType bs' effTy'

-- Still using InlineM because we may call back into inlining, and we wish to
-- retain our output binding environment.
reconstruct :: Emits o => Context e1 e2 o -> e1 o -> InlineM i o (e2 o)
reconstruct ctx e = case ctx of
  Stop -> return e
  TabAppCtx ix s ctx' -> withSubst s $ reconstructTabApp ctx' e ix
  CaseCtx alts resultTy effs s ctx' ->
    withSubst s $ reconstructCase ctx' e alts resultTy effs
  EmitToAtomCtx ctx' -> emit  e >>= reconstruct ctx'
  EmitToNameCtx ctx' -> emitToVar e >>= reconstruct ctx'
{-# INLINE reconstruct #-}

reconstructTabApp :: Emits o
  => Context SExpr e o -> SExpr o -> SAtom i -> InlineM i o (e o)
reconstructTabApp ctx expr i = case expr of
  PrimOp (Hof (TypedHof _ (For _ _ (UnaryLamExpr b body)))) -> do
    -- See NoteReconstructTabAppDecisions
    AtomVar i' _ <- inline (EmitToNameCtx Stop) i
    dropSubst $ extendSubst (b@>Rename i') do
      inlineExpr ctx body
  _ -> do
    array' <- emit expr
    i' <- inline Stop i
    reconstruct ctx =<< mkTabApp array' i'

reconstructCase :: Emits o
  => Context SExpr e o -> SExpr o -> [SAlt i] -> SType i -> EffectRow SimpIR i
  -> InlineM i o (e o)
reconstructCase ctx scrutExpr alts resultTy effs =
  case scrutExpr of
    Case sscrut salts _ -> do
      -- Perform case-of-case optimization
      -- TODO Add join points to reduce code duplication (and repeated inlining)
      -- of the arms of the outer case
      resultTy' <- inline Stop resultTy
      reconstruct ctx =<< (buildCase' sscrut resultTy' \i val -> do
        ans <- applyAbs (sink $ salts !! i) (SubstVal val) >>= emit
        buildCase ans (sink resultTy') \j jval -> do
          Abs b body <- return $ alts !! j
          extendSubst (b @> (SubstVal $ DoneEx $ Atom jval)) do
            inlineExpr Stop body >>= emit)
    _ -> do
      -- Attempt case-of-known-constructor optimization
      -- I can't use `buildCase` here because I want to propagate the incoming
      -- context `ctx` into the selected alternative if the optimization fires,
      -- but leave it around the whole reconstructed `Case` if it doesn't.
      scrut <- emit scrutExpr
      case scrut of
        Con con -> do
          SumCon _ i val <- return con
          Abs b body <- return $ alts !! i
          extendSubst (b @> (SubstVal $ DoneEx $ Atom val)) do
            inlineExpr ctx body
        Stuck _ _ -> do
          alts' <- mapM visitAlt alts
          resultTy' <- inline Stop resultTy
          effs' <- inline Stop effs
          reconstruct ctx $ Case scrut alts' (EffTy effs' resultTy')

instance Inlinable (EffectRow SimpIR)
instance Inlinable (EffTy     SimpIR)

-- === NoteReconstructTabAppDecisions ===

-- There's a decision here. Is it ok to inline the atoms in `ixsPref` into the
-- body `decls`? If so, should we pre-process them and carry them in `DoneEx`,
-- or suspend them in `SuspEx`? (If not, we can emit fresh bindings and use
-- `Rename`.) We can't make this decision properly without annotating the `for`
-- binders with occurrence information; even though `ixsPref` itself are atoms,
-- we may be carrying suspended inlining decisions that would want to make one
-- an expression, and thus force-inlining it may duplicate work.
--
-- There remains a decision between just emitting bindings, or running `mapM
-- (inline $ EmitToAtomCtx Stop)` and inlining the resulting atoms. In the
-- work-heavy case where an element of `ixsPref` becomes an expression after
-- inlining, the result will be the same; but in the work-light case where the
-- element remains an atom, more inlining can proceed. This decision only
-- affects the runtime of the inliner and the code size of the IR the inliner
-- produces.
--
-- Current status: Emitting bindings in the interest if "launch and iterate";
-- have not tried `EmitToAtomCtx`. Decision here. These decls have already been
-- processed by the inliner once, so their occurrence information is stale (and
-- should have been erased). Do we rerun occurrence analysis, or just complete
-- the pass without inlining any of them?
-- - Con rerunning: Slower
-- - Con completing: No detection of erroneous lack of occurrence info
-- For now went with "completing"; to detect erroneous lack of
-- occurrence info, change the relevant PlainLet cases above.
--
-- There's also a missed opportunity here to do more inlining in one pass: we
-- lost the occurrence information of the bindings, so we lost the ability to
-- inline them into the result, so in the common case that the result is a
-- variable reference, we will find ourselves emitting a rename, _which will
-- inhibit downstream inlining_ because a rename is not indexable.

-- === NoteConditionalInlining ===

-- TODO For now, this inliner does not do any conditional inlining. In order to
-- do it, we would need to augment the environment at this point, associating
-- name' to (expr', presInfo) so name' could be inlined at use sites.
--
-- Conditional inlining is different in Dex vs Haskell because Dex is strict. To
-- wit, once we have emitted the bidning for `expr'`, we are committed to doing
-- the work it represents unless it's inlined _everywhere_. For example,
--   xs = <something>
--   case <foo> of
--     Nothing -> xs  -- ok to inline here
--     Just _ -> xs ... xs  -- not ok here
-- If this were Haskell, it would be work-preserving for GHC to inline
-- `xs` into the `Nothing` arm, but in Dex it's not, unless we first
-- explicitly push the binding into the case like
--   case <foo> of
--     Nothing -> xs = <something>; xs
--     Just _ -> xs = <something>; xs ... xs
--
-- That said, the Secrets paper says that GHC only conditionally inlines
-- zero-work bindings anyway (or, more precisely, "bounded finite work"
-- bindings). All the heuristics about whether to inline at a particular site
-- are about code size and not increasing it overmuch. But, of course, inlining
-- even zero-work bindings can help runtime performance because it can unblock
-- other optimizations that otherwise could not occur across the binding.

-- === NoteSecretsSubtlety ===

-- A subtlety from the Secrets paper. In Haskell, it is feasible to have a
-- binding whose occurrence information indicates multiple uses, but which does
-- a small, bounded amount of runtime work. GHC will inline such a binding, but
-- not into contexts where GHC knows that no further optimizations are possible.
-- The example given in the paper is
--   f = \x -> E
--   g = \ys -> map f ys
-- Inlining f here is useless because it's not applied, and mildly costly
-- because it causes the closure to be allocated at every call to g rather than
-- just once.
-- TODO If we want to track this subtlety, we should make room for it in
-- the SizePreservationInfo ADT (maybe rename it), maybe with a
-- OnceButDuplicatesBoundedWork constructor.  Then only the true UsedOnce
-- would be inlined unconditionally here, and the
-- OnceButDuplicatesBoundedWork constructor could be inlined or not
-- depending on its usage context.  (This would correspond to the case
-- OnceUnsafe with whnfOrBot == True in the Secrets paper.)