{-# LANGUAGE BlockArguments #-}

-- | Space measurement as a Circuit.
--
-- GHC RTS allocation statistics read before and after a computation.
-- 'allocX' measures allocated bytes; 'SpaceStats' is available for
-- users who want the full RTS snapshot.
module Circuit.Meter.Space
  ( -- * Space meter
    allocX,
    allocGC,

    -- * Types
    SpaceStats (..),
    Bytes (..),
  )
where

import Circuit.Category (K (..))
import Circuit.Meter
import Control.Category ((.))
import Data.Word (Word32, Word64)
import GHC.Stats
import System.Mem (performGC)
import Prelude hiding (id, (.))

-- | Allocation statistics from the GHC RTS.
data SpaceStats = SpaceStats
  { SpaceStats -> Word64
allocated :: !Word64,
    SpaceStats -> Word64
copied :: !Word64,
    SpaceStats -> Word64
maxmem :: !Word64,
    SpaceStats -> Word32
minorgcs :: !Word32,
    SpaceStats -> Word32
majorgcs :: !Word32
  }
  deriving (ReadPrec [SpaceStats]
ReadPrec SpaceStats
Int -> ReadS SpaceStats
ReadS [SpaceStats]
(Int -> ReadS SpaceStats)
-> ReadS [SpaceStats]
-> ReadPrec SpaceStats
-> ReadPrec [SpaceStats]
-> Read SpaceStats
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS SpaceStats
readsPrec :: Int -> ReadS SpaceStats
$creadList :: ReadS [SpaceStats]
readList :: ReadS [SpaceStats]
$creadPrec :: ReadPrec SpaceStats
readPrec :: ReadPrec SpaceStats
$creadListPrec :: ReadPrec [SpaceStats]
readListPrec :: ReadPrec [SpaceStats]
Read, Int -> SpaceStats -> ShowS
[SpaceStats] -> ShowS
SpaceStats -> String
(Int -> SpaceStats -> ShowS)
-> (SpaceStats -> String)
-> ([SpaceStats] -> ShowS)
-> Show SpaceStats
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SpaceStats -> ShowS
showsPrec :: Int -> SpaceStats -> ShowS
$cshow :: SpaceStats -> String
show :: SpaceStats -> String
$cshowList :: [SpaceStats] -> ShowS
showList :: [SpaceStats] -> ShowS
Show, SpaceStats -> SpaceStats -> Bool
(SpaceStats -> SpaceStats -> Bool)
-> (SpaceStats -> SpaceStats -> Bool) -> Eq SpaceStats
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SpaceStats -> SpaceStats -> Bool
== :: SpaceStats -> SpaceStats -> Bool
$c/= :: SpaceStats -> SpaceStats -> Bool
/= :: SpaceStats -> SpaceStats -> Bool
Eq)

instance Semigroup SpaceStats where
  <> :: SpaceStats -> SpaceStats -> SpaceStats
(<>) = SpaceStats -> SpaceStats -> SpaceStats
addSpace

instance Monoid SpaceStats where
  mempty :: SpaceStats
mempty = Word64 -> Word64 -> Word64 -> Word32 -> Word32 -> SpaceStats
SpaceStats Word64
0 Word64
0 Word64
0 Word32
0 Word32
0

addSpace :: SpaceStats -> SpaceStats -> SpaceStats
addSpace :: SpaceStats -> SpaceStats -> SpaceStats
addSpace (SpaceStats Word64
a1 Word64
c1 Word64
m1 Word32
g1 Word32
g1') (SpaceStats Word64
a2 Word64
c2 Word64
m2 Word32
g2 Word32
g2') =
  Word64 -> Word64 -> Word64 -> Word32 -> Word32 -> SpaceStats
SpaceStats (Word64
a1 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
a2) (Word64
c1 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
c2) (Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
max Word64
m1 Word64
m2) (Word32
g1 Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
g2) (Word32
g1' Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
g2')

-- | Number of bytes.
newtype Bytes = Bytes {Bytes -> Word64
unbytes :: Word64}
  deriving (Int -> Bytes -> ShowS
[Bytes] -> ShowS
Bytes -> String
(Int -> Bytes -> ShowS)
-> (Bytes -> String) -> ([Bytes] -> ShowS) -> Show Bytes
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Bytes -> ShowS
showsPrec :: Int -> Bytes -> ShowS
$cshow :: Bytes -> String
show :: Bytes -> String
$cshowList :: [Bytes] -> ShowS
showList :: [Bytes] -> ShowS
Show, ReadPrec [Bytes]
ReadPrec Bytes
Int -> ReadS Bytes
ReadS [Bytes]
(Int -> ReadS Bytes)
-> ReadS [Bytes]
-> ReadPrec Bytes
-> ReadPrec [Bytes]
-> Read Bytes
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS Bytes
readsPrec :: Int -> ReadS Bytes
$creadList :: ReadS [Bytes]
readList :: ReadS [Bytes]
$creadPrec :: ReadPrec Bytes
readPrec :: ReadPrec Bytes
$creadListPrec :: ReadPrec [Bytes]
readListPrec :: ReadPrec [Bytes]
Read, Bytes -> Bytes -> Bool
(Bytes -> Bytes -> Bool) -> (Bytes -> Bytes -> Bool) -> Eq Bytes
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Bytes -> Bytes -> Bool
== :: Bytes -> Bytes -> Bool
$c/= :: Bytes -> Bytes -> Bool
/= :: Bytes -> Bytes -> Bool
Eq, Eq Bytes
Eq Bytes =>
(Bytes -> Bytes -> Ordering)
-> (Bytes -> Bytes -> Bool)
-> (Bytes -> Bytes -> Bool)
-> (Bytes -> Bytes -> Bool)
-> (Bytes -> Bytes -> Bool)
-> (Bytes -> Bytes -> Bytes)
-> (Bytes -> Bytes -> Bytes)
-> Ord Bytes
Bytes -> Bytes -> Bool
Bytes -> Bytes -> Ordering
Bytes -> Bytes -> Bytes
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
$ccompare :: Bytes -> Bytes -> Ordering
compare :: Bytes -> Bytes -> Ordering
$c< :: Bytes -> Bytes -> Bool
< :: Bytes -> Bytes -> Bool
$c<= :: Bytes -> Bytes -> Bool
<= :: Bytes -> Bytes -> Bool
$c> :: Bytes -> Bytes -> Bool
> :: Bytes -> Bytes -> Bool
$c>= :: Bytes -> Bytes -> Bool
>= :: Bytes -> Bytes -> Bool
$cmax :: Bytes -> Bytes -> Bytes
max :: Bytes -> Bytes -> Bytes
$cmin :: Bytes -> Bytes -> Bytes
min :: Bytes -> Bytes -> Bytes
Ord, Integer -> Bytes
Bytes -> Bytes
Bytes -> Bytes -> Bytes
(Bytes -> Bytes -> Bytes)
-> (Bytes -> Bytes -> Bytes)
-> (Bytes -> Bytes -> Bytes)
-> (Bytes -> Bytes)
-> (Bytes -> Bytes)
-> (Bytes -> Bytes)
-> (Integer -> Bytes)
-> Num Bytes
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
$c+ :: Bytes -> Bytes -> Bytes
+ :: Bytes -> Bytes -> Bytes
$c- :: Bytes -> Bytes -> Bytes
- :: Bytes -> Bytes -> Bytes
$c* :: Bytes -> Bytes -> Bytes
* :: Bytes -> Bytes -> Bytes
$cnegate :: Bytes -> Bytes
negate :: Bytes -> Bytes
$cabs :: Bytes -> Bytes
abs :: Bytes -> Bytes
$csignum :: Bytes -> Bytes
signum :: Bytes -> Bytes
$cfromInteger :: Integer -> Bytes
fromInteger :: Integer -> Bytes
Num, Num Bytes
Ord Bytes
(Num Bytes, Ord Bytes) => (Bytes -> Rational) -> Real Bytes
Bytes -> Rational
forall a. (Num a, Ord a) => (a -> Rational) -> Real a
$ctoRational :: Bytes -> Rational
toRational :: Bytes -> Rational
Real, Int -> Bytes
Bytes -> Int
Bytes -> [Bytes]
Bytes -> Bytes
Bytes -> Bytes -> [Bytes]
Bytes -> Bytes -> Bytes -> [Bytes]
(Bytes -> Bytes)
-> (Bytes -> Bytes)
-> (Int -> Bytes)
-> (Bytes -> Int)
-> (Bytes -> [Bytes])
-> (Bytes -> Bytes -> [Bytes])
-> (Bytes -> Bytes -> [Bytes])
-> (Bytes -> Bytes -> Bytes -> [Bytes])
-> Enum Bytes
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: Bytes -> Bytes
succ :: Bytes -> Bytes
$cpred :: Bytes -> Bytes
pred :: Bytes -> Bytes
$ctoEnum :: Int -> Bytes
toEnum :: Int -> Bytes
$cfromEnum :: Bytes -> Int
fromEnum :: Bytes -> Int
$cenumFrom :: Bytes -> [Bytes]
enumFrom :: Bytes -> [Bytes]
$cenumFromThen :: Bytes -> Bytes -> [Bytes]
enumFromThen :: Bytes -> Bytes -> [Bytes]
$cenumFromTo :: Bytes -> Bytes -> [Bytes]
enumFromTo :: Bytes -> Bytes -> [Bytes]
$cenumFromThenTo :: Bytes -> Bytes -> Bytes -> [Bytes]
enumFromThenTo :: Bytes -> Bytes -> Bytes -> [Bytes]
Enum, Enum Bytes
Real Bytes
(Real Bytes, Enum Bytes) =>
(Bytes -> Bytes -> Bytes)
-> (Bytes -> Bytes -> Bytes)
-> (Bytes -> Bytes -> Bytes)
-> (Bytes -> Bytes -> Bytes)
-> (Bytes -> Bytes -> (Bytes, Bytes))
-> (Bytes -> Bytes -> (Bytes, Bytes))
-> (Bytes -> Integer)
-> Integral Bytes
Bytes -> Integer
Bytes -> Bytes -> (Bytes, Bytes)
Bytes -> Bytes -> Bytes
forall a.
(Real a, Enum a) =>
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> (a, a))
-> (a -> a -> (a, a))
-> (a -> Integer)
-> Integral a
$cquot :: Bytes -> Bytes -> Bytes
quot :: Bytes -> Bytes -> Bytes
$crem :: Bytes -> Bytes -> Bytes
rem :: Bytes -> Bytes -> Bytes
$cdiv :: Bytes -> Bytes -> Bytes
div :: Bytes -> Bytes -> Bytes
$cmod :: Bytes -> Bytes -> Bytes
mod :: Bytes -> Bytes -> Bytes
$cquotRem :: Bytes -> Bytes -> (Bytes, Bytes)
quotRem :: Bytes -> Bytes -> (Bytes, Bytes)
$cdivMod :: Bytes -> Bytes -> (Bytes, Bytes)
divMod :: Bytes -> Bytes -> (Bytes, Bytes)
$ctoInteger :: Bytes -> Integer
toInteger :: Bytes -> Integer
Integral)

instance Semigroup Bytes where
  <> :: Bytes -> Bytes -> Bytes
(<>) = Bytes -> Bytes -> Bytes
forall a. Num a => a -> a -> a
(+)

instance Monoid Bytes where
  mempty :: Bytes
mempty = Bytes
0

-- | Measure only allocated bytes.
allocX :: Meter (K IO) Bytes Bytes
allocX :: Meter (K IO) Bytes Bytes
allocX =
  IO Bytes -> (Bytes -> IO Bytes) -> Meter (K IO) Bytes Bytes
forall (m :: * -> *) a b. m a -> (a -> m b) -> Meter (K m) a b
mkMeter
    ((RTSStats -> Bytes) -> IO RTSStats -> IO Bytes
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Word64 -> Bytes
Bytes (Word64 -> Bytes) -> (RTSStats -> Word64) -> RTSStats -> Bytes
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. RTSStats -> Word64
allocated_bytes) IO RTSStats
getRTSStats)
    ( \Bytes
s -> do
        s' <- (RTSStats -> Bytes) -> IO RTSStats -> IO Bytes
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Word64 -> Bytes
Bytes (Word64 -> Bytes) -> (RTSStats -> Word64) -> RTSStats -> Bytes
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. RTSStats -> Word64
allocated_bytes) IO RTSStats
getRTSStats
        pure (s' - s)
    )
{-# INLINEABLE allocX #-}

-- | Measure allocated bytes with GC-forced boundaries.
--
-- Forces a major GC before reading the counter at both start and stop.
-- This gives accurate per-interval allocation at the cost of GC overhead
-- (a stop-the-world collection on every measurement boundary).
--
-- Use 'allocX' for lightweight (but potentially stale) measurements;
-- use 'allocGC' when you need accurate per-stage numbers.
allocGC :: Meter (K IO) Bytes Bytes
allocGC :: Meter (K IO) Bytes Bytes
allocGC =
  IO Bytes -> (Bytes -> IO Bytes) -> Meter (K IO) Bytes Bytes
forall (m :: * -> *) a b. m a -> (a -> m b) -> Meter (K m) a b
mkMeter
    (IO ()
performGC IO () -> IO Bytes -> IO Bytes
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (RTSStats -> Bytes) -> IO RTSStats -> IO Bytes
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Word64 -> Bytes
Bytes (Word64 -> Bytes) -> (RTSStats -> Word64) -> RTSStats -> Bytes
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. RTSStats -> Word64
allocated_bytes) IO RTSStats
getRTSStats)
    ( \Bytes
s -> do
        IO ()
performGC
        s' <- (RTSStats -> Bytes) -> IO RTSStats -> IO Bytes
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Word64 -> Bytes
Bytes (Word64 -> Bytes) -> (RTSStats -> Word64) -> RTSStats -> Bytes
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. RTSStats -> Word64
allocated_bytes) IO RTSStats
getRTSStats
        pure (s' - s)
    )
{-# INLINEABLE allocGC #-}