{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE GHCForeignImportPrim #-}
{-# LANGUAGE UnliftedFFITypes #-}
module HeapSize (
recursiveSize,
recursiveSizeNF,
closureSize,
Heapsize,
runHeapsize
)
where
import Control.DeepSeq (NFData, force)
import Control.Exception (throwIO)
import Control.Monad
import Control.Monad.Catch
import Control.Monad.IO.Class
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.Reader
import Data.IORef
import Data.Hashable
import qualified Data.HashTable.IO as HT
import Data.Maybe (isJust, isNothing)
import Data.Typeable (Typeable)
import GHC.Exts hiding (closureSize#)
import GHC.Arr
import GHC.Exts.Heap hiding (size)
import qualified Data.Foldable as F
import System.Mem
import System.Mem.Weak
import Debug.Trace
foreign import prim "aToWordzh" aToWord# :: Any -> Word#
foreign import prim "unpackClosurePtrs" unpackClosurePtrs# :: Any -> Array# b
foreign import prim "closureSize" closureSize# :: Any -> Int#
closureSize :: a -> IO Int
closureSize :: a -> IO Int
closureSize a
x = Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int# -> Int
I# (Any -> Int#
closureSize# (a -> Any
unsafeCoerce# a
x)))
getClosures :: a -> IO (Array Int Box)
getClosures :: a -> IO (Array Int Box)
getClosures a
x = case Any -> Array# Any
forall b. Any -> Array# b
unpackClosurePtrs# (a -> Any
unsafeCoerce# a
x) of
Array# Any
pointers ->
let nelems :: Int
nelems = Int# -> Int
I# (Array# Any -> Int#
forall k1. Array# k1 -> Int#
sizeofArray# Array# Any
pointers)
in Array Int Box -> IO (Array Int Box)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Any -> Box) -> Array Int Any -> Array Int Box
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Any -> Box
Box (Array Int Any -> Array Int Box) -> Array Int Any -> Array Int Box
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Array# Any -> Array Int Any
forall i e. i -> i -> Int -> Array# e -> Array i e
Array Int
0 (Int
nelems Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
nelems Array# Any
pointers)
data HeapsizeState = HeapsizeState
{
HeapsizeState -> BasicHashTable HashableBox ()
closuresSeen :: HT.BasicHashTable HashableBox (),
HeapsizeState -> IORef Int
seenSizeRef :: IORef Int,
HeapsizeState -> GcDetector
gcDetect :: GcDetector
}
newtype GcDetector = GcDetector {GcDetector -> IO Bool
gcSinceCreation :: IO Bool}
gcDetector :: IO GcDetector
gcDetector :: IO GcDetector
gcDetector = do
IORef ()
ref <- () -> IO (IORef ())
forall a. a -> IO (IORef a)
newIORef ()
Weak (IORef ())
w <- IORef () -> IO () -> IO (Weak (IORef ()))
forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef ()
ref (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
GcDetector -> IO GcDetector
forall (m :: * -> *) a. Monad m => a -> m a
return (GcDetector -> IO GcDetector) -> GcDetector -> IO GcDetector
forall a b. (a -> b) -> a -> b
$ IO Bool -> GcDetector
GcDetector (IO Bool -> GcDetector) -> IO Bool -> GcDetector
forall a b. (a -> b) -> a -> b
$ Maybe (IORef ()) -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe (IORef ()) -> Bool) -> IO (Maybe (IORef ())) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Weak (IORef ()) -> IO (Maybe (IORef ()))
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak (IORef ())
w
newtype Heapsize a = Heapsize
{ Heapsize a -> ReaderT HeapsizeState (MaybeT IO) a
_unHeapsize :: ReaderT HeapsizeState (MaybeT IO) a}
deriving (Functor Heapsize
a -> Heapsize a
Functor Heapsize
-> (forall a. a -> Heapsize a)
-> (forall a b. Heapsize (a -> b) -> Heapsize a -> Heapsize b)
-> (forall a b c.
(a -> b -> c) -> Heapsize a -> Heapsize b -> Heapsize c)
-> (forall a b. Heapsize a -> Heapsize b -> Heapsize b)
-> (forall a b. Heapsize a -> Heapsize b -> Heapsize a)
-> Applicative Heapsize
Heapsize a -> Heapsize b -> Heapsize b
Heapsize a -> Heapsize b -> Heapsize a
Heapsize (a -> b) -> Heapsize a -> Heapsize b
(a -> b -> c) -> Heapsize a -> Heapsize b -> Heapsize c
forall a. a -> Heapsize a
forall a b. Heapsize a -> Heapsize b -> Heapsize a
forall a b. Heapsize a -> Heapsize b -> Heapsize b
forall a b. Heapsize (a -> b) -> Heapsize a -> Heapsize b
forall a b c.
(a -> b -> c) -> Heapsize a -> Heapsize b -> Heapsize c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: Heapsize a -> Heapsize b -> Heapsize a
$c<* :: forall a b. Heapsize a -> Heapsize b -> Heapsize a
*> :: Heapsize a -> Heapsize b -> Heapsize b
$c*> :: forall a b. Heapsize a -> Heapsize b -> Heapsize b
liftA2 :: (a -> b -> c) -> Heapsize a -> Heapsize b -> Heapsize c
$cliftA2 :: forall a b c.
(a -> b -> c) -> Heapsize a -> Heapsize b -> Heapsize c
<*> :: Heapsize (a -> b) -> Heapsize a -> Heapsize b
$c<*> :: forall a b. Heapsize (a -> b) -> Heapsize a -> Heapsize b
pure :: a -> Heapsize a
$cpure :: forall a. a -> Heapsize a
$cp1Applicative :: Functor Heapsize
Applicative, a -> Heapsize b -> Heapsize a
(a -> b) -> Heapsize a -> Heapsize b
(forall a b. (a -> b) -> Heapsize a -> Heapsize b)
-> (forall a b. a -> Heapsize b -> Heapsize a) -> Functor Heapsize
forall a b. a -> Heapsize b -> Heapsize a
forall a b. (a -> b) -> Heapsize a -> Heapsize b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> Heapsize b -> Heapsize a
$c<$ :: forall a b. a -> Heapsize b -> Heapsize a
fmap :: (a -> b) -> Heapsize a -> Heapsize b
$cfmap :: forall a b. (a -> b) -> Heapsize a -> Heapsize b
Functor, Applicative Heapsize
a -> Heapsize a
Applicative Heapsize
-> (forall a b. Heapsize a -> (a -> Heapsize b) -> Heapsize b)
-> (forall a b. Heapsize a -> Heapsize b -> Heapsize b)
-> (forall a. a -> Heapsize a)
-> Monad Heapsize
Heapsize a -> (a -> Heapsize b) -> Heapsize b
Heapsize a -> Heapsize b -> Heapsize b
forall a. a -> Heapsize a
forall a b. Heapsize a -> Heapsize b -> Heapsize b
forall a b. Heapsize a -> (a -> Heapsize b) -> Heapsize b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: a -> Heapsize a
$creturn :: forall a. a -> Heapsize a
>> :: Heapsize a -> Heapsize b -> Heapsize b
$c>> :: forall a b. Heapsize a -> Heapsize b -> Heapsize b
>>= :: Heapsize a -> (a -> Heapsize b) -> Heapsize b
$c>>= :: forall a b. Heapsize a -> (a -> Heapsize b) -> Heapsize b
$cp1Monad :: Applicative Heapsize
Monad, Monad Heapsize
Monad Heapsize
-> (forall a. IO a -> Heapsize a) -> MonadIO Heapsize
IO a -> Heapsize a
forall a. IO a -> Heapsize a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
liftIO :: IO a -> Heapsize a
$cliftIO :: forall a. IO a -> Heapsize a
$cp1MonadIO :: Monad Heapsize
MonadIO, MonadThrow Heapsize
MonadThrow Heapsize
-> (forall e a.
Exception e =>
Heapsize a -> (e -> Heapsize a) -> Heapsize a)
-> MonadCatch Heapsize
Heapsize a -> (e -> Heapsize a) -> Heapsize a
forall e a.
Exception e =>
Heapsize a -> (e -> Heapsize a) -> Heapsize a
forall (m :: * -> *).
MonadThrow m
-> (forall e a. Exception e => m a -> (e -> m a) -> m a)
-> MonadCatch m
catch :: Heapsize a -> (e -> Heapsize a) -> Heapsize a
$ccatch :: forall e a.
Exception e =>
Heapsize a -> (e -> Heapsize a) -> Heapsize a
$cp1MonadCatch :: MonadThrow Heapsize
MonadCatch, MonadCatch Heapsize
MonadCatch Heapsize
-> (forall b.
((forall a. Heapsize a -> Heapsize a) -> Heapsize b) -> Heapsize b)
-> (forall b.
((forall a. Heapsize a -> Heapsize a) -> Heapsize b) -> Heapsize b)
-> (forall a b c.
Heapsize a
-> (a -> ExitCase b -> Heapsize c)
-> (a -> Heapsize b)
-> Heapsize (b, c))
-> MonadMask Heapsize
Heapsize a
-> (a -> ExitCase b -> Heapsize c)
-> (a -> Heapsize b)
-> Heapsize (b, c)
((forall a. Heapsize a -> Heapsize a) -> Heapsize b) -> Heapsize b
((forall a. Heapsize a -> Heapsize a) -> Heapsize b) -> Heapsize b
forall b.
((forall a. Heapsize a -> Heapsize a) -> Heapsize b) -> Heapsize b
forall a b c.
Heapsize a
-> (a -> ExitCase b -> Heapsize c)
-> (a -> Heapsize b)
-> Heapsize (b, c)
forall (m :: * -> *).
MonadCatch m
-> (forall b. ((forall a. m a -> m a) -> m b) -> m b)
-> (forall b. ((forall a. m a -> m a) -> m b) -> m b)
-> (forall a b c.
m a -> (a -> ExitCase b -> m c) -> (a -> m b) -> m (b, c))
-> MonadMask m
generalBracket :: Heapsize a
-> (a -> ExitCase b -> Heapsize c)
-> (a -> Heapsize b)
-> Heapsize (b, c)
$cgeneralBracket :: forall a b c.
Heapsize a
-> (a -> ExitCase b -> Heapsize c)
-> (a -> Heapsize b)
-> Heapsize (b, c)
uninterruptibleMask :: ((forall a. Heapsize a -> Heapsize a) -> Heapsize b) -> Heapsize b
$cuninterruptibleMask :: forall b.
((forall a. Heapsize a -> Heapsize a) -> Heapsize b) -> Heapsize b
mask :: ((forall a. Heapsize a -> Heapsize a) -> Heapsize b) -> Heapsize b
$cmask :: forall b.
((forall a. Heapsize a -> Heapsize a) -> Heapsize b) -> Heapsize b
$cp1MonadMask :: MonadCatch Heapsize
MonadMask, Monad Heapsize
e -> Heapsize a
Monad Heapsize
-> (forall e a. Exception e => e -> Heapsize a)
-> MonadThrow Heapsize
forall e a. Exception e => e -> Heapsize a
forall (m :: * -> *).
Monad m -> (forall e a. Exception e => e -> m a) -> MonadThrow m
throwM :: e -> Heapsize a
$cthrowM :: forall e a. Exception e => e -> Heapsize a
$cp1MonadThrow :: Monad Heapsize
MonadThrow)
runHeapsize :: Int -> Heapsize a -> IO (Maybe a)
runHeapsize :: Int -> Heapsize a -> IO (Maybe a)
runHeapsize Int
initSize (Heapsize ReaderT HeapsizeState (MaybeT IO) a
comp) = do
!HashTable RealWorld HashableBox ()
closuresSeen <- Int -> IO (BasicHashTable HashableBox ())
forall (h :: * -> * -> * -> *) k v.
HashTable h =>
Int -> IO (IOHashTable h k v)
HT.newSized Int
initSize
IORef Int
seenSizeRef <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0
IO ()
performMajorGC
GcDetector
gcDetect <- IO GcDetector
gcDetector
MaybeT IO a -> IO (Maybe a)
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT IO a -> IO (Maybe a)) -> MaybeT IO a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ ReaderT HeapsizeState (MaybeT IO) a -> HeapsizeState -> MaybeT IO a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT HeapsizeState (MaybeT IO) a
comp HeapsizeState :: BasicHashTable HashableBox ()
-> IORef Int -> GcDetector -> HeapsizeState
HeapsizeState{IORef Int
HashTable RealWorld HashableBox ()
BasicHashTable HashableBox ()
GcDetector
gcDetect :: GcDetector
seenSizeRef :: IORef Int
closuresSeen :: HashTable RealWorld HashableBox ()
gcDetect :: GcDetector
seenSizeRef :: IORef Int
closuresSeen :: BasicHashTable HashableBox ()
..}
recursiveSize :: a -> Heapsize Int
recursiveSize :: a -> Heapsize Int
recursiveSize a
x = ReaderT HeapsizeState (MaybeT IO) Int -> Heapsize Int
forall a. ReaderT HeapsizeState (MaybeT IO) a -> Heapsize a
Heapsize (ReaderT HeapsizeState (MaybeT IO) Int -> Heapsize Int)
-> ReaderT HeapsizeState (MaybeT IO) Int -> Heapsize Int
forall a b. (a -> b) -> a -> b
$ do
HeapsizeState{IORef Int
BasicHashTable HashableBox ()
GcDetector
gcDetect :: GcDetector
seenSizeRef :: IORef Int
closuresSeen :: BasicHashTable HashableBox ()
gcDetect :: HeapsizeState -> GcDetector
seenSizeRef :: HeapsizeState -> IORef Int
closuresSeen :: HeapsizeState -> BasicHashTable HashableBox ()
..} <- ReaderT HeapsizeState (MaybeT IO) HeapsizeState
forall (m :: * -> *) r. Monad m => ReaderT r m r
ask
IORef Int
accSizeRef <- IO (IORef Int) -> ReaderT HeapsizeState (MaybeT IO) (IORef Int)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef Int) -> ReaderT HeapsizeState (MaybeT IO) (IORef Int))
-> IO (IORef Int) -> ReaderT HeapsizeState (MaybeT IO) (IORef Int)
forall a b. (a -> b) -> a -> b
$ Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0
let checkGC :: IO ()
checkGC = GcDetector -> IO Bool
gcSinceCreation GcDetector
gcDetect IO Bool -> (Bool -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
abort -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
abort (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Interrupted -> IO ()
forall e a. Exception e => e -> IO a
throwIO Interrupted
Interrupted
let
go :: [ Box ] -> IO ()
go :: [Box] -> IO ()
go [] = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
go (b :: Box
b@(Box Any
y) : [Box]
rest) = do
let addr :: Word
addr = Word# -> Word
W# (Any -> Word#
aToWord# Any
y)
!Bool
seen <- Maybe () -> Bool
forall a. Maybe a -> Bool
isJust (Maybe () -> Bool) -> IO (Maybe ()) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BasicHashTable HashableBox () -> HashableBox -> IO (Maybe ())
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> IO (Maybe v)
HT.lookup BasicHashTable HashableBox ()
closuresSeen (Box -> HashableBox
HashableBox Box
b)
[Box]
next <- if Bool
seen Bool -> Bool -> Bool
|| Word -> Bool
isBadAddress Word
addr then [Box] -> IO [Box]
forall (m :: * -> *) a. Monad m => a -> m a
return [] else do
IO ()
checkGC
Int
thisSize <- Any -> IO Int
forall a. a -> IO Int
closureSize Any
y
IO ()
checkGC
Array Int Box
next <- Any -> IO (Array Int Box)
forall a. a -> IO (Array Int Box)
getClosures Any
y
BasicHashTable HashableBox () -> HashableBox -> () -> IO ()
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> v -> IO ()
HT.insert BasicHashTable HashableBox ()
closuresSeen (Box -> HashableBox
HashableBox Box
b) ()
IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
accSizeRef (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
thisSize)
IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
seenSizeRef Int -> Int
forall a. Enum a => a -> a
succ
[Box] -> IO [Box]
forall (m :: * -> *) a. Monad m => a -> m a
return (Array Int Box -> [Box]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList Array Int Box
next)
[Box] -> IO ()
go ([Box]
next [Box] -> [Box] -> [Box]
forall a. [a] -> [a] -> [a]
++ [Box]
rest)
IO () -> ReaderT HeapsizeState (MaybeT IO) ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO ([Box] -> IO ()
go [a -> Box
forall a. a -> Box
asBox a
x]) ReaderT HeapsizeState (MaybeT IO) ()
-> (Interrupted -> ReaderT HeapsizeState (MaybeT IO) ())
-> ReaderT HeapsizeState (MaybeT IO) ()
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> (e -> m a) -> m a
`catch` \Interrupted
Interrupted -> do
Int
seen <- IO Int -> ReaderT HeapsizeState (MaybeT IO) Int
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> ReaderT HeapsizeState (MaybeT IO) Int)
-> IO Int -> ReaderT HeapsizeState (MaybeT IO) Int
forall a b. (a -> b) -> a -> b
$ IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
seenSizeRef
IO () -> ReaderT HeapsizeState (MaybeT IO) ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ReaderT HeapsizeState (MaybeT IO) ())
-> IO () -> ReaderT HeapsizeState (MaybeT IO) ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
traceIO (String
"SEEN: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
seen)
ReaderT HeapsizeState (MaybeT IO) ()
forall (m :: * -> *) a. MonadPlus m => m a
mzero
IO Int -> ReaderT HeapsizeState (MaybeT IO) Int
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
accSizeRef)
data Interrupted = Interrupted deriving (Int -> Interrupted -> String -> String
[Interrupted] -> String -> String
Interrupted -> String
(Int -> Interrupted -> String -> String)
-> (Interrupted -> String)
-> ([Interrupted] -> String -> String)
-> Show Interrupted
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
showList :: [Interrupted] -> String -> String
$cshowList :: [Interrupted] -> String -> String
show :: Interrupted -> String
$cshow :: Interrupted -> String
showsPrec :: Int -> Interrupted -> String -> String
$cshowsPrec :: Int -> Interrupted -> String -> String
Show, Typeable)
instance Exception Interrupted
recursiveSizeNF :: NFData a => a -> Heapsize Int
recursiveSizeNF :: a -> Heapsize Int
recursiveSizeNF = a -> Heapsize Int
forall a. a -> Heapsize Int
recursiveSize (a -> Heapsize Int) -> (a -> a) -> a -> Heapsize Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
forall a. NFData a => a -> a
force
isBadAddress :: Word -> Bool
isBadAddress :: Word -> Bool
isBadAddress Word
0 = Bool
True
isBadAddress Word
_ = Bool
False
newtype HashableBox = HashableBox Box
deriving newtype Int -> HashableBox -> String -> String
[HashableBox] -> String -> String
HashableBox -> String
(Int -> HashableBox -> String -> String)
-> (HashableBox -> String)
-> ([HashableBox] -> String -> String)
-> Show HashableBox
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
showList :: [HashableBox] -> String -> String
$cshowList :: [HashableBox] -> String -> String
show :: HashableBox -> String
$cshow :: HashableBox -> String
showsPrec :: Int -> HashableBox -> String -> String
$cshowsPrec :: Int -> HashableBox -> String -> String
Show
instance Eq HashableBox where
(HashableBox (Box Any
a1)) == :: HashableBox -> HashableBox -> Bool
== (HashableBox (Box Any
a2)) =
Word# -> Word
W# (Any -> Word#
aToWord# Any
a1) Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word# -> Word
W# (Any -> Word#
aToWord# Any
a2)
instance Hashable HashableBox where
hashWithSalt :: Int -> HashableBox -> Int
hashWithSalt Int
n (HashableBox (Box Any
a)) = Int -> Word -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
n (Word# -> Word
W# (Any -> Word#
aToWord# Any
a))