module GHC.Stack.Profiler.Internal.Decode (
  CallStackSample (..),
  StackSymbolTable,
  SymbolTableWriter,
  initMessages,
  serializeCallStack,
  serializeMessage,
  serializeMessages,
  decodeToCallStack,
  definitions,
) where

import Control.Concurrent.STM
import Control.Exception (assert)
import Data.Binary
import Data.Binary.Put
import qualified Data.ByteString.Lazy as LBS
import qualified Data.List.NonEmpty as NonEmpty
import GHC.Generics (Generic)
import GHC.Stack.CloneStack (StackSnapshot)
import GHC.Stack.Profiler.Core
import GHC.Stack.Profiler.Internal.Stack.Decode (decodeStackWithIpProvId)
import GHC.Stack.Profiler.Internal.SymbolTable

-- | A 'CallStackSample' is a snapshot of a threads RTS callstack.
-- This callstack is a copy of the original callstack, so can be traversed and
-- decoded without affecting the running thread.
--
-- The 'StackSnapshot' is a boxed value and needs to be garbage collected.
-- Note, as long as 'StackSnapshot' is alive, you keep the full callstack
-- alive, which might be quite expensive.
data CallStackSample = CallStackSample
  { CallStackSample -> ThreadId
callStackSampleThreadId :: !ThreadId
  , CallStackSample -> CapabilityId
callStackSampleCapabilityId :: !CapabilityId
  , CallStackSample -> StackSnapshot
callStackSampleStackSnapshot :: !StackSnapshot
  }
  deriving ((forall x. CallStackSample -> Rep CallStackSample x)
-> (forall x. Rep CallStackSample x -> CallStackSample)
-> Generic CallStackSample
forall x. Rep CallStackSample x -> CallStackSample
forall x. CallStackSample -> Rep CallStackSample x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. CallStackSample -> Rep CallStackSample x
from :: forall x. CallStackSample -> Rep CallStackSample x
$cto :: forall x. Rep CallStackSample x -> CallStackSample
to :: forall x. Rep CallStackSample x -> CallStackSample
Generic)

decodeToCallStack :: CallStackSample -> IO CallStack
decodeToCallStack :: CallStackSample -> IO CallStack
decodeToCallStack CallStackSample
sample = do
  frames <- StackSnapshot -> IO [StackItem]
decodeStackWithIpProvId (StackSnapshot -> IO [StackItem])
-> StackSnapshot -> IO [StackItem]
forall a b. (a -> b) -> a -> b
$ CallStackSample -> StackSnapshot
callStackSampleStackSnapshot CallStackSample
sample
  let
    -- removes immediate duplicates
    callStackItems = (NonEmpty StackItem -> StackItem)
-> [NonEmpty StackItem] -> [StackItem]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NonEmpty StackItem -> StackItem
forall a. NonEmpty a -> a
NonEmpty.head ([NonEmpty StackItem] -> [StackItem])
-> [NonEmpty StackItem] -> [StackItem]
forall a b. (a -> b) -> a -> b
$ [StackItem] -> [NonEmpty StackItem]
forall (f :: * -> *) a. (Foldable f, Eq a) => f a -> [NonEmpty a]
NonEmpty.group [StackItem]
frames

  pure
    MkCallStack
      { callThreadId = callStackSampleThreadId sample
      , callCapabilityId = callStackSampleCapabilityId sample
      , callStack = callStackItems
      }

serializeCallStack :: StackSymbolTable -> CallStack -> STM [Message]
serializeCallStack :: StackSymbolTable -> CallStack -> STM [Message]
serializeCallStack StackSymbolTable
tableRef CallStack
callStackMessage = do
  table <- StackSymbolTable -> STM (SymbolTableWriter MapTable)
readSymbolTable StackSymbolTable
tableRef
  let
    (eventlogMessages, newTable) = dehydrateCallStack table callStackMessage
  writeSymbolTable newTable tableRef
  pure eventlogMessages

serializeMessage :: Message -> LBS.ByteString
serializeMessage :: Message -> ByteString
serializeMessage = Put -> ByteString
runPut (Put -> ByteString) -> (Message -> Put) -> Message -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Message -> Put
forall t. Binary t => t -> Put
put

serializeMessages :: [Message] -> [LBS.ByteString]
serializeMessages :: [Message] -> [ByteString]
serializeMessages = (Message -> ByteString) -> [Message] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map Message -> ByteString
serializeMessage

