content: -- 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
  import Types.Top
  
  -- === 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.)
  