{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE RecordWildCards #-}

-- | Stopwatch/interval metering as a first-class circuit wire.
--
-- Instead of wrapping each stage individually, drop named markers into a
-- left-to-right pipeline. The timing log travels on a cartesian @(,)@ wire
-- alongside the payload.
--
-- A single watch is active at any point. 'start' sets the active watch;
-- 'lap' records an interval since the previous 'start' or 'lap' under a
-- given label and starts a fresh interval on the same active watch;
-- 'stop' records the final interval and keeps the log.
module Circuit.Meter.Stopwatch
  ( -- * Timing log
    Watches,
    allLaps,
    watchLaps,

    -- * Markers
    start,
    lap,
    stop,

    -- * Stage sugar
    carry,
    carryT,
    meterIt,
    timeIt,
    meterItN,
    timeItN,
  )
where

import Circuit hiding (eval)
import Circuit.Category (K (..), (.))
import Circuit.Meter (Meter, firstK)
import Circuit.Meter qualified as Meter
import Circuit.Meter.Time (Nanos, timeX)
import Circuit.Syntax (eval)
import Control.Exception (evaluate)
import Control.Monad (replicateM_)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Prelude hiding (id, (.))

-- ---------------------------------------------------------------------------
-- Timing log
-- ---------------------------------------------------------------------------

-- | Stopwatch state: the active watch name, the meter state for the current
-- in-flight interval, and the log of completed intervals.
data Watches x y = Watches
  { forall x y. Watches x y -> String
active :: String,
    forall x y. Watches x y -> x
startState :: x,
    forall x y. Watches x y -> Map String [y]
laps :: Map String [y]
  }

-- | Extract all measurements for a named label, oldest first.
watchLaps :: Watches x y -> String -> [y]
watchLaps :: forall x y. Watches x y -> String -> [y]
watchLaps Watches {x
String
Map String [y]
active :: forall x y. Watches x y -> String
startState :: forall x y. Watches x y -> x
laps :: forall x y. Watches x y -> Map String [y]
active :: String
startState :: x
laps :: Map String [y]
..} String
name = [y] -> ([y] -> [y]) -> Maybe [y] -> [y]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] [y] -> [y]
forall a. [a] -> [a]
reverse (Maybe [y] -> [y]) -> Maybe [y] -> [y]
forall a b. (a -> b) -> a -> b
$ String -> Map String [y] -> Maybe [y]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
name Map String [y]
laps

-- | Extract all measurements for all labels, oldest first.
allLaps :: Watches x y -> Map String [y]
allLaps :: forall x y. Watches x y -> Map String [y]
allLaps Watches {x
String
Map String [y]
active :: forall x y. Watches x y -> String
startState :: forall x y. Watches x y -> x
laps :: forall x y. Watches x y -> Map String [y]
active :: String
startState :: x
laps :: Map String [y]
..} = ([y] -> [y]) -> Map String [y] -> Map String [y]
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map [y] -> [y]
forall a. [a] -> [a]
reverse Map String [y]
laps

-- ---------------------------------------------------------------------------
-- Markers
-- ---------------------------------------------------------------------------

-- | Start a named watch. Sets the active watch and initializes the meter
-- state for the first interval.
start :: Meter (K IO) x y -> String -> Trace (,) (K IO) a (a, Watches x y)
start :: forall x y a.
Meter (K IO) x y -> String -> Trace (,) (K IO) a (a, Watches x y)
start Meter (K IO) x y
m String
name = K IO a (a, Watches x y) -> Trace (,) (K IO) a (a, Watches x y)
forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base (K IO a (a, Watches x y) -> Trace (,) (K IO) a (a, Watches x y))
-> K IO a (a, Watches x y) -> Trace (,) (K IO) a (a, Watches x y)
forall a b. (a -> b) -> a -> b
$ (a -> IO (a, Watches x y)) -> K IO a (a, Watches x y)
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K ((a -> IO (a, Watches x y)) -> K IO a (a, Watches x y))
-> (a -> IO (a, Watches x y)) -> K IO a (a, Watches x y)
forall a b. (a -> b) -> a -> b
$ \a
a -> do
  x <- K IO () x -> () -> IO x
