{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fno-cse #-}

-- |
-- Module      :   Grisette.Internal.SymPrim.Prim.Internal.Caches
-- Copyright   :   (c) Sirui Lu 2021-2024
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Internal.SymPrim.Prim.Internal.Caches
  ( SomeStableName (..),
    Id,
    StableIdent,
    Digest,
    CachedInfo (..),
    Interned (..),
    intern,
    haveCache,
    threadCacheSize,
    -- dumpThreadCache,
    threadCacheLiveSize,
  )
where

import Control.Concurrent
  ( MVar,
    ThreadId,
    myThreadId,
    newMVar,
    putMVar,
    takeMVar,
  )
import Control.Monad (replicateM)
import qualified Data.Array as A
import Data.Atomics (atomicModifyIORefCAS, atomicModifyIORefCAS_)
import Data.Data (Proxy (Proxy), Typeable, typeRepFingerprint)
import Data.Foldable (traverse_)
import qualified Data.HashMap.Strict as HM
import Data.Hashable (Hashable)
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Maybe (isJust)
import qualified Data.Vector.Unboxed.Mutable as M
import Data.Word (Word32)
import GHC.Base (Any)
import GHC.Fingerprint (Fingerprint)
import GHC.IO (unsafePerformIO)
import GHC.StableName (makeStableName)
import GHC.Weak (Weak, deRefWeak, finalize)
import Grisette.Internal.SymPrim.Prim.Internal.Utils
  ( SomeStableName (SomeStableName),
    WeakThreadId,
    WeakThreadIdRef,
    mkWeakStableNameRefWithFinalizer,
    mkWeakThreadIdRefWithFinalizer,
    myWeakThreadId,
    weakThreadId,
  )
import System.Mem.StableName (StableName)
import Type.Reflection (someTypeRep)
import Unsafe.Coerce (unsafeCoerce)

-- | A unique identifier for a term.
type Id = Word32

-- | The identity of a term.
type StableIdent = StableName Any

-- | A digest of a term.
type Digest = Word32

-- | Information about a cached term.
data CachedInfo = CachedInfo
  { CachedInfo -> WeakThreadId
cachedThreadId :: {-# UNPACK #-} !WeakThreadId,
    CachedInfo -> Word32
cachedDigest :: {-# UNPACK #-} !Digest,
    CachedInfo -> Word32
cachedId :: {-# UNPACK #-} !Id,
    CachedInfo -> StableIdent
cachedStableIdent :: {-# UNPACK #-} !StableIdent
  }

newtype Cache t = Cache {forall t. Cache t -> Array Int (CacheState t)
getCache :: A.Array Int (CacheState t)}

type HashTable k v = IORef (HM.HashMap k v)

data CacheState t where
  CacheState ::
    { forall t. CacheState t -> MVar ()
_sem :: MVar (),
      forall t. CacheState t -> IOVector Word32
_nextId :: M.IOVector Id,
      forall t.
CacheState t
-> HashTable (Description t) (Word32, Weak StableIdent)
_currentThread :: HashTable (Description t) (Id, Weak StableIdent)
    } ->
    CacheState t

finalizeCacheState :: CacheState t -> IO ()
finalizeCacheState :: forall t. CacheState t -> IO ()
finalizeCacheState (CacheState MVar ()
_ IOVector Word32
_ HashTable (Description t) (Word32, Weak StableIdent)
s) = do
  m <- HashTable (Description t) (Word32, Weak StableIdent)
-> IO (HashMap (Description t) (Word32, Weak StableIdent))
forall a. IORef a -> IO a
readIORef HashTable (Description t) (Word32, Weak StableIdent)
s
  traverse_ (\(Word32
_, Weak StableIdent
w) -> Weak StableIdent -> IO ()
forall v. Weak v -> IO ()
finalize Weak StableIdent
w) m

finalizeCache :: Cache t -> IO ()
finalizeCache :: forall t. Cache t -> IO ()
finalizeCache (Cache Array Int (CacheState t)
a) = (CacheState t -> IO ()) -> [CacheState t] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ CacheState t -> IO ()
forall t. CacheState t -> IO ()
finalizeCacheState (Array Int (CacheState t) -> [CacheState t]
forall i e. Array i e -> [e]
A.elems Array Int (CacheState t)
a)

-- | A class for interning terms.
class Interned t where
  data Description t
  type Uninterned t
  describe :: Uninterned t -> Description t
  identify :: CachedInfo -> Uninterned t -> t
  threadId :: t -> WeakThreadId
  descriptionDigest :: Description t -> Digest

{-# NOINLINE termCacheCell #-}
termCacheCell ::
  IORef
    ( HM.HashMap
        WeakThreadId
        ( WeakThreadIdRef,
          IORef (HM.HashMap Fingerprint (Cache Any))
        )
    )
termCacheCell :: IORef
  (HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
termCacheCell = IO
  (IORef
     (HashMap
        WeakThreadId
        (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))))
-> IORef
     (HashMap
        WeakThreadId
        (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
forall a. IO a -> a
unsafePerformIO (IO
   (IORef
      (HashMap
         WeakThreadId
         (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))))
 -> IORef
      (HashMap
         WeakThreadId
         (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))))
-> IO
     (IORef
        (HashMap
           WeakThreadId
           (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))))
-> IORef
     (HashMap
        WeakThreadId
        (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
forall a b. (a -> b) -> a -> b
$ HashMap
  WeakThreadId
  (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))
-> IO
     (IORef
        (HashMap
           WeakThreadId
           (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))))
forall a. a -> IO (IORef a)
newIORef HashMap
  WeakThreadId
  (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))
forall k v. HashMap k v
HM.empty

cacheWidth :: Word32
cacheWidth :: Word32
cacheWidth = Word32
10
{-# INLINE cacheWidth #-}

mkCache :: forall t. (Interned t) => IO (Cache t)
mkCache :: forall t. Interned t => IO (Cache t)
mkCache = IO (Cache t)
forall {t}. IO (Cache t)
result
  where
    element :: IO (CacheState t)
element =
      MVar ()
-> IOVector Word32
-> HashTable (Description t) (Word32, Weak StableIdent)
-> CacheState t
forall t.
MVar ()
-> IOVector Word32
-> HashTable (Description t) (Word32, Weak StableIdent)
-> CacheState t
CacheState
        (MVar ()
 -> IOVector Word32
 -> HashTable (Description t) (Word32, Weak StableIdent)
 -> CacheState t)
-> IO (MVar ())
-> IO
     (IOVector Word32
      -> HashTable (Description t) (Word32, Weak StableIdent)
      -> CacheState t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> () -> IO (MVar ())
forall a. a -> IO (MVar a)
newMVar ()
        IO
  (IOVector Word32
   -> HashTable (Description t) (Word32, Weak StableIdent)
   -> CacheState t)
-> IO (IOVector Word32)
-> IO
     (HashTable (Description t) (Word32, Weak StableIdent)
      -> CacheState t)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> Word32 -> IO (MVector (PrimState IO) Word32)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
M.replicate Int
1 Word32
0
        IO
  (HashTable (Description t) (Word32, Weak StableIdent)
   -> CacheState t)
-> IO (HashTable (Description t) (Word32, Weak StableIdent))
-> IO (CacheState t)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> HashMap (Description t) (Word32, Weak StableIdent)
-> IO (HashTable (Description t) (Word32, Weak StableIdent))
forall a. a -> IO (IORef a)
newIORef HashMap (Description t) (Word32, Weak StableIdent)
forall k v. HashMap k v
HM.empty
    result :: IO (Cache t)
result = do
      elements <- Int -> IO (CacheState t) -> IO [CacheState t]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
cacheWidth) IO (CacheState t)
forall {t}. IO (CacheState t)
element
      return $ Cache $ A.listArray (0, fromIntegral cacheWidth - 1) elements

-- | Internal cache for memoization of term construction. Different types have
-- different caches and they may share names, ids, or representations, but they
-- are not the same term.
typeMemoizedCache ::
  forall a. (Interned a) => ThreadId -> Fingerprint -> IO (Cache a)
typeMemoizedCache :: forall a. Interned a => ThreadId -> Fingerprint -> IO (Cache a)
typeMemoizedCache ThreadId
tid Fingerprint
tyFingerprint = do
  caches <- IORef
  (HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
-> IO
     (HashMap
        WeakThreadId
        (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
forall a. IORef a -> IO a
readIORef IORef
  (HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
termCacheCell
  let wtid = ThreadId -> WeakThreadId
weakThreadId ThreadId
tid
  case HM.lookup wtid caches of
    Just (WeakThreadIdRef
_, IORef (HashMap Fingerprint (Cache Any))
cref) -> do
      cache <- IORef (HashMap Fingerprint (Cache Any))
-> IO (HashMap Fingerprint (Cache Any))
forall a. IORef a -> IO a
readIORef IORef (HashMap Fingerprint (Cache Any))
cref
      case HM.lookup tyFingerprint cache of
        Just Cache Any
d -> Cache a -> IO (Cache a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Cache a -> IO (Cache a)) -> Cache a -> IO (Cache a)
forall a b. (a -> b) -> a -> b
$ Cache Any -> Cache a
forall a b. a -> b
unsafeCoerce Cache Any
d
        Maybe (Cache Any)
Nothing -> do
          r1 <- IO (Cache a)
forall t. Interned t => IO (Cache t)
mkCache
          writeIORef cref $!
            HM.insert tyFingerprint (unsafeCoerce r1) cache
          return r1
    Maybe (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))
Nothing -> do
      r1 <- IO (Cache a)
forall t. Interned t => IO (Cache t)
mkCache
      wtidRef <-
        mkWeakThreadIdRefWithFinalizer tid $ do
          finalizeCache r1
          atomicModifyIORefCAS_ termCacheCell (HM.delete wtid)
      r <- newIORef $ HM.singleton tyFingerprint (unsafeCoerce r1)
      atomicModifyIORefCAS termCacheCell $
        \HashMap
  WeakThreadId
  (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))
m -> (WeakThreadId
-> (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))
-> HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))
-> HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert WeakThreadId
wtid (WeakThreadIdRef
wtidRef, IORef (HashMap Fingerprint (Cache Any))
r) HashMap
  WeakThreadId
  (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))
m, Cache a
r1)

reclaimTerm ::
  forall t.
  (Interned t, Hashable (Description t), Eq (Description t)) =>
  WeakThreadId ->
  Fingerprint ->
  Int ->
  Description t ->
  IO ()
reclaimTerm :: forall t.
(Interned t, Hashable (Description t), Eq (Description t)) =>
WeakThreadId -> Fingerprint -> Int -> Description t -> IO ()
reclaimTerm WeakThreadId
id Fingerprint
tyFingerprint Int
grp Description t
dt = do
  caches <- IORef
  (HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
-> IO
     (HashMap
        WeakThreadId
        (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
forall a. IORef a -> IO a
readIORef IORef
  (HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
termCacheCell
  case HM.lookup id caches of
    Just (WeakThreadIdRef
_, IORef (HashMap Fingerprint (Cache Any))
cref) -> do
      cache <- IORef (HashMap Fingerprint (Cache Any))
-> IO (HashMap Fingerprint (Cache Any))
forall a. IORef a -> IO a
readIORef IORef (HashMap Fingerprint (Cache Any))
cref
      case HM.lookup tyFingerprint cache of
        Just Cache Any
c -> do
          let Cache Array Int (CacheState t)
a = Cache Any -> Cache t
forall a b. a -> b
unsafeCoerce Cache Any
c :: Cache t
          let CacheState MVar ()
sem IOVector Word32
_ HashTable (Description t) (Word32, Weak StableIdent)
s = Array Int (CacheState t)
a Array Int (CacheState t) -> Int -> CacheState t
forall i e. Ix i => Array i e -> i -> e
A.! Int
grp
          MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
sem
          current <- HashTable (Description t) (Word32, Weak StableIdent)
-> IO (HashMap (Description t) (Word32, Weak StableIdent))
forall a. IORef a -> IO a
readIORef HashTable (Description t) (Word32, Weak StableIdent)
s
          case HM.lookup dt current of
            Maybe (Word32, Weak StableIdent)
Nothing -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            Just (Word32
_, Weak StableIdent
wr) -> do
              t <- Weak StableIdent -> IO (Maybe StableIdent)
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak StableIdent
wr
              case t of
                Maybe StableIdent
Nothing -> HashTable (Description t) (Word32, Weak StableIdent)
-> HashMap (Description t) (Word32, Weak StableIdent) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef HashTable (Description t) (Word32, Weak StableIdent)
s (HashMap (Description t) (Word32, Weak StableIdent) -> IO ())
-> HashMap (Description t) (Word32, Weak StableIdent) -> IO ()
forall a b. (a -> b) -> a -> b
$ Description t
-> HashMap (Description t) (Word32, Weak StableIdent)
-> HashMap (Description t) (Word32, Weak StableIdent)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
HM.delete Description t
dt HashMap (Description t) (Word32, Weak StableIdent)
current
                Just StableIdent
_ -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
          putMVar sem ()
        Maybe (Cache Any)
Nothing -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Maybe (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))
Nothing -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Internalize a term.
intern ::
  forall t.
  (Interned t, Typeable t, Hashable (Description t), Eq (Description t)) =>
  Uninterned t ->
  IO t
intern :: forall t.
(Interned t, Typeable t, Hashable (Description t),
 Eq (Description t)) =>
Uninterned t -> IO t
intern !Uninterned t
bt = do
  tid <- IO ThreadId
myThreadId
  let wtid = ThreadId -> WeakThreadId
weakThreadId ThreadId
tid
  let fingerprint = TypeRep -> Fingerprint
typeRepFingerprint (TypeRep -> Fingerprint) -> TypeRep -> Fingerprint
forall a b. (a -> b) -> a -> b
$ Proxy t -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
someTypeRep (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @t)
  cache <- typeMemoizedCache tid fingerprint
  let !dt = Uninterned t -> Description t
forall t. Interned t => Uninterned t -> Description t
describe Uninterned t
bt :: Description t
      !hdt = Description t -> Word32
forall t. Interned t => Description t -> Word32
descriptionDigest Description t
dt
      !r = Word32
hdt Word32 -> Word32 -> Word32
forall a. Integral a => a -> a -> a
`mod` Word32
cacheWidth
      CacheState sem nextId s = getCache cache A.! (fromIntegral r)
  takeMVar sem
  -- print ("intern", wtid, dt, r)
  current <- readIORef s
  case HM.lookup dt current of
    Maybe (Word32, Weak StableIdent)
Nothing -> do
      newId0 <- MVector (PrimState IO) Word32 -> Int -> IO Word32
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
M.unsafeRead IOVector Word32
MVector (PrimState IO) Word32
nextId Int
0
      M.unsafeWrite nextId 0 (newId0 + 1)
      let newId = Word32
newId0 Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
* Word32
cacheWidth Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
r
      newIdent <- makeStableName dt
      let anyNewIdent = StableName (Description t) -> StableIdent
forall a b. a -> b
unsafeCoerce StableName (Description t)
newIdent :: StableIdent
      identRef <-
        mkWeakStableNameRefWithFinalizer anyNewIdent $
          reclaimTerm wtid fingerprint (fromIntegral r) dt
      let !t = CachedInfo -> Uninterned t -> t
forall t. Interned t => CachedInfo -> Uninterned t -> t
identify (WeakThreadId -> Word32 -> Word32 -> StableIdent -> CachedInfo
CachedInfo (ThreadId -> WeakThreadId
weakThreadId ThreadId
tid) Word32
hdt Word32
newId StableIdent
anyNewIdent) Uninterned t
bt
      writeIORef s $ HM.insert dt (newId, identRef) current
      putMVar sem ()
      return t
    Just (Word32
oldId, Weak StableIdent
oldIdentRef) -> do
      t1 <- Weak StableIdent -> IO (Maybe StableIdent)
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak StableIdent
oldIdentRef
      case t1 of
        Maybe StableIdent
Nothing -> do
          newId0 <- MVector (PrimState IO) Word32 -> Int -> IO Word32
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
M.unsafeRead IOVector Word32
MVector (PrimState IO) Word32
nextId Int
0
          M.unsafeWrite nextId 0 (newId0 + 1)
          let newId = Word32
newId0 Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
* Word32
cacheWidth Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
r
          newIdent <- makeStableName dt
          let anyNewIdent = StableName (Description t) -> StableIdent
forall a b. a -> b
unsafeCoerce StableName (Description t)
newIdent :: StableIdent
          identRef <-
            mkWeakStableNameRefWithFinalizer anyNewIdent $
              reclaimTerm wtid fingerprint (fromIntegral r) dt
          let !term =
                CachedInfo -> Uninterned t -> t
forall t. Interned t => CachedInfo -> Uninterned t -> t
identify
                  (WeakThreadId -> Word32 -> Word32 -> StableIdent -> CachedInfo
CachedInfo (ThreadId -> WeakThreadId
weakThreadId ThreadId
tid) Word32
hdt Word32
newId StableIdent
anyNewIdent)
                  Uninterned t
bt
          writeIORef s $ HM.insert dt (newId, identRef) current
          putMVar sem ()
          return term
        Just StableIdent
t1 -> do
          MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
sem ()
          t -> IO t
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (t -> IO t) -> t -> IO t
forall a b. (a -> b) -> a -> b
$! CachedInfo -> Uninterned t -> t
forall t. Interned t => CachedInfo -> Uninterned t -> t
identify (WeakThreadId -> Word32 -> Word32 -> StableIdent -> CachedInfo
CachedInfo (ThreadId -> WeakThreadId
weakThreadId ThreadId
tid) Word32
hdt Word32
oldId StableIdent
t1) Uninterned t
bt
{-# NOINLINE intern #-}

-- | Check if the current thread has a cache.
haveCache :: IO Bool
haveCache :: IO Bool
haveCache = do
  caches <- IORef
  (HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
-> IO
     (HashMap
        WeakThreadId
        (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
forall a. IORef a -> IO a
readIORef IORef
  (HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
termCacheCell
  tid <- myWeakThreadId
  return $ HM.member tid caches

cacheStateSize :: CacheState t -> IO Int
cacheStateSize :: forall t. CacheState t -> IO Int
cacheStateSize (CacheState MVar ()
_ IOVector Word32
_ HashTable (Description t) (Word32, Weak StableIdent)
s) = HashMap (Description t) (Word32, Weak StableIdent) -> Int
forall k v. HashMap k v -> Int
HM.size (HashMap (Description t) (Word32, Weak StableIdent) -> Int)
-> IO (HashMap (Description t) (Word32, Weak StableIdent))
-> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashTable (Description t) (Word32, Weak StableIdent)
-> IO (HashMap (Description t) (Word32, Weak StableIdent))
forall a. IORef a -> IO a
readIORef HashTable (Description t) (Word32, Weak StableIdent)
s

cacheStateLiveSize :: CacheState t -> IO Int
cacheStateLiveSize :: forall t. CacheState t -> IO Int
cacheStateLiveSize (CacheState MVar ()
sem IOVector Word32
_ HashTable (Description t) (Word32, Weak StableIdent)
s) = do
  MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
sem
  v <- ((Description t, (Word32, Weak StableIdent))
 -> (Word32, Weak StableIdent))
-> [(Description t, (Word32, Weak StableIdent))]
-> [(Word32, Weak StableIdent)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Description t, (Word32, Weak StableIdent))
-> (Word32, Weak StableIdent)
forall a b. (a, b) -> b
snd ([(Description t, (Word32, Weak StableIdent))]
 -> [(Word32, Weak StableIdent)])
-> (HashMap (Description t) (Word32, Weak StableIdent)
    -> [(Description t, (Word32, Weak StableIdent))])
-> HashMap (Description t) (Word32, Weak StableIdent)
-> [(Word32, Weak StableIdent)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashMap (Description t) (Word32, Weak StableIdent)
-> [(Description t, (Word32, Weak StableIdent))]
forall k v. HashMap k v -> [(k, v)]
HM.toList (HashMap (Description t) (Word32, Weak StableIdent)
 -> [(Word32, Weak StableIdent)])
-> IO (HashMap (Description t) (Word32, Weak StableIdent))
-> IO [(Word32, Weak StableIdent)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashTable (Description t) (Word32, Weak StableIdent)
-> IO (HashMap (Description t) (Word32, Weak StableIdent))
forall a. IORef a -> IO a
readIORef HashTable (Description t) (Word32, Weak StableIdent)
s
  r <-
    sum
      <$> mapM
        ( \(Word32
_, Weak StableIdent
x) -> do
            x <- Weak StableIdent -> IO (Maybe StableIdent)
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak StableIdent
x
            if isJust x then return 1 else return 0
        )
        v
  putMVar sem ()
  return r

{-
dumpCacheState :: CacheState t -> IO ()
dumpCacheState (CacheState sem s) = do
  takeMVar sem
  v <- HM.toList <$> readIORef s
  mapM_
    ( \(k, (i, v)) -> do
        v1 <- deRefWeak v
        case v1 of
          Nothing -> print (k, i, "dead")
          Just r -> print (k, i, r)
    )
    v
  putMVar sem ()

dumpCache :: Cache t -> IO ()
dumpCache (Cache a) = mapM_ dumpCacheState (A.elems a)
-}

cacheSize :: Cache t -> IO Int
cacheSize :: forall t. Cache t -> IO Int
cacheSize (Cache Array Int (CacheState t)
a) = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> IO [Int] -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (CacheState t -> IO Int) -> [CacheState t] -> IO [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM CacheState t -> IO Int
forall t. CacheState t -> IO Int
cacheStateSize (Array Int (CacheState t) -> [CacheState t]
forall i e. Array i e -> [e]
A.elems Array Int (CacheState t)
a)

cacheLiveSize :: Cache t -> IO Int
cacheLiveSize :: forall t. Cache t -> IO Int
cacheLiveSize (Cache Array Int (CacheState t)
a) = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> IO [Int] -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (CacheState t -> IO Int) -> [CacheState t] -> IO [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM CacheState t -> IO Int
forall t. CacheState t -> IO Int
cacheStateLiveSize (Array Int (CacheState t) -> [CacheState t]
forall i e. Array i e -> [e]
A.elems Array Int (CacheState t)
a)

-- | Get the size of the current thread's cache.
threadCacheSize :: WeakThreadId -> IO Int
threadCacheSize :: WeakThreadId -> IO Int
threadCacheSize WeakThreadId
tid = do
  caches <- IORef
  (HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
-> IO
     (HashMap
        WeakThreadId
        (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
forall a. IORef a -> IO a
readIORef IORef
  (HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
termCacheCell
  case HM.lookup tid caches of
    Just (WeakThreadIdRef
_, IORef (HashMap Fingerprint (Cache Any))
cref) -> do
      cache <- IORef (HashMap Fingerprint (Cache Any))
-> IO (HashMap Fingerprint (Cache Any))
forall a. IORef a -> IO a
readIORef IORef (HashMap Fingerprint (Cache Any))
cref
      sum <$> mapM cacheSize (HM.elems cache)
    Maybe (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))
Nothing -> Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0

-- | Get the live size of the current thread's cache.
threadCacheLiveSize :: WeakThreadId -> IO Int
threadCacheLiveSize :: WeakThreadId -> IO Int
threadCacheLiveSize WeakThreadId
tid = do
  caches <- IORef
  (HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
-> IO
     (HashMap
        WeakThreadId
        (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
forall a. IORef a -> IO a
readIORef IORef
  (HashMap
     WeakThreadId
     (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any))))
termCacheCell
  case HM.lookup tid caches of
    Just (WeakThreadIdRef
_, IORef (HashMap Fingerprint (Cache Any))
cref) -> do
      cache <- IORef (HashMap Fingerprint (Cache Any))
-> IO (HashMap Fingerprint (Cache Any))
forall a. IORef a -> IO a
readIORef IORef (HashMap Fingerprint (Cache Any))
cref
      sum <$> mapM cacheLiveSize (HM.elems cache)
    Maybe (WeakThreadIdRef, IORef (HashMap Fingerprint (Cache Any)))
Nothing -> Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0

{-
-- | Dump the current thread's cache.
dumpThreadCache :: WeakThreadId -> IO ()
dumpThreadCache tid = do
  caches <- readIORef termCacheCell
  case HM.lookup tid caches of
    Just (_, cref) -> do
      cache <- readIORef cref
      mapM_ dumpCache (HM.elems cache)
    Nothing -> return ()
-}