{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

{- |
   Module      : GHC.DataSize
   Copyright   : (c) Dennis Felsing
   License     : 3-Clause BSD-style
   Maintainer  : dennis@felsin9.de
 -}
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

-- Inspired by Simon Marlow:
-- https://ghcmutterings.wordpress.com/2009/02/12/53/

-- | Calculate size of GHC objects in Bytes. Note that an object may not be
--   evaluated yet and only the size of the initial closure is returned.
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

-- | Calculate the recursive size of GHC objects in Bytes. Note that the actual
--   size in memory is calculated, so shared values are only counted once.
--
--   Call with
--   @
--    recursiveSize $! 2
--   @
--   to force evaluation to WHNF before calculating the size.
--
--   Call with
--   @
--    recursiveSize $!! \"foobar\"
--   @
--   ($!! from Control.DeepSeq) to force full evaluation before calculating the
--   size.
--
--   A garbage collection is performed before the size is calculated, because
--   the garbage collector would make heap walks difficult.
--
--   This function works very quickly on small data structures, but can be slow
--   on large and complex ones. If speed is an issue it's probably possible to
--   get the exact size of a small portion of the data structure and then
--   estimate the total size from that.

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

-- | Calculate the recursive size of GHC objects in Bytes after calling
-- Control.DeepSeq.force on the data structure to force it into Normal Form.
-- Using this function requires that the data structure has an `NFData`
-- typeclass instance.

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

-- | Adapted from 'GHC.Exts.Heap.getClosureRaw' which isn't exported.
--
-- This returns the raw words of the closure on the heap. Once back in the
-- Haskell world, the raw words that hold pointers may be outdated after a
-- garbage collector run.
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] ]