{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
module GHC.DataSize (
closureSize,
recursiveSize,
recursiveSizeNF
)
where
import Control.DeepSeq (NFData, ($!!))
import GHC.Exts
import GHC.Exts.Heap hiding (size)
import GHC.Exts.Heap.Constants (wORD_SIZE)
import Control.Monad
import System.Mem
closureSize :: a -> IO Word
closureSize :: a -> IO Word
closureSize a
x = do
[Word]
rawWds <- a -> IO [Word]
forall a. a -> IO [Word]
getClosureRawWords a
x
Word -> IO Word
forall (m :: * -> *) a. Monad m => a -> m a
return (Word -> IO Word) -> (Int -> Word) -> Int -> IO Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> IO Word) -> Int -> IO Word
forall a b. (a -> b) -> a -> b
$ [Word] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word]
rawWds Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
wORD_SIZE
recursiveSize :: a -> IO Word
recursiveSize :: a -> IO Word
recursiveSize a
x = do
IO ()
performGC
(([Box], Word) -> Word) -> IO ([Box], Word) -> IO Word
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM ([Box], Word) -> Word
forall a b. (a, b) -> b
snd (IO ([Box], Word) -> IO Word) -> IO ([Box], Word) -> IO Word
forall a b. (a -> b) -> a -> b
$ ([Box], Word) -> Box -> IO ([Box], Word)
go ([], Word
0) (Box -> IO ([Box], Word)) -> Box -> IO ([Box], Word)
forall a b. (a -> b) -> a -> b
$ a -> Box
forall a. a -> Box
asBox a
x
where go :: ([Box], Word) -> Box -> IO ([Box], Word)
go (![Box]
vs, !Word
acc) b :: Box
b@(Box Any
y) = do
Bool
isElem <- ([Bool] -> Bool) -> IO [Bool] -> IO Bool
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or (IO [Bool] -> IO Bool) -> IO [Bool] -> IO Bool
forall a b. (a -> b) -> a -> b
$ (Box -> IO Bool) -> [Box] -> IO [Bool]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Box -> Box -> IO Bool
areBoxesEqual Box
b) [Box]
vs
if Bool
isElem
then ([Box], Word) -> IO ([Box], Word)
forall (m :: * -> *) a. Monad m => a -> m a
return ([Box]
vs, Word
acc)
else do
Word
size <- Any -> IO Word
forall a. a -> IO Word
closureSize Any
y
Closure
closure <- Any -> IO Closure
forall a. HasHeapRep a => a -> IO Closure
getClosureData Any
y
(([Box], Word) -> Box -> IO ([Box], Word))
-> ([Box], Word) -> [Box] -> IO ([Box], Word)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ([Box], Word) -> Box -> IO ([Box], Word)
go (Box
b Box -> [Box] -> [Box]
forall a. a -> [a] -> [a]
: [Box]
vs, Word
acc Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
size) ([Box] -> IO ([Box], Word)) -> [Box] -> IO ([Box], Word)
forall a b. (a -> b) -> a -> b
$ Closure -> [Box]
forall b. GenClosure b -> [b]
allClosures Closure
closure
recursiveSizeNF :: NFData a => a -> IO Word
recursiveSizeNF :: a -> IO Word
recursiveSizeNF a
x = a -> IO Word
forall a. a -> IO Word
recursiveSize (a -> IO Word) -> a -> IO Word
forall a b. NFData a => (a -> b) -> a -> b
$!! a
x
getClosureRawWords :: a -> IO [Word]
getClosureRawWords :: a -> IO [Word]
getClosureRawWords a
x = do
case a -> (# Addr#, ByteArray#, Array# Any #)
forall a b. a -> (# Addr#, ByteArray#, Array# b #)
unpackClosure# a
x of
(# Addr#
_iptr, ByteArray#
dat, Array# Any
_pointers #) -> do
let nelems :: Int
nelems = (Int# -> Int
I# (ByteArray# -> Int#
sizeofByteArray# ByteArray#
dat)) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
wORD_SIZE
end :: Int
end = Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
nelems Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
[Word] -> IO [Word]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Word# -> Word
W# (ByteArray# -> Int# -> Word#
indexWordArray# ByteArray#
dat Int#
i) | I# Int#
i <- [Int
0.. Int
end] ]