forall {k} (m :: k -> *) a (b :: k). K m a b -> a -> m b
runK (Meter (K IO) x y -> K IO () x
forall (arr :: * -> * -> *) a b. Meter arr a b -> arr () a
Meter.start Meter (K IO) x y
m) ()
  pure (a, Watches name x Map.empty)

-- | Record a lap: stop the current interval, store the measurement under
-- @label@, and start a fresh interval on the same active watch.
lap :: Meter (K IO) x y -> String -> Trace (,) (K IO) (a, Watches x y) (a, Watches x y)
lap :: forall x y a.
Meter (K IO) x y
-> String -> Trace (,) (K IO) (a, Watches x y) (a, Watches x y)
lap Meter (K IO) x y
m String
label = K IO (a, Watches x y) (a, Watches x y)
-> Trace (,) (K IO) (a, Watches x y) (a, Watches x y)
forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base (K IO (a, Watches x y) (a, Watches x y)
 -> Trace (,) (K IO) (a, Watches x y) (a, Watches x y))
-> K IO (a, Watches x y) (a, Watches x y)
-> Trace (,) (K IO) (a, Watches x y) (a, Watches x y)
forall a b. (a -> b) -> a -> b
$ ((a, Watches x y) -> IO (a, Watches x y))
-> K IO (a, Watches x y) (a, Watches x y)
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K (((a, Watches x y) -> IO (a, Watches x y))
 -> K IO (a, Watches x y) (a, Watches x y))
-> ((a, Watches x y) -> IO (a, Watches x y))
-> K IO (a, Watches x y) (a, Watches x y)
forall a b. (a -> b) -> a -> b
$ \(a
a, Watches x y
ws) -> do
  y <- K IO x y -> x -> IO y
