-- Copyright 2013 Evan Laforge
-- This program is distributed under the terms of the GNU General Public
-- License 3.0, see COPYING or http://www.gnu.org/licenses/gpl-3.0.txt

{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-- | Utilities dealing with memory usage.
--
-- TODO the SizeOf part can maybe be replaced by the weigh package.
module Util.Memory where
import qualified Data.Map as Map
import qualified Data.Word as Word

import qualified Foreign
import qualified GHC.Stats
import qualified System.Mem
import qualified System.Posix.Process as Posix.Process
import qualified System.Process as Process

import qualified Util.Num as Num
import qualified Util.Pretty as Pretty


-- * Size

-- | Size in bytes.
newtype Size = Size Integer deriving (Size -> Size -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Size -> Size -> Bool
$c/= :: Size -> Size -> Bool
== :: Size -> Size -> Bool
$c== :: Size -> Size -> Bool
Eq, Eq Size
Size -> Size -> Bool
Size -> Size -> Ordering
Size -> Size -> Size
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Size -> Size -> Size
$cmin :: Size -> Size -> Size
max :: Size -> Size -> Size
$cmax :: Size -> Size -> Size
>= :: Size -> Size -> Bool
$c>= :: Size -> Size -> Bool
> :: Size -> Size -> Bool
$c> :: Size -> Size -> Bool
<= :: Size -> Size -> Bool
$c<= :: Size -> Size -> Bool
< :: Size -> Size -> Bool
$c< :: Size -> Size -> Bool
compare :: Size -> Size -> Ordering
$ccompare :: Size -> Size -> Ordering
Ord, Integer -> Size
Size -> Size
Size -> Size -> Size
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
fromInteger :: Integer -> Size
$cfromInteger :: Integer -> Size
signum :: Size -> Size
$csignum :: Size -> Size
abs :: Size -> Size
$cabs :: Size -> Size
negate :: Size -> Size
$cnegate :: Size -> Size
* :: Size -> Size -> Size
$c* :: Size -> Size -> Size
- :: Size -> Size -> Size
$c- :: Size -> Size -> Size
+ :: Size -> Size -> Size
$c+ :: Size -> Size -> Size
Num, Int -> Size -> ShowS
[Size] -> ShowS
Size -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Size] -> ShowS
$cshowList :: [Size] -> ShowS
show :: Size -> String
$cshow :: Size -> String
showsPrec :: Int -> Size -> ShowS
$cshowsPrec :: Int -> Size -> ShowS
Show)

instance Pretty.Pretty Size where
    pretty :: Size -> Text
pretty (Size Integer
n) =
        forall a. Pretty a => a -> Text
Pretty.pretty (forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n forall a. Fractional a => a -> a -> a
/ Double
1024 forall a. Fractional a => a -> a -> a
/ Double
1024 :: Double) forall a. Semigroup a => a -> a -> a
<> Text
"m"

fromBytes :: Integral a => a -> Size
fromBytes :: forall a. Integral a => a -> Size
fromBytes = Integer -> Size
Size forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral

toBytes :: Size -> Integer
toBytes :: Size -> Integer
toBytes (Size Integer
b) = Integer
b

fromK :: RealFrac a => a -> Size
fromK :: forall a. RealFrac a => a -> Size
fromK = forall a. Integral a => a -> Size
fromBytes forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (RealFrac a, Integral b) => a -> b
round forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Num a => a -> a -> a
*a
1024)

toK :: Size -> Double
toK :: Size -> Double
toK = (forall a. Fractional a => a -> a -> a
/Double
1024) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. Size -> Integer
toBytes

fromM :: RealFrac a => a -> Size
fromM :: forall a. RealFrac a => a -> Size
fromM = forall a. RealFrac a => a -> Size
fromK forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Num a => a -> a -> a
*a
1024)

toM :: Size -> Double
toM :: Size -> Double
toM = (forall a. Fractional a => a -> a -> a
/Double
1024) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Size -> Double
toK

fromWords :: Integral a => a -> Size
fromWords :: forall a. Integral a => a -> Size
fromWords = forall a. Integral a => a -> Size
fromBytes forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Storable a => a -> Int
Foreign.sizeOf forall a. Ptr a
Foreign.nullPtr))

-- * SizeOf

class SizeOf a where
    sizeOf :: a -> Size

instance SizeOf a => SizeOf [a] where
    sizeOf :: [a] -> Size
sizeOf [a]
xs = forall a. Integral a => a -> Size
fromBytes (forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) forall a. Num a => a -> a -> a
* Size
taggedBox forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
Num.sum (forall a b. (a -> b) -> [a] -> [b]
map forall a. SizeOf a => a -> Size
sizeOf [a]
xs)

