{-# 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.EnforceEpt.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`

  }

instance Outputable ExprCtx where
  ppr :: ExprCtx -> SDoc
ppr ExprCtx
g = SDoc -> Int -> SDoc -> SDoc
hang (String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"ExprCtx") Int
2 (SDoc -> SDoc) -> SDoc -> SDoc
forall a b. (a -> b) -> a -> b
$ [SDoc] -> SDoc
forall doc. IsDoc doc => [doc] -> doc
vcat
            [ [SDoc] -> SDoc
forall doc. IsLine doc => [doc] -> doc
hcat [String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"ctxTop: ", Id -> SDoc
forall a. Outputable a => a -> SDoc
ppr (ExprCtx -> Id
ctxTop ExprCtx
g)]
            , [SDoc] -> SDoc
forall doc. IsLine doc => [doc] -> doc
hcat [String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"ctxTarget:", [TypedExpr] -> SDoc
forall a. Outputable a => a -> SDoc
ppr (ExprCtx -> [TypedExpr]
ctxTarget ExprCtx
g)]
            , [SDoc] -> SDoc
forall doc. IsLine doc => [doc] -> doc
hcat [String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"ctxSrcSpan:", Maybe RealSrcSpan -> SDoc
forall a. Outputable a => a -> SDoc
ppr (ExprCtx -> Maybe RealSrcSpan
ctxSrcSpan ExprCtx
g)]
            , [SDoc] -> SDoc
forall doc. IsLine doc => [doc] -> doc
hcat [String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"ctxLneFrameBs:", UniqFM Id Int -> SDoc
forall a. Outputable a => a -> SDoc
ppr (ExprCtx -> UniqFM Id Int
ctxLneFrameBs ExprCtx
g)]
            , [SDoc] -> SDoc
forall doc. IsLine doc => [doc] -> doc
hcat [String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"ctxLneFrameVars:", [(Id, Int)] -> SDoc
forall a. Outputable a => a -> SDoc
ppr (ExprCtx -> [(Id, Int)]
ctxLneFrameVars ExprCtx
g)]
            , [SDoc] -> SDoc
forall doc. IsLine doc => [doc] -> doc
hcat [String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"ctxLneFrameSize:", Int -> SDoc
forall a. Outputable a => a -> SDoc
ppr (ExprCtx -> Int
ctxLneFrameSize ExprCtx
g)]
            ]

-- | 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
           }
      )