{-# LANGUAGE BlockArguments #-}

-- | Performance measurement reimagined as a Circuit.
--
-- A 'Meter' @arr a b@ is a pair of arrows: 'start' produces the initial
-- state, 'stop' observes it and produces a measurement. The canonical
-- use is 'meterAction', which sandwiches a payload arrow between the
-- two meter arrows.
module Circuit.Meter
  ( -- * Meter
    Meter (..),
    mkMeter,

    -- * Cartesian helpers for Kleisli arrows
    firstK,
    secondK,
    dimapK,

    -- * Meter composition
    both,

    -- * Plugin metering
    meterAction,

    -- * Hold
    hold,
  )
where

import Circuit.Category (Category (..), K (..))
import Circuit.Trace (Trace, base)
import Prelude hiding (id, (.))

-- ---------------------------------------------------------------------------
-- Meter
-- ---------------------------------------------------------------------------

-- | A 'Meter' @arr a b@ is a stopwatch: a pair of arrows.
--
-- * 'start' — @() → arr a@ — capture initial state (e.g. read the clock).
-- * 'stop' — @a → arr b@ — observe the state and produce a measurement
--   (e.g. read the clock again and subtract).  May be called multiple
--   times; each call measures elapsed time since the single 'start'.
data Meter arr a b = Meter
  { forall (arr :: * -> * -> *) a b. Meter arr a b -> arr () a
start :: arr () a,
    forall (arr :: * -> * -> *) a b. Meter arr a b -> arr a b
stop :: arr a b
  }

-- | Construct a 'Meter' from raw monadic actions.
mkMeter :: m a -> (a -> m b) -> Meter (K m) a b
mkMeter :: forall (m :: * -> *) a b. m a -> (a -> m b) -> Meter (K m) a b
mkMeter m a
pre a -> m b
post = K m () a -> K m a b -> Meter (K m) a b
forall (arr :: * -> * -> *) a b.
arr () a -> arr a b -> Meter arr a b
Meter ((() -> m a) -> K m () a
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K (m a -> () -> m a
forall a b. a -> b -> a
const m a
pre)) ((a -> m b) -> K m a b
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K a -> m b
post)
{-# INLINEABLE mkMeter #-}

-- ---------------------------------------------------------------------------
-- Cartesian helpers for Kleisli arrows
-- ---------------------------------------------------------------------------

-- | First component for @K m@.
firstK :: (Functor m) => K m a b -> K m (a, c) (b, c)
firstK :: forall (m :: * -> *) a b c.
Functor m =>
K m a b -> K m (a, c) (b, c)
firstK (K a -> m b
k) = ((a, c) -> m (b, c)) -> K m (a, c) (b, c)
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K (\(a
a, c
c) -> (b -> (b, c)) -> m b -> m (b, c)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\b
b -> (b
b, c
c)) (a -> m b
k a
a))
{-# INLINEABLE firstK #-}

-- | Second component for @K m@.
secondK :: (Functor m) => K m a b -> K m (c, a) (c, b)
secondK :: forall (m :: * -> *) a b c.
Functor m =>
K m a b -> K m (c, a) (c, b)
secondK (K a -> m b
k) = ((c, a) -> m (c, b)) -> K m (c, a) (c, b)
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K (\(c
c, a
a) -> (b -> (c, b)) -> m b -> m (c, b)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\b
b -> (c
c, b
b)) (a -> m b
k a
a))
{-# INLINEABLE secondK #-}

-- | Profunctor-style pre/post composition for @K m@.
dimapK :: (Functor m) => (a' -> a) -> (b -> b') -> K m a b -> K m a' b'
dimapK :: forall (m :: * -> *) a' a b b'.
Functor m =>
(a' -> a) -> (b -> b') -> K m a b -> K m a' b'
dimapK a' -> a
f b -> b'
g (K a -> m b
k) = (a' -> m b') -> K m a' b'
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K ((b -> b') -> m b -> m b'
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b'
g (m b -> m b') -> (a -> m b) -> a -> m b'
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. a -> m b
k (a -> m b') -> (a' -> a) -> a' -> m b'
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. a' -> a
f)
{-# INLINEABLE dimapK #-}

-- | Run two meters simultaneously.
--
-- The state wires are independent; the @(,)@ tensor handles the
-- wiring automatically.
both :: (Monad m) => Meter (K m) a1 b1 -> Meter (K m) a2 b2 -> Meter (K m) (a1, a2) (b1, b2)
both :: forall (m :: * -> *) a1 b1 a2 b2.
Monad m =>
Meter (K m) a1 b1
-> Meter (K m) a2 b2 -> Meter (K m) (a1, a2) (b1, b2)
both Meter (K m) a1 b1
m1 Meter (K m) a2 b2
m2 =
  Meter
    { start :: K m () (a1, a2)
start = (() -> ((), ()))
-> ((a1, a2) -> (a1, a2))
-> K m ((), ()) (a1, a2)
-> K m () (a1, a2)
forall (m :: * -> *) a' a b b'.
Functor m =>
(a' -> a) -> (b -> b') -> K m a b -> K m a' b'
dimapK (\() -> ((), ())) (a1, a2) -> (a1, a2)
forall a. a -> a
forall k (arr :: k -> k -> *) (a :: k). Category arr => arr a a
id (K m () a1 -> K m ((), a2) (a1, a2)
forall (m :: * -> *) a b c.
Functor m =>
K m a b -> K m (a, c) (b, c)
firstK (Meter (K m) a1 b1 -> K m () a1
forall (arr :: * -> * -> *) a b. Meter arr a b -> arr () a
start Meter (K m) a1 b1
m1) K m ((), a2) (a1, a2)
-> K m ((), ()) ((), a2) -> K m ((), ()) (a1, a2)
forall b c a. K m b c -> K m a b -> K m a c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. K m () a2 -> K m ((), ()) ((), a2)
forall (m :: * -> *) a b c.
Functor m =>
K m a b -> K m (c, a) (c, b)
secondK (Meter (K m) a2 b2 -> K m () a2
forall (arr :: * -> * -> *) a b. Meter arr a b -> arr () a
start Meter (K m) a2 b2
m2)),
      stop :: K m (a1, a2) (b1, b2)
stop = K m a1 b1 -> K m (a1, b2) (b1, b2)
forall (m :: * -> *) a b c.
Functor m =>
K m a b -> K m (a, c) (b, c)
firstK (Meter (K m) a1 b1 -> K m a1 b1
forall (arr :: * -> * -> *) a b. Meter arr a b -> arr a b
stop Meter (K m) a1 b1
m1) K m (a1, b2) (b1, b2)
-> K m (a1, a2) (a1, b2) -> K m (a1, a2) (b1, b2)
forall b c a. K m b c -> K m a b -> K m a c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. K m a2 b2 -> K m (a1, a2) (a1, b2)
forall (m :: * -> *) a b c.
Functor m =>
K m a b -> K m (c, a) (c, b)
secondK (Meter (K m) a2 b2 -> K m a2 b2
forall (arr :: * -> * -> *) a b. Meter arr a b -> arr a b
stop Meter (K m) a2 b2
m2)
    }
{-# INLINEABLE both #-}

-- ---------------------------------------------------------------------------
-- Plugin metering
-- ---------------------------------------------------------------------------

-- | Meter a Kleisli arrow action, keeping the measurement.
--
-- Tensor-agnostic: the bracket is built directly in the base arrow
-- and lifted with 'base', so the meter state is introduced and consumed
-- locally. The result is a 'Circuit' polymorphic in the tensor @t@.
--
-- For arrow-level extraction, use 'eval' with your chosen tensor.
meterAction :: (Monad m) => Meter (K m) a b -> K m c d -> Trace t (K m) c (b, d)
meterAction :: forall (m :: * -> *) a b c d (t :: * -> * -> *).
Monad m =>
Meter (K m) a b -> K m c d -> Trace t (K m) c (b, d)
meterAction Meter (K m) a b
m K m c d
k =
  K m c (b, d) -> Trace t (K m) c (b, d)
forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base (K m a b -> K m (a, d) (b, d)
forall (m :: * -> *) a b c.
Functor m =>
K m a b -> K m (a, c) (b, c)
firstK (Meter (K m) a b -> K m a b
forall (arr :: * -> * -> *) a b. Meter arr a b -> arr a b
stop Meter (K m) a b
m) K m (a, d) (b, d) -> K m (a, c) (a, d) -> K m (a, c) (b, d)
forall b c a. K m b c -> K m a b -> K m a c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. K m c d -> K m (a, c) (a, d)
forall (m :: * -> *) a b c.
Functor m =>
K m a b -> K m (c, a) (c, b)
secondK K m c d
k K m (a, c) (b, d) -> K m c (a, c) -> K m c (b, d)
forall b c a. K m b c -> K m a b -> K m a c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. (c -> ((), c))
-> ((a, c) -> (a, c)) -> K m ((), c) (a, c) -> K m c (a, c)
forall (m :: * -> *) a' a b b'.
Functor m =>
(a' -> a) -> (b -> b') -> K m a b -> K m a' b'
dimapK ((),) (a, c) -> (a, c)
forall a. a -> a
forall k (arr :: k -> k -> *) (a :: k). Category arr => arr a a
id (K m () a -> K m ((), c) (a, c)
forall (m :: * -> *) a b c.
Functor m =>
K m a b -> K m (a, c) (b, c)
firstK (Meter (K m) a b -> K m () a
forall (arr :: * -> * -> *) a b. Meter arr a b -> arr () a
start Meter (K m) a b
m)))
{-# INLINEABLE meterAction #-}

-- | Hold back a value so GHC cannot float a function application past
-- the meter boundary. Re-exported for custom meter authors.
hold :: a -> a
hold :: forall a. a -> a
hold a
x = a
x
{-# NOINLINE hold #-}