{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}

module Control.Tracer.Transformers
  ( Counting(..)
  , Folding(..)
  , counting
  , fanning
  , folding
  ) where

import           Control.Monad (join)
import           Control.Monad.IO.Class (MonadIO (..))
import           Data.IORef (IORef, atomicModifyIORef', newIORef)

import           Control.Tracer

-- | A pure tracer combinator that allows to decide a further tracer to use,
--   based on the message being processed.
fanning
  :: forall m a
  . (a -> Tracer m a) -> Tracer m a
fanning :: (a -> Tracer m a) -> Tracer m a
fanning a -> Tracer m a
fan = (a -> m ()) -> Tracer m a
forall (m :: * -> *) a. (a -> m ()) -> Tracer m a
Tracer ((a -> m ()) -> Tracer m a) -> (a -> m ()) -> Tracer m a
forall a b. (a -> b) -> a -> b
$ \a
x -> Tracer m a -> a -> m ()
forall (m :: * -> *) a. Tracer m a -> a -> m ()
traceWith (a -> Tracer m a
fan a
x) a
x

newtype Counting a = Counting Int

-- | A stateful tracer transformer that substitutes messages with
--   a monotonically incrementing occurence count.
counting
  :: forall m a . (MonadIO m)
  => Tracer m (Counting a) -> m (Tracer m a)
counting :: Tracer m (Counting a) -> m (Tracer m a)
counting Tracer m (Counting a)
tr =
  IORef Int -> Tracer m a
mkTracer (IORef Int -> Tracer m a) -> m (IORef Int) -> m (Tracer m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (IORef Int) -> m (IORef Int)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0)
 where
    mkTracer :: IORef Int -> Tracer m a
    mkTracer :: IORef Int -> Tracer m a
mkTracer IORef Int
ctrref = (a -> m ()) -> Tracer m a
forall (m :: * -> *) a. (a -> m ()) -> Tracer m a
Tracer ((a -> m ()) -> Tracer m a) -> (a -> m ()) -> Tracer m a
forall a b. (a -> b) -> a -> b
$ \a
_ -> do
      Int
ctr <- IO Int -> m Int
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> m Int) -> IO Int -> m Int
forall a b. (a -> b) -> a -> b
$ IORef Int -> (Int -> (Int, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
ctrref ((Int -> (Int, Int)) -> IO Int) -> (Int -> (Int, Int)) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Int
n -> (Int -> Int -> (Int, Int)) -> Int -> (Int, Int)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (,) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      Tracer m (Counting a) -> Counting a -> m ()
forall (m :: * -> *) a. Tracer m a -> a -> m ()
traceWith Tracer m (Counting a)
tr (Int -> Counting a
forall a. Int -> Counting a
Counting Int
ctr)

newtype Folding a f = Folding f

-- | A generalised trace transformer that provides evolving state,
--   defined as a strict left fold.
folding
  :: forall m f a
  .  (MonadIO m)
  => (f -> a -> f) -> f -> Tracer m (Folding a f) -> m (Tracer m a)
folding :: (f -> a -> f) -> f -> Tracer m (Folding a f) -> m (Tracer m a)
folding f -> a -> f
cata f
initial Tracer m (Folding a f)
tr =
  IORef f -> Tracer m a
mkTracer (IORef f -> Tracer m a) -> m (IORef f) -> m (Tracer m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (IORef f) -> m (IORef f)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (f -> IO (IORef f)
forall a. a -> IO (IORef a)
newIORef f
initial)
 where
    mkTracer :: IORef f -> Tracer m a
    mkTracer :: IORef f -> Tracer m a
mkTracer IORef f
ref = (a -> m ()) -> Tracer m a
forall (m :: * -> *) a. (a -> m ()) -> Tracer m a
Tracer ((a -> m ()) -> Tracer m a) -> (a -> m ()) -> Tracer m a
forall a b. (a -> b) -> a -> b
$ \a
a -> do
      f
x' <- IO f -> m f
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO f -> m f) -> IO f -> m f
forall a b. (a -> b) -> a -> b
$ IORef f -> (f -> (f, f)) -> IO f
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef f
ref ((f -> (f, f)) -> IO f) -> (f -> (f, f)) -> IO f
forall a b. (a -> b) -> a -> b
$ \f
x -> (f -> f -> (f, f)) -> f -> (f, f)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (,) (f -> a -> f
cata f
x a
a)
      Tracer m (Folding a f) -> Folding a f -> m ()
forall (m :: * -> *) a. Tracer m a -> a -> m ()
traceWith Tracer m (Folding a f)
tr (f -> Folding a f
forall a f. f -> Folding a f
Folding f
x')