{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeAbstractions #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Free traced monoidal category syntax.
--
-- This module packages the traced-monoidal layer on top of the generic
-- free-construction substrate in "Circuit.Syntax". The signature sum is
--
-- @
-- 'SigCompose' ':+:' 'SigYank' t
-- @
--
-- where 'SigCompose' provides sequential composition and 'SigYank' provides
-- feedback / trace over the channel tensor @t@.
--
-- Higher-level signatures (parallel composition, braid, copy\/discard,
-- shared-medium fusion, mediators) live in "Circuit.SMC", "Circuit.Bimonoid",
-- and "Circuit.Shared".
module Circuit.Trace
  ( -- * Free traced category
    Trace,
    base,
    yank,

    -- * Trace signature
    SigYank (..),
  )
where

import Circuit.Category (Category (..))
import Circuit.Channel (Channel (..), Strength (..), Traced (..))
import Circuit.Syntax
  ( Algebra (..),
    SigCompose (..),
    Syntax (..),
    eval,
    (:+:) (..),
  )
import Data.Kind (Type)
import Prelude hiding (id, (.))

-- Trace signature

-- | Feedback loop / trace over tensor @t@.
data SigYank (t :: Type -> Type -> Type) arr rec a b where
  Yank ::
    rec (t s a) (t s b) ->
    SigYank t arr rec a b

instance (Traced t arr') => Algebra (SigYank t) arr arr' where
  type Ctx (SigYank t) arr arr' = Traced t arr'
  alg ::
    forall (rec :: Type -> Type -> Type) (b :: Type) (c :: Type).
    (forall x y. arr x y -> arr' x y) ->
    (forall x y. rec x y -> arr' x y) ->
    SigYank t arr rec b c ->
    arr' b c
  alg :: forall (rec :: * -> * -> *) b c.
(forall x y. arr x y -> arr' x y)
-> (forall x y. rec x y -> arr' x y)
-> SigYank t arr rec b c
-> arr' b c
alg forall x y. arr x y -> arr' x y
_ forall x y. rec x y -> arr' x y
rec (Yank @_ @_ @_ @_ @_ @_ rec (t s b) (t s c)
f) = arr' (t s b) (t s c) -> arr' b c
forall a b c. arr' (t a b) (t a c) -> arr' b c
forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k)
       (b :: k) (c :: k).
Traced t arr =>
arr (t a b) (t a c) -> arr b c
trace (rec (t s b) (t s c) -> arr' (t s b) (t s c)
forall x y. rec x y -> arr' x y
rec rec (t s b) (t s c)
f)

-- Free traced monoidal category over tensor @t@

-- | Free traced monoidal category over tensor @t@.
type Trace t arr = Syntax (SigCompose :+: SigYank t) arr

-- | Lift a base arrow into the free traced category.
base :: arr a b -> Trace t arr a b
base :: forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base = arr a b -> Syntax (SigCompose :+: SigYank t) arr a b
forall (arr :: * -> * -> *) a b (sig :: Sig).
arr a b -> Syntax sig arr a b
Lift

-- | Close a feedback loop over the channel tensor @t@.
yank :: Trace t arr (t s a) (t s b) -> Trace t arr a b
yank :: forall (t :: * -> * -> *) (arr :: * -> * -> *) s a b.
Trace t arr (t s a) (t s b) -> Trace t arr a b
yank Trace t arr (t s a) (t s b)
body = (:+:)
  SigCompose
  (SigYank t)
  arr
  (Syntax (SigCompose :+: SigYank t) arr)
  a
  b
-> Syntax (SigCompose :+: SigYank t) arr a b
forall (sig :: Sig) (arr :: * -> * -> *) a b.
sig arr (Syntax sig arr) a b -> Syntax sig arr a b
Op (SigYank t arr (Syntax (SigCompose :+: SigYank t) arr) a b
-> (:+:)
     SigCompose
     (SigYank t)
     arr
     (Syntax (SigCompose :+: SigYank t) arr)
     a
     b
forall {k} {k1} {k2} {k3} (sig2 :: k -> k1 -> k2 -> k3 -> *)
       (arr :: k) (rec :: k1) (a :: k2) (b :: k3)
       (sig1 :: k -> k1 -> k2 -> k3 -> *).
sig2 arr rec a b -> (:+:) sig1 sig2 arr rec a b
R (Trace t arr (t s a) (t s b)
-> SigYank t arr (Syntax (SigCompose :+: SigYank t) arr) a b
forall {k} (rec :: * -> * -> *) (t :: * -> * -> *) s a b
       (arr :: k).
rec (t s a) (t s b) -> SigYank t arr rec a b
Yank Trace t arr (t s a) (t s b)
body))

-- Instances for the free traced category

instance (Category arr) => Category (Trace t arr) where
  id :: forall a. Trace t arr a a
id = arr a a -> Trace t arr a a
forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base arr a a
forall a. arr a a
forall k (arr :: k -> k -> *) (a :: k). Category arr => arr a a
id
  Trace t arr b c
f . :: forall b c a. Trace t arr b c -> Trace t arr a b -> Trace t arr a c
. Trace t arr a b
g = (:+:) SigCompose (SigYank t) arr (Trace t arr) a c
-> Syntax (SigCompose :+: SigYank t) arr a c
forall (sig :: Sig) (arr :: * -> * -> *) a b.
sig arr (Syntax sig arr) a b -> Syntax sig arr a b
Op (SigCompose arr (Trace t arr) a c
-> (:+:) SigCompose (SigYank t) arr (Trace t arr) a c
forall {k} {k1} {k2} {k3} (sig1 :: k -> k1 -> k2 -> k3 -> *)
       (arr :: k) (rec :: k1) (a :: k2) (b :: k3)
       (sig2 :: k -> k1 -> k2 -> k3 -> *).
sig1 arr rec a b -> (:+:) sig1 sig2 arr rec a b
L (Trace t arr b c
-> Trace t arr a b -> SigCompose arr (Trace t arr) a c
forall {k1} {k} (rec :: k1 -> k1 -> *) (b1 :: k1) (b :: k1)
       (a :: k1) (arr :: k).
rec b1 b -> rec a b1 -> SigCompose arr rec a b
SigCompose Trace t arr b c
f Trace t arr a b
g))

instance (Category arr, Channel t arr) => Channel t (Trace t arr) where
  assoc :: forall a b c. Trace t arr (t (t a b) c) (t a (t b c))
assoc = arr (t (t a b) c) (t a (t b c))
-> Trace t arr (t (t a b) c) (t a (t b c))
forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base arr (t (t a b) c) (t a (t b c))
forall a b c. arr (t (t a b) c) (t a (t b c))
forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k)
       (b :: k) (c :: k).
Channel t arr =>
arr (t (t a b) c) (t a (t b c))
assoc
  assoc' :: forall a b c. Trace t arr (t a (t b c)) (t (t a b) c)
assoc' = arr (t a (t b c)) (t (t a b) c)
-> Trace t arr (t a (t b c)) (t (t a b) c)
forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base arr (t a (t b c)) (t (t a b) c)
forall a b c. arr (t a (t b c)) (t (t a b) c)
forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k)
       (b :: k) (c :: k).
Channel t arr =>
arr (t a (t b c)) (t (t a b) c)
assoc'
  slide :: forall a b c. Trace t arr (t a (t b c)) (t b (t a c))
slide = arr (t a (t b c)) (t b (t a c))
-> Trace t arr (t a (t b c)) (t b (t a c))
forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base arr (t a (t b c)) (t b (t a c))
forall a b c. arr (t a (t b c)) (t b (t a c))
forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k)
       (b :: k) (c :: k).
Channel t arr =>
arr (t a (t b c)) (t b (t a c))
slide

instance (Strength t arr, Traced t arr) => Strength t (Trace t arr) where
  strength :: forall b c a. Trace t arr b c -> Trace t arr (t a b) (t a c)
strength Trace t arr b c
f = arr (t a b) (t a c) -> Trace t arr (t a b) (t a c)
forall (arr :: * -> * -> *) a b (t :: * -> * -> *).
arr a b -> Trace t arr a b
base (arr b c -> arr (t a b) (t a c)
forall b c a. arr b c -> arr (t a b) (t a c)
forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (b :: k)
       (c :: k) (a :: k).
Strength t arr =>
arr b c -> arr (t a b) (t a c)
strength (Trace t arr b c -> arr b c
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 arr b c
f))

instance (Traced t arr) => Traced t (Trace t arr) where
  trace :: forall a b c. Trace t arr (t a b) (t a c) -> Trace t arr b c
trace = Trace t arr (t a b) (t a c) -> Trace t arr b c
forall (t :: * -> * -> *) (arr :: * -> * -> *) s a b.
Trace t arr (t s a) (t s b) -> Trace t arr a b
yank