module Debug.Provenance.Internal (
    -- * Callsites
    CallSite(..)
  , prettyCallSite
  , briefSrcLoc
  , callSite
  , callSiteWithLabel
    -- * Invocations
  , Invocation -- opaque
  , prettyInvocation
  , newInvocation
  , newInvocationFrom
  ) where

import Control.Monad.IO.Class
import Data.Hashable
import Data.HashMap.Strict (HashMap)
import Data.HashMap.Strict qualified as HashMap
import Data.IORef
import Data.List (intercalate)
import Data.Maybe (fromMaybe)
import GHC.Generics
import GHC.Stack
import System.IO.Unsafe (unsafePerformIO)

{-------------------------------------------------------------------------------
  Callsites
-------------------------------------------------------------------------------}

-- | Callsite
--
-- A callsite tells you where something was called: a location in the source,
-- and the name of the function that did the calling. Optionally, they can be
-- given an additional user-defined label also.
--
-- /NOTE/: If you are seeing @{unknown}@ instead of the function name,
-- the calling function does not have a 'HasCallStack' annotation:
--
-- > yourFunction :: HasCallStack => IO () -- 'HasCallStack' probably missing
-- > yourFunction = do
-- >     let cs = callSite
-- >     ..
--
-- Once you add this annotation, you should see @yourFunction@ instead of
-- @{unknown}@. Similarly, if you have local function definitions, it may
-- be useful to give those 'HasCallStack' constraints of their own:
--
-- > yourFunction :: HasCallStack => IO ()
-- > yourFunction = ..
-- >   where
-- >     someLocalFn :: HasCallStack => IO ()
-- >     someLocalFn = do
-- >         let cs = callSite
-- >         ..
--
-- In this example the 'HasCallStack' constraint on @someLocalFn@ means that the
-- calling function will be reported as @someLocalFn@ instead of @yourFunction@.
data CallSite = CallSite {
      CallSite -> Maybe SrcLoc
callSiteSrcLoc :: Maybe SrcLoc
    , CallSite -> Maybe String
callSiteCaller :: Maybe String
    , CallSite -> Maybe String
callSiteCallee :: Maybe String
    , CallSite -> Label
callSiteLabel  :: Label
    }
  deriving stock (CallSite -> CallSite -> Bool
(CallSite -> CallSite -> Bool)
-> (CallSite -> CallSite -> Bool) -> Eq CallSite
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CallSite -> CallSite -> Bool
== :: CallSite -> CallSite -> Bool
$c/= :: CallSite -> CallSite -> Bool
/= :: CallSite -> CallSite -> Bool
Eq)

instance Show CallSite where
  show :: CallSite -> String
show = CallSite -> String
prettyCallSite

-- | Label associated with 'CallSite'
--
-- This is an internal type.
data Label = Label String | NoLabel
  deriving stock (Label -> Label -> Bool
(Label -> Label -> Bool) -> (Label -> Label -> Bool) -> Eq Label
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Label -> Label -> Bool
== :: Label -> Label -> Bool
$c/= :: Label -> Label -> Bool
/= :: Label -> Label -> Bool
Eq, (forall x. Label -> Rep Label x)
-> (forall x. Rep Label x -> Label) -> Generic Label
forall x. Rep Label x -> Label
forall x. Label -> Rep Label x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Label -> Rep Label x
from :: forall x. Label -> Rep Label x
$cto :: forall x. Rep Label x -> Label
to :: forall x. Rep Label x -> Label
Generic)
  deriving anyclass (Eq Label
Eq Label =>
(Int -> Label -> Int) -> (Label -> Int) -> Hashable Label
Int -> Label -> Int
Label -> Int
forall a. Eq a => (Int -> a -> Int) -> (a -> Int) -> Hashable a
$chashWithSalt :: Int -> Label -> Int
hashWithSalt :: Int -> Label -> Int
$chash :: Label -> Int
hash :: Label -> Int
Hashable)

