{-|
Module      : Control.TracerA.Arrow
Copyright   : (c) Alexander Vieth, 2019
Licence     : Apache-2.0
Maintainer  : aovieth@gmail.com
-}

{-# LANGUAGE GADTs        #-}
{-# LANGUAGE Arrows       #-}
{-# LANGUAGE RankNTypes   #-}
{-# LANGUAGE BangPatterns #-}

module Control.Tracer.Arrow
  ( TracerA (..)
  , runTracerA
  , compute
  , emit
  , effect
  , squelch
  , nat
  ) where

import Prelude hiding ((.), id)
import Control.Arrow
import Control.Category

-- | Formal representation of a tracer arrow as a Kleisli arrow over some
-- monad, but tagged so that we know whether it has any effects which will emit
-- a trace.
data TracerA m a b where
  -- | An emitting part, followed by a non-emitting part.
  -- The non-emitting part is there so that later emitting parts can be
  -- tacked-on later.
  Emitting   :: Kleisli m a x -> Kleisli m x b -> TracerA m a b
  -- | No emitting. There may be side-effects, but they are assumed to be
  -- benign and will be discarded by 'runTracerA'.
  Squelching :: Kleisli m a b                  -> TracerA m a b

-- | The resulting Kleisli arrow includes all of the effects required to do
-- the emitting part.
runTracerA :: Monad m => TracerA m a () -> Kleisli m a ()
runTracerA :: TracerA m a () -> Kleisli m a ()
runTracerA (Emitting Kleisli m a x
emits Kleisli m x ()
_noEmits) = Kleisli m a x
emits Kleisli m a x -> Kleisli m x () -> Kleisli m a ()
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> (x -> ()) -> Kleisli m x ()
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (() -> x -> ()
forall a b. a -> b -> a
const ())
runTracerA (Squelching     Kleisli m a ()
_       ) =           (a -> ()) -> Kleisli m a ()
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (() -> a -> ()
forall a b. a -> b -> a
const ())

-- | Ignore the input and do not emit. The name is intended to lead to clear
-- and suggestive arrow expressions.
squelch :: Applicative m => TracerA m a ()
squelch :: TracerA m a ()
squelch = (a -> ()) -> TracerA m a ()
forall (m :: * -> *) a b.
Applicative m =>
(a -> b) -> TracerA m a b
compute (() -> a -> ()
forall a b. a -> b -> a
const ())

-- | Do an emitting effect. Contrast with 'effect' which does not make the
-- tracer an emitting tracer.
emit :: Applicative m => (a -> m ()) -> TracerA m a ()
emit :: (a -> m ()) -> TracerA m a ()
emit a -> m ()
f = Kleisli m a () -> Kleisli m () () -> TracerA m a ()
forall (m :: * -> *) a x b.
Kleisli m a x -> Kleisli m x b -> TracerA m a b
Emitting ((a -> m ()) -> Kleisli m a ()
forall (m :: * -> *) a b. (a -> m b) -> Kleisli m a b
Kleisli a -> m ()
f) ((() -> m ()) -> Kleisli m () ()
forall (m :: * -> *) a b. (a -> m b) -> Kleisli m a b
Kleisli (m () -> () -> m ()
forall a b. a -> b -> a
const (() -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())))

-- | Do a non-emitting effect. This effect will only be run if some part of
-- the tracer downstream emits (see 'emit').
effect :: (a -> m b) -> TracerA m a b
effect :: (a -> m b) -> TracerA m a b
effect = Kleisli m a b -> TracerA m a b
forall (m :: * -> *) a b. Kleisli m a b -> TracerA m a b
Squelching (Kleisli m a b -> TracerA m a b)
-> ((a -> m b) -> Kleisli m a b) -> (a -> m b) -> TracerA m a b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (a -> m b) -> Kleisli m a b
forall (m :: * -> *) a b. (a -> m b) -> Kleisli m a b
Kleisli

-- | Pure computation in a tracer: no side effects or emits.
compute :: Applicative m => (a -> b) -> TracerA m a b
compute :: (a -> b) -> TracerA m a b
compute a -> b
f = (a -> m b) -> TracerA m a b
forall a (m :: * -> *) b. (a -> m b) -> TracerA m a b
effect (b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> m b) -> (a -> b) -> a -> m b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. a -> b
f)

instance Monad m => Category (TracerA m) where
  id :: TracerA m a a
id = (a -> a) -> TracerA m a a
forall (m :: * -> *) a b.
Applicative m =>
(a -> b) -> TracerA m a b
compute a -> a
forall k (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id
  Squelching Kleisli m b c
l     . :: TracerA m b c -> TracerA m a b -> TracerA m a c
. Squelching Kleisli m a b
r     = Kleisli m a c -> TracerA m a c
forall (m :: * -> *) a b. Kleisli m a b -> TracerA m a b
Squelching (Kleisli m b c
l  Kleisli m b c -> Kleisli m a b -> Kleisli m a c
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Kleisli m a b
r)
  -- Crucial: the squelching parts stay together. Could also have written
  --                                  = Emitting   (rp . re)      l
  -- but that would miss opportunities to skip doing work.
  Squelching Kleisli m b c
l     . Emitting   Kleisli m a x
re Kleisli m x b
rp = Kleisli m a x -> Kleisli m x c -> TracerA m a c
forall (m :: * -> *) a x b.
Kleisli m a x -> Kleisli m x b -> TracerA m a b
Emitting   Kleisli m a x
re             (Kleisli m b c
l Kleisli m b c -> Kleisli m x b -> Kleisli m x c
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Kleisli m x b
rp)
  -- Contrast with the above clause: here the emitting part comes _after_ the
  -- squelching part, so the squelching part becomes part of the emitting part.
  Emitting   Kleisli m b x
le Kleisli m x c
lp . Squelching Kleisli m a b
r     = Kleisli m a x -> Kleisli m x c -> TracerA m a c
forall (m :: * -> *) a x b.
Kleisli m a x -> Kleisli m x b -> TracerA m a b
Emitting   (Kleisli m b x
le Kleisli m b x -> Kleisli m a b -> Kleisli m a x
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Kleisli m a b
r)       Kleisli m x c
lp
  Emitting   Kleisli m b x
le Kleisli m x c
lp . Emitting   Kleisli m a x
re Kleisli m x b
rp = Kleisli m a x -> Kleisli m x c -> TracerA m a c
forall (m :: * -> *) a x b.
Kleisli m a x -> Kleisli m x b -> TracerA m a b
Emitting   (Kleisli m b x
le Kleisli m b x -> Kleisli m a b -> Kleisli m a x
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Kleisli m x b
rp Kleisli m x b -> Kleisli m a x -> Kleisli m a b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Kleisli m a x
re) Kleisli m x c
lp

instance Monad m => Arrow (TracerA m) where
  arr :: (b -> c) -> TracerA m b c
arr = (b -> c) -> TracerA m b c
forall (m :: * -> *) a b.
Applicative m =>
(a -> b) -> TracerA m a b
compute
  Squelching Kleisli m b c
l     *** :: TracerA m b c -> TracerA m b' c' -> TracerA m (b, b') (c, c')
*** Squelching Kleisli m b' c'
r     = Kleisli m (b, b') (c, c') -> TracerA m (b, b') (c, c')
forall (m :: * -> *) a b. Kleisli m a b -> TracerA m a b
Squelching (Kleisli m b c
l  Kleisli m b c -> Kleisli m b' c' -> Kleisli m (b, b') (c, c')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Kleisli m b' c'
r )
  Squelching Kleisli m b c
l     *** Emitting   Kleisli m b' x
re Kleisli m x c'
rp = Kleisli m (b, b') (b, x)
-> Kleisli m (b, x) (c, c') -> TracerA m (b, b') (c, c')
forall (m :: * -> *) a x b.
Kleisli m a x -> Kleisli m x b -> TracerA m a b
Emitting   (Kleisli m b b
forall k (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id Kleisli m b b -> Kleisli m b' x -> Kleisli m (b, b') (b, x)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Kleisli m b' x
re) (Kleisli m b c
l  Kleisli m b c -> Kleisli m x c' -> Kleisli m (b, x) (c, c')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Kleisli m x c'
rp)
  Emitting   Kleisli m b x
le Kleisli m x c
lp *** Squelching Kleisli m b' c'
r     = Kleisli m (b, b') (x, b')
-> Kleisli m (x, b') (c, c') -> TracerA m (b, b') (c, c')
forall (m :: * -> *) a x b.
Kleisli m a x -> Kleisli m x b -> TracerA m a b
Emitting   (Kleisli m b x
le Kleisli m b x -> Kleisli m b' b' -> Kleisli m (b, b') (x, b')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Kleisli m b' b'
forall k (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id) (Kleisli m x c
lp Kleisli m x c -> Kleisli m b' c' -> Kleisli m (x, b') (c, c')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Kleisli m b' c'
r )
  Emitting   Kleisli m b x
le Kleisli m x c
lp *** Emitting   Kleisli m b' x
re Kleisli m x c'
rp = Kleisli m (b, b') (x, x)
-> Kleisli m (x, x) (c, c') -> TracerA m (b, b') (c, c')
forall (m :: * -> *) a x b.
Kleisli m a x -> Kleisli m x b -> TracerA m a b
Emitting   (Kleisli m b x
le Kleisli m b x -> Kleisli m b' x -> Kleisli m (b, b') (x, x)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Kleisli m b' x
re) (Kleisli m x c
lp Kleisli m x c -> Kleisli m x c' -> Kleisli m (x, x) (c, c')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Kleisli m x c'
rp)

instance Monad m => ArrowChoice (TracerA m) where
  Squelching Kleisli m b c
l     +++ :: TracerA m b c
-> TracerA m b' c' -> TracerA m (Either b b') (Either c c')
+++ Squelching Kleisli m b' c'
r     = Kleisli m (Either b b') (Either c c')
-> TracerA m (Either b b') (Either c c')
forall (m :: * -> *) a b. Kleisli m a b -> TracerA m a b
Squelching (Kleisli m b c
l Kleisli m b c
-> Kleisli m b' c' -> Kleisli m (Either b b') (Either c c')
forall (a :: * -> * -> *) b c b' c'.
ArrowChoice a =>
a b c -> a b' c' -> a (Either b b') (Either c c')
+++ Kleisli m b' c'
r)
  Squelching Kleisli m b c
l     +++ Emitting   Kleisli m b' x
re Kleisli m x c'
rp = Kleisli m (Either b b') (Either b x)
-> Kleisli m (Either b x) (Either c c')
-> TracerA m (Either b b') (Either c c')
forall (m :: * -> *) a x b.
Kleisli m a x -> Kleisli m x b -> TracerA m a b
Emitting   (Kleisli m b b
forall k (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id Kleisli m b b
-> Kleisli m b' x -> Kleisli m (Either b b') (Either b x)
forall (a :: * -> * -> *) b c b' c'.
ArrowChoice a =>
a b c -> a b' c' -> a (Either b b') (Either c c')
+++ Kleisli m b' x
re) (Kleisli m b c
l  Kleisli m b c
-> Kleisli m x c' -> Kleisli m (Either b x) (Either c c')
forall (a :: * -> * -> *) b c b' c'.
ArrowChoice a =>
a b c -> a b' c' -> a (Either b b') (Either c c')
+++ Kleisli m x c'
rp)
  Emitting   Kleisli m b x
le Kleisli m x c
lp +++ Squelching Kleisli m b' c'
r     = Kleisli m (Either b b') (Either x b')
-> Kleisli m (Either x b') (Either c c')
-> TracerA m (Either b b') (Either c c')
forall (m :: * -> *) a x b.
Kleisli m a x -> Kleisli m x b -> TracerA m a b
Emitting   (Kleisli m b x
le Kleisli m b x
-> Kleisli m b' b' -> Kleisli m (Either b b') (Either x b')
forall (a :: * -> * -> *) b c b' c'.
ArrowChoice a =>
a b c -> a b' c' -> a (Either b b') (Either c c')
+++ Kleisli m b' b'
forall k (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id) (Kleisli m x c
lp Kleisli m x c
-> Kleisli m b' c' -> Kleisli m (Either x b') (Either c c')
forall (a :: * -> * -> *) b c b' c'.
ArrowChoice a =>
a b c -> a b' c' -> a (Either b b') (Either c c')
+++ Kleisli m b' c'
r )
  Emitting   Kleisli m b x
le Kleisli m x c
lp +++ Emitting   Kleisli m b' x
re Kleisli m x c'
rp = Kleisli m (Either b b') (Either x x)
-> Kleisli m (Either x x) (Either c c')
-> TracerA m (Either b b') (Either c c')
forall (m :: * -> *) a x b.
Kleisli m a x -> Kleisli m x b -> TracerA m a b
Emitting   (Kleisli m b x
le Kleisli m b x
-> Kleisli m b' x -> Kleisli m (Either b b') (Either x x)
forall (a :: * -> * -> *) b c b' c'.
ArrowChoice a =>
a b c -> a b' c' -> a (Either b b') (Either c c')
+++ Kleisli m b' x
re) (Kleisli m x c
lp Kleisli m x c
-> Kleisli m x c' -> Kleisli m (Either x x) (Either c c')
forall (a :: * -> * -> *) b c b' c'.
ArrowChoice a =>
a b c -> a b' c' -> a (Either b b') (Either c c')
+++ Kleisli m x c'
rp)

-- | Use a natural transformation to change the underlying monad.
nat :: (forall x . m x -> n x) -> TracerA m a b -> TracerA n a b
nat :: (forall x. m x -> n x) -> TracerA m a b -> TracerA n a b
nat forall x. m x -> n x
h (Squelching (Kleisli a -> m b
k))             = Kleisli n a b -> TracerA n a b
forall (m :: * -> *) a b. Kleisli m a b -> TracerA m a b
Squelching ((a -> n b) -> Kleisli n a b
forall (m :: * -> *) a b. (a -> m b) -> Kleisli m a b
Kleisli (m b -> n b
forall x. m x -> n x
h (m b -> n b) -> (a -> m b) -> a -> n b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. a -> m b
k))
nat forall x. m x -> n x
h (Emitting   (Kleisli a -> m x
k) (Kleisli x -> m b
l)) = Kleisli n a x -> Kleisli n x b -> TracerA n a b
forall (m :: * -> *) a x b.
Kleisli m a x -> Kleisli m x b -> TracerA m a b
Emitting   ((a -> n x) -> Kleisli n a x
forall (m :: * -> *) a b. (a -> m b) -> Kleisli m a b
Kleisli (m x -> n x
forall x. m x -> n x
h (m x -> n x) -> (a -> m x) -> a -> n x
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. a -> m x
k)) ((x -> n b) -> Kleisli n x b
forall (m :: * -> *) a b. (a -> m b) -> Kleisli m a b
Kleisli (m b -> n b
forall x. m x -> n x
h (m b -> n b) -> (x -> m b) -> x -> n b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. x -> m b
l))