initMessages :: SymbolTableWriter MapTable -> [Message]
initMessages :: SymbolTableWriter MapTable -> [Message]
initMessages SymbolTableWriter MapTable
symbolTable =
  let
    ([StringDef]
stringDefs, [SourceLocationDef]
srcLocDefs) = SymbolTableWriter MapTable -> ([StringDef], [SourceLocationDef])
definitions SymbolTableWriter MapTable
symbolTable
  in
    ( (StringDef -> Message) -> [StringDef] -> [Message]
forall a b. (a -> b) -> [a] -> [b]
map StringDef -> Message
StringDef [StringDef]
stringDefs
        [Message] -> [Message] -> [Message]
forall a. [a] -> [a] -> [a]
++ (SourceLocationDef -> Message) -> [SourceLocationDef] -> [Message]
forall a b. (a -> b) -> [a] -> [b]
map SourceLocationDef -> Message
SourceLocationDef [SourceLocationDef]
srcLocDefs
    )

definitions :: SymbolTableWriter MapTable -> ([StringDef], [SourceLocationDef])
definitions :: SymbolTableWriter MapTable -> ([StringDef], [SourceLocationDef])
definitions SymbolTableWriter MapTable
table =
  let
    knownStrings :: [(StringId, Text)]
knownStrings = MapTable -> [(StringId, Text)]
getKnownStrings (MapTable -> [(StringId, Text)]) -> MapTable -> [(StringId, Text)]
forall a b. (a -> b) -> a -> b
$ SymbolTableWriter MapTable -> MapTable
forall tbl. SymbolTableWriter tbl -> tbl
writerTable SymbolTableWriter MapTable
table
    knownSrcLocs :: [(SourceLocationId, SourceLocation)]
knownSrcLocs = MapTable -> [(SourceLocationId, SourceLocation)]
getKnownSourceLocations (MapTable -> [(SourceLocationId, SourceLocation)])
-> MapTable -> [(SourceLocationId, SourceLocation)]
forall a b. (a -> b) -> a -> b
$ SymbolTableWriter MapTable -> MapTable
forall tbl. SymbolTableWriter tbl -> tbl
writerTable SymbolTableWriter MapTable
table

    stringDefs :: [StringDef]
stringDefs =
      ((StringId, Text) -> StringDef)
-> [(StringId, Text)] -> [StringDef]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((StringId -> Text -> StringDef) -> (StringId, Text) -> StringDef
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry StringId -> Text -> StringDef
MkStringDef) [(StringId, Text)]
knownStrings

    srcLocDefs :: [SourceLocationDef]
srcLocDefs =
      ((SourceLocationId, SourceLocation) -> SourceLocationDef)
-> [(SourceLocationId, SourceLocation)] -> [SourceLocationDef]
forall a b. (a -> b) -> [a] -> [b]
map ((SourceLocationId -> SourceLocation -> SourceLocationDef)
-> (SourceLocationId, SourceLocation) -> SourceLocationDef
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry SourceLocationId -> SourceLocation -> SourceLocationDef
go) [(SourceLocationId, SourceLocation)]
knownSrcLocs
  in
    ( [StringDef]
stringDefs
    , [SourceLocationDef]
srcLocDefs
    )
 where
  go :: SourceLocationId -> SourceLocation -> SourceLocationDef
  go :: SourceLocationId -> SourceLocation -> SourceLocationDef
go SourceLocationId
sid SourceLocation
s =
    let
      (StringId
fileId, Bool
newFileName, MapTable
_) = SymbolTableWriter MapTable
-> MapTable -> Text -> (StringId, Bool, MapTable)
forall tbl.
SymbolTableWriter tbl -> tbl -> Text -> (StringId, Bool, tbl)
lookupOrInsertText SymbolTableWriter MapTable
table (SymbolTableWriter MapTable -> MapTable
forall tbl. SymbolTableWriter tbl -> tbl
writerTable SymbolTableWriter MapTable
table) (SourceLocation -> Text
fileName SourceLocation
s)
    in
      -- These should always be found
      Bool -> SourceLocationDef -> SourceLocationDef
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Bool -> Bool
not Bool
newFileName) (SourceLocationDef -> SourceLocationDef)
-> SourceLocationDef -> SourceLocationDef
forall a b. (a -> b) -> a -> b
$
        MkSourceLocationDef
          { sourceLocationDefId :: SourceLocationId
sourceLocationDefId = SourceLocationId
sid
          , sourceLocationDefRow :: Word32
sourceLocationDefRow = SourceLocation -> Word32
line SourceLocation
s
          , sourceLocationDefColumn :: Word32
sourceLocationDefColumn = SourceLocation -> Word32
column SourceLocation
s
          , sourceLocationDefFilename :: StringId
sourceLocationDefFilename = StringId
fileId
          }