-- | Render 'CallSite' to human-readable format
prettyCallSite :: CallSite -> String
prettyCallSite :: CallSite -> String
prettyCallSite CallSite
cs =
    [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [
        String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
" -> " [
            String -> Maybe String -> String
forall a. a -> Maybe a -> a
fromMaybe String
"{unknown}" Maybe String
callSiteCaller
          , String -> Maybe String -> String
forall a. a -> Maybe a -> a
fromMaybe String
"{unknown}" Maybe String
callSiteCallee
          ]
      , String
" ("
      , String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ [[String]] -> [String]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [
            [ SrcLoc -> String
briefSrcLoc SrcLoc
loc
            | Just SrcLoc
loc <- [Maybe SrcLoc
callSiteSrcLoc]
            ]
          , [ ShowS
forall a. Show a => a -> String
show String
label
            | Label String
label <- [Label
callSiteLabel]
            ]
          ]
      , String
")"
      ]
  where
    CallSite{
        Maybe SrcLoc
callSiteSrcLoc :: CallSite -> Maybe SrcLoc
callSiteSrcLoc :: Maybe SrcLoc
callSiteSrcLoc
      , Maybe String
callSiteCaller :: CallSite -> Maybe String
callSiteCaller :: Maybe String
callSiteCaller
      , Maybe String
callSiteCallee :: CallSite -> Maybe String
callSiteCallee :: Maybe String
callSiteCallee
      , Label
callSiteLabel :: CallSite -> Label
callSiteLabel :: Label
callSiteLabel
      } = CallSite
cs

-- | Variant on 'prettySrcLoc' which omits the package and module name
briefSrcLoc :: SrcLoc -> [Char]
briefSrcLoc :: SrcLoc -> String
briefSrcLoc SrcLoc
loc = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
":" [
      SrcLoc -> String
srcLocFile SrcLoc
loc
    , Int -> String
forall a. Show a => a -> String
show (Int -> String) -> Int -> String
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Int
srcLocStartLine SrcLoc
loc
    , Int -> String
forall a. Show a => a -> String
show (Int -> String) -> Int -> String
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Int
srcLocStartCol SrcLoc
loc
    ]

instance Hashable CallSite where
  hashWithSalt :: Int -> CallSite -> Int
hashWithSalt Int
salt CallSite
cs =
      Int -> (Maybe String, Maybe String, Maybe String, Label) -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
salt (
          SrcLoc -> String
prettySrcLoc (SrcLoc -> String) -> Maybe SrcLoc -> Maybe String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe SrcLoc
callSiteSrcLoc
        , Maybe String
callSiteCaller
        , Maybe String
callSiteCallee
        , Label
callSiteLabel
        )
    where
      CallSite{
          Maybe SrcLoc
callSiteSrcLoc :: CallSite -> Maybe SrcLoc
callSiteSrcLoc :: Maybe SrcLoc
callSiteSrcLoc
        , Maybe String
callSiteCaller :: CallSite -> Maybe String
callSiteCaller :: Maybe String
callSiteCaller
        , Maybe String
callSiteCallee :: CallSite -> Maybe String
callSiteCallee :: Maybe String
callSiteCallee
        , Label
callSiteLabel :: CallSite -> Label
callSiteLabel :: Label
callSiteLabel
        } = CallSite
cs


-- | Current 'CallSite'
callSite :: HasCallStack => CallSite
callSite :: HasCallStack => CallSite
callSite = (HasCallStack => CallSite) -> CallSite
forall a. HasCallStack => (HasCallStack => a) -> a
withFrozenCallStack ((HasCallStack => CallSite) -> CallSite)
-> (HasCallStack => CallSite) -> CallSite
forall a b. (a -> b) -> a -> b
$ HasCallStack => Label -> CallSite
Label -> CallSite
mkCallSite Label
NoLabel

-- | Current 'CallSite' with user-defined label
callSiteWithLabel :: HasCallStack => String -> CallSite
callSiteWithLabel :: HasCallStack => String -> CallSite
callSiteWithLabel String
label = (HasCallStack => CallSite) -> CallSite
forall a. HasCallStack => (HasCallStack => a) -> a
withFrozenCallStack ((HasCallStack => CallSite) -> CallSite)
-> (HasCallStack => CallSite) -> CallSite
forall a b. (a -> b) -> a -> b
$ HasCallStack => Label -> CallSite
Label -> CallSite
mkCallSite (String -> Label
Label String
label)

-- | Internal auxiliary to 'callSite' and 'callSiteWithLabel'
mkCallSite  :: HasCallStack => Label -> CallSite
mkCallSite :: HasCallStack => Label -> CallSite
mkCallSite Label
callSiteLabel = CallStack -> CallSite
aux CallStack
HasCallStack => CallStack
callStack
  where
    aux :: CallStack -> CallSite
    aux :: CallStack -> CallSite
aux CallStack
cs =
        -- drop the call to @callSite{withLabel}@
        case CallStack -> [(String, SrcLoc)]
getCallStack CallStack
cs of
          (String, SrcLoc)
_ : (String
callee, SrcLoc
loc) : [] -> CallSite {
              callSiteSrcLoc :: Maybe SrcLoc
callSiteSrcLoc = SrcLoc -> Maybe SrcLoc
forall a. a -> Maybe a
Just SrcLoc
loc
            , callSiteCaller :: Maybe String
callSiteCaller = Maybe String
forall a. Maybe a
Nothing
            , callSiteCallee :: Maybe String
callSiteCallee = String -> Maybe String
forall a. a -> Maybe a
Just String
callee
            , Label
callSiteLabel :: Label
callSiteLabel :: Label
callSiteLabel
            }
          (String, SrcLoc)
_ : (String
callee, SrcLoc
loc) : (String
caller, SrcLoc
_) : [(String, SrcLoc)]
_ -> CallSite {
              callSiteSrcLoc :: Maybe SrcLoc
callSiteSrcLoc   = SrcLoc -> Maybe SrcLoc
forall a. a -> Maybe a
Just SrcLoc
loc
            , callSiteCaller :: Maybe String
callSiteCaller = String -> Maybe String
forall a. a -> Maybe a
Just String
caller
            , callSiteCallee :: Maybe String
callSiteCallee = String -> Maybe String
forall a. a -> Maybe a
Just String
callee
            , Label
callSiteLabel :: Label
callSiteLabel :: Label
callSiteLabel
            }
          [(String, SrcLoc)]
_otherwise -> CallSite {
              callSiteSrcLoc :: Maybe SrcLoc
callSiteSrcLoc = Maybe SrcLoc
forall a. Maybe a
Nothing
            , callSiteCaller :: Maybe String
callSiteCaller = Maybe String
forall a. Maybe a
Nothing
            , callSiteCallee :: Maybe String
callSiteCallee = Maybe String
forall a. Maybe a
Nothing
            , Label
callSiteLabel :: Label
callSiteLabel :: Label
callSiteLabel
            }

{-------------------------------------------------------------------------------
  Invocations
-------------------------------------------------------------------------------}

-- | Invocation
--
-- An invocation not only tells you the /where/, but also the /when/: it pairs a
-- 'CallSite' with a count, automatically incremented on each call to
-- 'newInvocation'. Each 'CallSite' uses its own counter.
data Invocation = Invocation CallSite Int
  deriving stock (Invocation -> Invocation -> Bool
(Invocation -> Invocation -> Bool)
-> (Invocation -> Invocation -> Bool) -> Eq Invocation
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Invocation -> Invocation -> Bool
== :: Invocation -> Invocation -> Bool
$c/= :: Invocation -> Invocation -> Bool
/= :: Invocation -> Invocation -> Bool
Eq)

instance Show Invocation where
  show :: Invocation -> String
show = Invocation -> String
prettyInvocation

-- | Render 'Invocation' to human-readable format
prettyInvocation :: Invocation -> String
prettyInvocation :: Invocation -> String
prettyInvocation (Invocation CallSite
cs Int
n) =
    [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [
        String -> Maybe String -> String
forall a. a -> Maybe a -> a
fromMaybe String
"{unknown}" Maybe String
callSiteCaller
      , String
" ("
      , String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ [[String]] -> [String]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [
            [ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
":" [
                  SrcLoc -> String
srcLocFile SrcLoc
loc
                , Int -> String
forall a. Show a => a -> String
show (Int -> String) -> Int -> String
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Int
srcLocStartLine SrcLoc
loc
                , Int -> String
forall a. Show a => a -> String
show (Int -> String) -> Int -> String
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Int
srcLocStartCol SrcLoc
loc
                ]
            | Just SrcLoc
loc <- [Maybe SrcLoc
callSiteSrcLoc]
            ]
          , [ ShowS
forall a. Show a => a -> String
show String
label
            | Label String
label <- [Label
callSiteLabel]
            ]
          ]
      , String
") #"
      , Int -> String
forall a. Show a => a -> String
show Int
n
      ]
  where
    -- the callee is 'newInvocation'
    CallSite{
        Maybe SrcLoc
callSiteSrcLoc :: CallSite -> Maybe SrcLoc
callSiteSrcLoc :: Maybe SrcLoc
callSiteSrcLoc
      , Maybe String
callSiteCaller :: CallSite -> Maybe String
callSiteCaller :: Maybe String
callSiteCaller
      , Label
callSiteLabel :: CallSite -> Label
callSiteLabel :: Label
callSiteLabel
      } = CallSite
cs

-- | New invocation
--
-- See 'Invocation' for discussion.
newInvocation :: (HasCallStack, MonadIO m) => m Invocation
newInvocation :: forall (m :: * -> *). (HasCallStack, MonadIO m) => m Invocation
newInvocation =
    -- We intentionally do /NOT/ freeze the callstack here: when function @foo@
    -- calls @newInvocation@, we want a 'CallSite' of @foo -> newInvocation@,
    -- not @bar -> foo@.
    CallSite -> m Invocation
forall (m :: * -> *). MonadIO m => CallSite -> m Invocation
newInvocationFrom CallSite
HasCallStack => CallSite
callSite

-- | Generalization of 'newInvocation'
newInvocationFrom :: MonadIO m => CallSite -> m Invocation
newInvocationFrom :: forall (m :: * -> *). MonadIO m => CallSite -> m Invocation
newInvocationFrom CallSite
cs = IO Invocation -> m Invocation
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Invocation -> m Invocation) -> IO Invocation -> m Invocation
forall a b. (a -> b) -> a -> b
$ do
    IORef (HashMap CallSite Int)
-> (HashMap CallSite Int -> (HashMap CallSite Int, Invocation))
-> IO Invocation
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (HashMap CallSite Int)
globalCounters ((HashMap CallSite Int -> (HashMap CallSite Int, Invocation))
 -> IO Invocation)
-> (HashMap CallSite Int -> (HashMap CallSite Int, Invocation))
-> IO Invocation
forall a b. (a -> b) -> a -> b
$ \HashMap CallSite Int
counters ->
      let i :: Int
i = Int -> CallSite -> HashMap CallSite Int -> Int
forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
HashMap.findWithDefault Int
1 CallSite
cs HashMap CallSite Int
counters
      in (CallSite -> Int -> HashMap CallSite Int -> HashMap CallSite Int
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert CallSite
cs (Int -> Int
forall a. Enum a => a -> a
succ Int
i) HashMap CallSite Int
counters, CallSite -> Int -> Invocation
Invocation CallSite
cs Int
i)

{-------------------------------------------------------------------------------
  Internal: globals
-------------------------------------------------------------------------------}

globalCounters :: IORef (HashMap CallSite Int)
{-# NOINLINE globalCounters #-}
globalCounters :: IORef (HashMap CallSite Int)
globalCounters = IO (IORef (HashMap CallSite Int)) -> IORef (HashMap CallSite Int)
forall a. IO a -> a
unsafePerformIO (IO (IORef (HashMap CallSite Int)) -> IORef (HashMap CallSite Int))
-> IO (IORef (HashMap CallSite Int))
-> IORef (HashMap CallSite Int)
forall a b. (a -> b) -> a -> b
$ HashMap CallSite Int -> IO (IORef (HashMap CallSite Int))
forall a. a -> IO (IORef a)
newIORef HashMap CallSite Int
forall k v. HashMap k v
HashMap.empty