instance SizeOf Char where
    sizeOf :: Char -> Size
sizeOf Char
_ = forall a. Integral a => a -> Size
fromWords Integer
0 -- chars from 0--255 are interned

instance SizeOf Int where sizeOf :: Int -> Size
sizeOf = forall a. Storable a => a -> Size
boxedStorable
instance SizeOf Word.Word8 where sizeOf :: Word8 -> Size
sizeOf = forall a. Storable a => a -> Size
boxedStorable
instance SizeOf Double where sizeOf :: Double -> Size
sizeOf = forall a. Storable a => a -> Size
boxedStorable
instance SizeOf Float where sizeOf :: Float -> Size
sizeOf = forall a. Storable a => a -> Size
boxedStorable

instance (SizeOf k, SizeOf a) => SizeOf (Map.Map k a) where
    sizeOf :: Map k a -> Size
sizeOf Map k a
m = forall a. Integral a => a -> Size
fromBytes Int
nodes forall a. Num a => a -> a -> a
* Size
taggedBox forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
Num.sum (forall a b. (a -> b) -> [a] -> [b]
map forall a. SizeOf a => a -> Size
sizeOf (forall k a. Map k a -> [k]
Map.keys Map k a
m))
            forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
Num.sum (forall a b. (a -> b) -> [a] -> [b]
map forall a. SizeOf a => a -> Size
sizeOf (forall k a. Map k a -> [a]
Map.elems Map k a
m))
        where
        nodes :: Int
nodes = forall k a. Map k a -> Int
Map.size Map k a
m forall a. Num a => a -> a -> a
+ forall a b. (RealFrac a, Integral b) => a -> b
ceiling (forall a. Floating a => a -> a -> a
logBase Double
2 (forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall k a. Map k a -> Int
Map.size Map k a
m)))

taggedBox :: Size
taggedBox :: Size
taggedBox = forall a. Integral a => a -> Size
fromWords Integer
3 -- gc overhead, tag, ptr

box :: Size
box :: Size
box = forall a. Integral a => a -> Size
fromWords Integer
2 -- gc overhead, ptr

-- | Boxed word-sized datatype.
boxedStorable :: (Foreign.Storable a) => a -> Size
boxedStorable :: forall a. Storable a => a -> Size
boxedStorable a
v =
    -- I think all types have to be word aligned.
    Size
box forall a. Num a => a -> a -> a
+ forall a. Integral a => a -> Size
fromBytes (forall a. Ord a => a -> a -> a
max (forall a. Storable a => a -> Int
Foreign.sizeOf forall a. Ptr a
Foreign.nullPtr) (forall a. Storable a => a -> Int
Foreign.sizeOf a
v))

-- * usage

-- | Return OS-reported (RSS, VSIZE).
rssVsize :: IO (Size, Size)
rssVsize :: IO (Size, Size)
rssVsize = do
    ProcessID
pid <- IO ProcessID
Posix.Process.getProcessID
    String
out <- String -> [String] -> String -> IO String
Process.readProcess String
"ps" [String
"-p", forall a. Show a => a -> String
show ProcessID
pid, String
"-orss,vsize"] String
""
    let [[String]
_, [String
rss, String
vsize]] = forall a b. (a -> b) -> [a] -> [b]
map String -> [String]
words (String -> [String]
lines String
out)
    forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. RealFrac a => a -> Size
fromK (forall a. Read a => String -> a
read String
rss), forall a. RealFrac a => a -> Size
fromK (forall a. Read a => String -> a
read String
vsize))

-- * RTS

-- | Allocated memory according to the GHC RTS.
rtsAllocated :: IO Size

#if GHC_VERSION >= 80401
rtsAllocated :: IO Size
rtsAllocated = do
    IO ()
System.Mem.performMajorGC
    RTSStats
stats <- IO RTSStats
GHC.Stats.getRTSStats
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Size
fromBytes forall a b. (a -> b) -> a -> b
$ GCDetails -> Word64
GHC.Stats.gcdetails_live_bytes forall a b. (a -> b) -> a -> b
$ RTSStats -> GCDetails
GHC.Stats.gc RTSStats
stats

stats :: IO GHC.Stats.RTSStats
stats :: IO RTSStats
stats = do
    IO ()
System.Mem.performMajorGC
    IO RTSStats
GHC.Stats.getRTSStats

#else

rtsAllocated = do
    System.Mem.performMajorGC
    stats <- GHC.Stats.getGCStats
    return $ fromBytes $ GHC.Stats.currentBytesUsed stats

stats :: IO GHC.Stats.GCStats
stats = do
    System.Mem.performMajorGC
    GHC.Stats.getGCStats

#endif