    {-# LANGUAGE TupleSections #-}

    -----------------------------------------------------------------------------
    -- |
    -- Module      :  GHC.StgToJS.ExprCtx
    -- Copyright   :  (c) The University of Glasgow 2001
    -- License     :  BSD-style (see the file LICENSE)
    --
    -- Maintainer  :  Jeffrey Young  <jeffrey.young@iohk.io>
    --                Luite Stegeman <luite.stegeman@iohk.io>
    --                Sylvain Henry  <sylvain.henry@iohk.io>
    --                Josh Meredith  <josh.meredith@iohk.io>
    -- Stability   :  experimental
    --
    -- TODO: Write my description!
    -----------------------------------------------------------------------------

    module GHC.StgToJS.ExprCtx
      ( ExprCtx
      , initExprCtx
      , ctxIsEvaluated
      , ctxSetSrcSpan
      , ctxSrcSpan
      , ctxSetTop
      , ctxTarget
      , ctxSetTarget
      -- * Let-no-escape
      , ctxClearLneFrame
      , ctxUpdateLneFrame
      , ctxLneFrameVars
      , ctxLneFrameSize
      , ctxIsLneBinding
      , ctxIsLneLiveVar
      , ctxLneBindingStackSize
      , ctxLneShrinkStack
      )
    where

    import GHC.Prelude

    import GHC.StgToJS.Types

    import GHC.Types.Unique.FM
    import GHC.Types.Var
    import GHC.Types.SrcLoc
    import GHC.Types.Id
    import GHC.Types.Id.Info

    import GHC.Stg.InferTags.TagSig

    import GHC.Utils.Outputable
    import GHC.Utils.Panic

    import Data.Maybe

    -- | Context into which an expression is evaluated
    data ExprCtx = ExprCtx
      { ExprCtx -> Id
    ctxTop        :: Id
        -- ^ Top-level binding Id

      , ExprCtx -> [TypedExpr]
    ctxTarget     :: [TypedExpr]
        -- ^ Target variables for the evaluated expression

      , ExprCtx -> Maybe RealSrcSpan
    ctxSrcSpan    :: Maybe RealSrcSpan
        -- ^ Source location

      ----------------------------
      -- Handling of let-no-escape

      , ExprCtx -> UniqFM Id Int
    ctxLneFrameBs :: UniqFM Id Int
        -- ^ LNE bindings with their expected stack size.
        --
        -- The Int is the size of the stack when the LNE binding was defined.
        -- We need to shrink the stack back to this size when we enter one of the
        -- associated binder rhs: it expects its free variables at certain offsets
        -- in the stack.

      , ExprCtx -> [(Id, Int)]
    ctxLneFrameVars :: [(Id,Int)]
        -- ^ Contents of current LNE frame
        --
        -- Variables and their index on the stack

      , ExprCtx -> Int
    ctxLneFrameSize :: {-# UNPACK #-} !Int
        -- ^ Cache the length of `ctxLneFrameVars`

      }

    -- | Initialize an expression context in the context of the given top-level
    -- binding Id
    initExprCtx :: Id -> ExprCtx
    initExprCtx :: Id -> ExprCtx
    initExprCtx Id
    i = ExprCtx
      { ctxTop :: Id
    ctxTop          = Id
    i
      , ctxTarget :: [TypedExpr]
    ctxTarget       = []
      , ctxLneFrameBs :: UniqFM Id Int
    ctxLneFrameBs   = UniqFM Id Int
    forall {k} (key :: k) elt. UniqFM key elt
    emptyUFM
      , ctxLneFrameVars :: [(Id, Int)]
    ctxLneFrameVars = []
      , ctxLneFrameSize :: Int
    ctxLneFrameSize = Int
    0
      , ctxSrcSpan :: Maybe RealSrcSpan
    ctxSrcSpan      = Maybe RealSrcSpan
    forall a. Maybe a
    Nothing
      }

    -- | Set target
    ctxSetTarget :: [TypedExpr] -> ExprCtx -> ExprCtx
    ctxSetTarget :: [TypedExpr] -> ExprCtx -> ExprCtx
    ctxSetTarget [TypedExpr]
    t ExprCtx
    ctx = ExprCtx
    ctx { ctxTarget = t }

    -- | Set top-level binding Id
    ctxSetTop :: Id -> ExprCtx -> ExprCtx
    ctxSetTop :: Id -> ExprCtx -> ExprCtx
    ctxSetTop Id
    i ExprCtx
    ctx = ExprCtx
    ctx { ctxTop = i }

    -- | Set source location
    ctxSetSrcSpan :: RealSrcSpan -> ExprCtx -> ExprCtx
    ctxSetSrcSpan :: RealSrcSpan -> ExprCtx -> ExprCtx
    ctxSetSrcSpan RealSrcSpan
    span ExprCtx
    ctx = ExprCtx
    ctx { ctxSrcSpan = Just span }

    -- | Update let-no-escape frame
    ctxUpdateLneFrame :: [(Id,Int)] -> [Id] -> ExprCtx -> ExprCtx
    ctxUpdateLneFrame :: [(Id, Int)] -> [Id] -> ExprCtx -> ExprCtx
    ctxUpdateLneFrame [(Id, Int)]
    new_spilled_vars [Id]
    new_lne_ids ExprCtx
    ctx =
      let old_frame_size :: Int
    old_frame_size = ExprCtx -> Int
    ctxLneFrameSize ExprCtx
    ctx
          new_frame_size :: Int
    new_frame_size = Int
    old_frame_size Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + [(Id, Int)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [(Id, Int)]
    new_spilled_vars
      in ExprCtx
    ctx
        { ctxLneFrameBs   = addListToUFM (ctxLneFrameBs ctx) (map (,new_frame_size) new_lne_ids)
        , ctxLneFrameSize = new_frame_size
        , ctxLneFrameVars = ctxLneFrameVars ctx ++ new_spilled_vars
        }

    -- | Remove information about the current LNE frame
    ctxClearLneFrame :: ExprCtx -> ExprCtx
    ctxClearLneFrame :: ExprCtx -> ExprCtx
    ctxClearLneFrame ExprCtx
    ctx =
      ExprCtx
    ctx
        { ctxLneFrameBs   = emptyUFM
        , ctxLneFrameVars = []
        , ctxLneFrameSize = 0
        }

    -- | Predicate: do we know for sure that the given Id is evaluated?
    ctxIsEvaluated :: Id -> Bool
    ctxIsEvaluated :: Id -> Bool
    ctxIsEvaluated Id
    i =
      Bool -> (TagSig -> Bool) -> Maybe TagSig -> Bool
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe Bool
    False TagSig -> Bool
    isTaggedSig (Id -> Maybe TagSig
    idTagSig_maybe Id
    i)
      Bool -> Bool -> Bool
    && IdDetails -> Bool
    go (Id -> IdDetails
    idDetails Id
    i)
      where
        go :: IdDetails -> Bool
    go JoinId{} = Bool
    False
        go IdDetails
    _        = Bool
    True

          -- DFunId new_type -> not new_type
          --    -- DFuns terminate, unless the dict is implemented
          --    -- with a newtype in which case they may not

          -- DataConWorkId {} -> True

          -- ClassOpId {} -> False
          --   -- suppose an argument, and we don't have one

          -- PrimOpId op _ -> primop_ok op
          --   -- probably already handled by StgOpApp

          -- JoinId {} -> False
          --   -- Don't speculate join points

          -- TickBoxOpId {} -> False
          --   -- Don't speculate box ticking

          -- -- Tagged (evaluated) ids
          -- _ | Just sig <- idTagSig_maybe i
          --   , isTaggedSig sig
          --   -> True

          -- _ -> False

    -- | Does the given Id correspond to a LNE binding
    ctxIsLneBinding :: ExprCtx -> Id -> Bool
    ctxIsLneBinding :: ExprCtx -> Id -> Bool
    ctxIsLneBinding ExprCtx
    ctx Id
    i = Maybe Int -> Bool
    forall a. Maybe a -> Bool
    isJust (ExprCtx -> Id -> Maybe Int
    ctxLneBindingStackSize ExprCtx
    ctx Id
    i)

    -- | Does the given Id correspond to a LNE live var on the stack
    ctxIsLneLiveVar :: ExprCtx -> Id -> Bool
    ctxIsLneLiveVar :: ExprCtx -> Id -> Bool
    ctxIsLneLiveVar ExprCtx
    ctx Id
    i = Id
    i Id -> [Id] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `elem` ((Id, Int) -> Id) -> [(Id, Int)] -> [Id]
    forall a b. (a -> b) -> [a] -> [b]
    map (Id, Int) -> Id
    forall a b. (a, b) -> a
    fst (ExprCtx -> [(Id, Int)]
    ctxLneFrameVars ExprCtx
    ctx)

    -- | Return the LNE stack size associated to the given Id.
    -- Return Nothing when the Id doesn't correspond to a LNE binding.
    ctxLneBindingStackSize :: ExprCtx -> Id -> Maybe Int
    ctxLneBindingStackSize :: ExprCtx -> Id -> Maybe Int
    ctxLneBindingStackSize ExprCtx
    ctx Id
    i = UniqFM Id Int -> Id -> Maybe Int
    forall key elt. Uniquable key => UniqFM key elt -> key -> Maybe elt
    lookupUFM (ExprCtx -> UniqFM Id Int
    ctxLneFrameBs ExprCtx
    ctx) Id
    i

    -- | Shrink the LNE stack to the given size
    ctxLneShrinkStack :: ExprCtx -> Int -> ExprCtx
    ctxLneShrinkStack :: ExprCtx -> Int -> ExprCtx
    ctxLneShrinkStack ExprCtx
    ctx Int
    n =
      let l :: Int
    l = ExprCtx -> Int
    ctxLneFrameSize ExprCtx
    ctx
      in Bool -> SDoc -> ExprCtx -> ExprCtx
    forall a. HasCallStack => Bool -> SDoc -> a -> a
    assertPpr
          (Int
    l Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Int
    n)
          ([SDoc] -> SDoc
    forall doc. IsDoc doc => [doc] -> doc
    vcat [ String -> SDoc
    forall doc. IsLine doc => String -> doc
    text String
    "ctxLneShrinkStack: let-no-escape stack too short:"
                , Int -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr Int
    l
                , String -> SDoc
    forall doc. IsLine doc => String -> doc
    text String
    " < "
                , Int -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr Int
    n
                ])
          (ExprCtx
    ctx { ctxLneFrameVars = take n (ctxLneFrameVars ctx)
               , ctxLneFrameSize = n
               }
          )