forall {k} (m :: k -> *) a (b :: k). K m a b -> a -> m b
runK (Meter (K IO) x y -> K IO x y
forall (arr :: * -> * -> *) a b. Meter arr a b -> arr a b
Meter.stop Meter (K IO) x y
m) (Watches x y -> x
forall x y. Watches x y -> x
startState Watches x y
ws)
  x' <- runK (Meter.start m) ()
  pure (a, ws {startState = x', laps = Map.insertWith (++) label [y] (laps ws)})

-- | Stop the active watch: record the final interval under @name@ and keep
-- the log.
stop :: Meter (K IO) x y -> String -> Trace (,) (K IO) (a, Watches x y) (a, Watches x y)
stop :: forall x y a.
Meter (K IO) x y
-> String -> Trace (,) (K IO) (a, Watches x y) (a, Watches x y)
stop Meter (K IO) x y
m String
name = K IO (a, Watches x y) (a, Watches x y)
-> Trace (,) (K IO) (a, Watches x y) (a, Watches x y)
forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base (K IO (a, Watches x y) (a, Watches x y)
 -> Trace (,) (K IO) (a, Watches x y) (a, Watches x y))
-> K IO (a, Watches x y) (a, Watches x y)
-> Trace (,) (K IO) (a, Watches x y) (a, Watches x y)
forall a b. (a -> b) -> a -> b
$ ((a, Watches x y) -> IO (a, Watches x y))
-> K IO (a, Watches x y) (a, Watches x y)
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K (((a, Watches x y) -> IO (a, Watches x y))
 -> K IO (a, Watches x y) (a, Watches x y))
-> ((a, Watches x y) -> IO (a, Watches x y))
-> K IO (a, Watches x y) (a, Watches x y)
forall a b. (a -> b) -> a -> b
$ \(a
a, Watches x y
ws) -> do
  y <- K IO x y -> x -> IO y
forall {k} (m :: k -> *) a (b :: k). K m a b -> a -> m b
runK (Meter (K IO) x y -> K IO x y
forall (arr :: * -> * -> *) a b. Meter arr a b -> arr a b
Meter.stop Meter (K IO) x y
m) (Watches x y -> x
forall x y. Watches x y -> x
startState Watches x y
ws)
  pure (a, ws {laps = Map.insertWith (++) name [y] (laps ws)})

-- ---------------------------------------------------------------------------
-- Stage sugar
-- ---------------------------------------------------------------------------

-- | Lift a base arrow so it carries the timing wire unchanged.
carry :: K IO a b -> Trace (,) (K IO) (a, Watches x y) (b, Watches x y)
carry :: forall a b x y.
K IO a b -> Trace (,) (K IO) (a, Watches x y) (b, Watches x y)
carry K IO a b
stage = K IO (a, Watches x y) (b, Watches x y)
-> Trace (,) (K IO) (a, Watches x y) (b, Watches x y)
forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base (K IO a b -> K IO (a, Watches x y) (b, Watches x y)
forall (m :: * -> *) a b c.
Functor m =>
K m a b -> K m (a, c) (b, c)
firstK K IO a b
stage)

-- | Lift an already-built 'Trace' stage so it carries the timing wire
-- unchanged. The stage is run at its own tensor and then threaded through the
-- cartesian timing wire.
carryT :: (Traced t (K IO)) => Trace t (K IO) a b -> Trace (,) (K IO) (a, Watches x y) (b, Watches x y)
carryT :: forall (t :: * -> * -> *) a b x y.
Traced t (K IO) =>
Trace t (K IO) a b
-> Trace (,) (K IO) (a, Watches x y) (b, Watches x y)
carryT Trace t (K IO) a b
stage = K IO (a, Watches x y) (b, Watches x y)
-> Trace (,) (K IO) (a, Watches x y) (b, Watches x y)
forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base (K IO a b -> K IO (a, Watches x y) (b, Watches x y)
forall (m :: * -> *) a b c.
Functor m =>
K m a b -> K m (a, c) (b, c)
firstK (Trace t (K IO) a b -> K IO a b
forall (arr :: * -> * -> *) (sig :: Sig) a b.
(Category arr, Algebra sig arr arr, Ctx sig arr arr) =>
Syntax sig arr a b -> arr a b
eval Trace t (K IO) a b
stage))

-- | Meter a single stage: start, run the stage, stop.
meterIt :: Meter (K IO) x y -> String -> K IO a b -> Trace (,) (K IO) a (b, Watches x y)
meterIt :: forall x y a b.
Meter (K IO) x y
-> String -> K IO a b -> Trace (,) (K IO) a (b, Watches x y)
meterIt Meter (K IO) x y
m String
name K IO a b
stage =
  Meter (K IO) x y -> String -> Trace (,) (K IO) a (a, Watches x y)
forall x y a.
Meter (K IO) x y -> String -> Trace (,) (K IO) a (a, Watches x y)
start Meter (K IO) x y
m String
name
    Trace (,) (K IO) a (a, Watches x y)
-> Syntax
     (SigCompose :+: SigYank (,))
     (K IO)
     (a, Watches x y)
     (b, Watches x y)
-> Syntax (SigCompose :+: SigYank (,)) (K IO) a (b, Watches x y)
forall {k} (arr :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category arr =>
arr a b -> arr b c -> arr a c
.> K IO a b
-> Syntax
     (SigCompose :+: SigYank (,))
     (K IO)
     (a, Watches x y)
     (b, Watches x y)
forall a b x y.
K IO a b -> Trace (,) (K IO) (a, Watches x y) (b, Watches x y)
carry K IO a b
stage
    Syntax (SigCompose :+: SigYank (,)) (K IO) a (b, Watches x y)
-> Syntax
     (SigCompose :+: SigYank (,))
     (K IO)
     (b, Watches x y)
     (b, Watches x y)
-> Syntax (SigCompose :+: SigYank (,)) (K IO) a (b, Watches x y)
forall {k} (arr :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category arr =>
arr a b -> arr b c -> arr a c
.> Meter (K IO) x y
-> String
-> Syntax
     (SigCompose :+: SigYank (,))
     (K IO)
     (b, Watches x y)
     (b, Watches x y)
forall x y a.
Meter (K IO) x y
-> String -> Trace (,) (K IO) (a, Watches x y) (a, Watches x y)
stop Meter (K IO) x y
m String
name

-- | 'meterIt' with the default time meter.
timeIt :: String -> K IO a b -> Trace (,) (K IO) a (b, Watches Nanos Nanos)
timeIt :: forall a b.
String -> K IO a b -> Trace (,) (K IO) a (b, Watches Nanos Nanos)
timeIt = Meter (K IO) Nanos Nanos
-> String
-> K IO a b
-> Trace (,) (K IO) a (b, Watches Nanos Nanos)
forall x y a b.
Meter (K IO) x y
-> String -> K IO a b -> Trace (,) (K IO) a (b, Watches x y)
meterIt Meter (K IO) Nanos Nanos
timeX

-- | Meter a stage over @n@ repetitions and record the total measurement.
--
-- The stage is run @n@ times between 'start' and 'stop'; the recorded lap is
-- the total time/space for all runs. Divide by @n@ for a per-iteration average.
-- The last result is kept and forced to WHNF; intermediate results are also forced
-- so the work cannot be floated out of the loop.
meterItN :: Int -> Meter (K IO) x y -> String -> K IO a b -> Trace (,) (K IO) a (b, Watches x y)
meterItN :: forall x y a b.
Int
-> Meter (K IO) x y
-> String
-> K IO a b
-> Trace (,) (K IO) a (b, Watches x y)
meterItN Int
n0 Meter (K IO) x y
m String
name K IO a b
stage =
  Meter (K IO) x y -> String -> Trace (,) (K IO) a (a, Watches x y)
forall x y a.
Meter (K IO) x y -> String -> Trace (,) (K IO) a (a, Watches x y)
start Meter (K IO) x y
m String
name
    Trace (,) (K IO) a (a, Watches x y)
-> Syntax
     (SigCompose :+: SigYank (,))
     (K IO)
     (a, Watches x y)
     (b, Watches x y)
-> Syntax (SigCompose :+: SigYank (,)) (K IO) a (b, Watches x y)
forall {k} (arr :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category arr =>
arr a b -> arr b c -> arr a c
.> K IO a b
-> Syntax
     (SigCompose :+: SigYank (,))
     (K IO)
     (a, Watches x y)
     (b, Watches x y)
forall a b x y.
K IO a b -> Trace (,) (K IO) (a, Watches x y) (b, Watches x y)
carry ((a -> IO b) -> K IO a b
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K a -> IO b
loop)
    Syntax (SigCompose :+: SigYank (,)) (K IO) a (b, Watches x y)
-> Syntax
     (SigCompose :+: SigYank (,))
     (K IO)
     (b, Watches x y)
     (b, Watches x y)
-> Syntax (SigCompose :+: SigYank (,)) (K IO) a (b, Watches x y)
forall {k} (arr :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category arr =>
arr a b -> arr b c -> arr a c
.> Meter (K IO) x y
-> String
-> Syntax
     (SigCompose :+: SigYank (,))
     (K IO)
     (b, Watches x y)
     (b, Watches x y)
forall x y a.
Meter (K IO) x y
-> String -> Trace (,) (K IO) (a, Watches x y) (a, Watches x y)
stop Meter (K IO) x y
m String
name
  where
    n :: Int
n = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
n0
    loop :: a -> IO b
loop a
a = do
      b <- K IO a b -> a -> IO b
forall {k} (m :: k -> *) a (b :: k). K m a b -> a -> m b
runK K IO a b
stage a
a IO b -> (b -> IO b) -> IO b
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> IO b
forall a. a -> IO a
evaluate
      replicateM_ (n - 1) $ do
        !_ <- runK stage a >>= evaluate
        pure ()
      pure b

-- | 'meterItN' with the default time meter.
timeItN :: Int -> String -> K IO a b -> Trace (,) (K IO) a (b, Watches Nanos Nanos)
timeItN :: forall a b.
Int
-> String
-> K IO a b
-> Trace (,) (K IO) a (b, Watches Nanos Nanos)
timeItN Int
n = Int
-> Meter (K IO) Nanos Nanos
-> String
-> K IO a b
-> Trace (,) (K IO) a (b, Watches Nanos Nanos)
forall x y a b.
Int
-> Meter (K IO) x y
-> String
-> K IO a b
-> Trace (,) (K IO) a (b, Watches x y)
meterItN Int
n Meter (K IO) Nanos Nanos
timeX