{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

-- | Shared-medium fusion: two bodies interleaved on one shared channel.
--
-- The connective here is the multiplicative disjunction @⅋@ in operational
-- form: two sub-loops share a single channel, and a 'Schedule' resolves the
-- interleaving.  This is the mixed-mode counterpart to
-- 'Circuit.Tensor.superpose', which keeps channels independent
-- (the @⊗@ product).
--
-- 'Bias' is re-exported from "Circuit.Tensor" because it is also used for
-- additive disjunction in "Circuit.Poles".
module Circuit.Shared
  ( -- * Schedule bias
    Bias (..),

    -- * Schedule decision
    Pick (..),

    -- * Schedule driver
    Schedule (..),

    -- * Shared fusion class
    Shared (..),

    -- * Shared-medium fusion signature
    SigShared (..),

    -- * Free traced category with shared-medium fusion
    AlgShared,
  )
where

import Circuit.Category (Category (..), K (..))
import Circuit.Channel (Channel (..))
import Circuit.Syntax
  ( Algebra (..),
    SigCompose,
    Syntax,
    (:+:),
  )
import Circuit.Tensor (Bias (..), Tensor (..))
import Circuit.Trace (SigYank)
import Control.Monad (Monad)
import Data.Kind (Type)
import Data.These (These (..))
import Prelude hiding (id, (.))

-- | A schedule decision: which poles advance on a shared channel.
--
-- * @L@ — advance the left body only; the right input is not consumed
--   (corresponds to 'This').
-- * @R@ — advance the right body only; the left input is not consumed
--   (corresponds to 'That').
-- * @Both b@ — advance both bodies, with the bias choosing the order
--   (corresponds to 'These').
data Pick = L | R | Both Bias
  deriving (Pick -> Pick -> Bool
(Pick -> Pick -> Bool) -> (Pick -> Pick -> Bool) -> Eq Pick
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Pick -> Pick -> Bool
== :: Pick -> Pick -> Bool
$c/= :: Pick -> Pick -> Bool
/= :: Pick -> Pick -> Bool
Eq, Int -> Pick -> ShowS
[Pick] -> ShowS
Pick -> String
(Int -> Pick -> ShowS)
-> (Pick -> String) -> ([Pick] -> ShowS) -> Show Pick
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Pick -> ShowS
showsPrec :: Int -> Pick -> ShowS
$cshow :: Pick -> String
show :: Pick -> String
$cshowList :: [Pick] -> ShowS
showList :: [Pick] -> ShowS
Show)

-- | A schedule drives shared-medium fusion.
--
-- The state @s@ is threaded through the fusion; in typical use it is the
-- shared channel.  At each step the schedule looks at the state and chooses
-- which poles advance, returning the updated schedule state.
newtype Schedule s = Schedule
  { -- | Given the current shared state, return the updated state and a 'Pick'
    -- value describing which poles advance and in what order.
    forall s. Schedule s -> s -> (s, Pick)
chooseS :: s -> (s, Pick)
  }

-- | Tensors that support shared-medium fusion of two knot bodies.
--
-- This is the operational content of the multiplicative disjunction: two
-- sub-loops share one channel, and a 'Schedule' resolves the interleaving.
-- Contrast 'Circuit.Tensor.superpose', which keeps the channels independent (⊗).
class (Tensor t arr) => Shared t arr where
  -- | Fuse two knot bodies over a shared channel.
  --
  -- The combined body has type @arr (t s (t a c)) (t s (These b d))@: one
  -- shared state @s@, paired inputs @a@ and @c@, and a partial output.  At
  -- each step the schedule chooses which body advances; the gated body's
  -- input is discarded and no output is produced for that side.
  sharedBy ::
    Schedule s ->
    arr (t s a) (t s b) ->
    arr (t s c) (t s d) ->
    arr (t s (t a c)) (t s (These b d))

-- | Cartesian shared fusion on functions.
--
-- The schedule chooses which bodies advance and in what order.  @L@/@R@
-- run only the chosen body and emit a partial 'This'/'That' product; the
-- other body's input is discarded.  @Both LeftFirst@ / @Both RightFirst@ run
-- both bodies, threading the shared state in the chosen order, and emit a
-- total 'These' product.  When both bodies read and write @s@, the two orders
-- are observationally different — this is the ⅋-vs-⊗ distinction.
instance Shared (,) (->) where
  sharedBy :: forall s a b c d.
Schedule s
-> ((s, a) -> (s, b))
-> ((s, c) -> (s, d))
-> (s, (a, c)) -> (s, These b d)
sharedBy Schedule s
sched (s, a) -> (s, b)
f (s, c) -> (s, d)
g (s
s, (a
a, c
c)) =
    let (s
s', Pick
pick) = Schedule s -> s -> (s, Pick)
forall s. Schedule s -> s -> (s, Pick)
chooseS Schedule s
sched s
s
     in case Pick
pick of
          Pick
L ->
            let (s
s'', b
b) = (s, a) -> (s, b)
f (s
s', a
a)
             in (s
s'', b -> These b d
forall a b. a -> These a b
This b
b)
          Pick
R ->
            let (s
s'', d
d) = (s, c) -> (s, d)
g (s
s', c
c)
             in (s
s'', d -> These b d
forall a b. b -> These a b
That d
d)
          Both Bias
LeftFirst ->
            let (s
s'', b
b) = (s, a) -> (s, b)
f (s
s', a
a)
                (s
s''', d
d) = (s, c) -> (s, d)
g (s
s'', c
c)
             in (s
s''', b -> d -> These b d
forall a b. a -> b -> These a b
These b
b d
d)
          Both Bias
RightFirst ->
            let (s
s'', d
d) = (s, c) -> (s, d)
g (s
s', c
c)
                (s
s''', b
b) = (s, a) -> (s, b)
f (s
s'', a
a)
             in (s
s''', b -> d -> These b d
forall a b. a -> b -> These a b
These b
b d
d)
  {-# INLINE sharedBy #-}

-- | Cartesian shared fusion on @K@ arrows.
instance (Monad m) => Shared (,) (K m) where
  sharedBy :: forall s a b c d.
Schedule s
-> K m (s, a) (s, b)
-> K m (s, c) (s, d)
-> K m (s, (a, c)) (s, These b d)
sharedBy Schedule s
sched (K (s, a) -> m (s, b)
f) (K (s, c) -> m (s, d)
g) =
    ((s, (a, c)) -> m (s, These b d)) -> K m (s, (a, c)) (s, These b d)
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K (((s, (a, c)) -> m (s, These b d))
 -> K m (s, (a, c)) (s, These b d))
-> ((s, (a, c)) -> m (s, These b d))
-> K m (s, (a, c)) (s, These b d)
forall a b. (a -> b) -> a -> b
$ \(s
s, (a
a, c
c)) -> do
      let (s
s', Pick
pick) = Schedule s -> s -> (s, Pick)
forall s. Schedule s -> s -> (s, Pick)
chooseS Schedule s
sched s
s
      case Pick
pick of
        Pick
L -> do
          (s'', b) <- (s, a) -> m (s, b)
f (s
s', a
a)
          pure (s'', This b)
        Pick
R -> do
          (s'', d) <- (s, c) -> m (s, d)
g (s
s', c
c)
          pure (s'', That d)
        Both Bias
LeftFirst -> do
          (s'', b) <- (s, a) -> m (s, b)
f (s
s', a
a)
          (s''', d) <- g (s'', c)
          pure (s''', These b d)
        Both Bias
RightFirst -> do
          (s'', d) <- (s, c) -> m (s, d)
g (s
s', c
c)
          (s''', b) <- f (s'', a)
          pure (s''', These b d)
  {-# INLINE sharedBy #-}

-- Shared-medium fusion signature

-- | Shared-medium fusion (the tensor product ⅋), parameterised by a schedule.
--
-- The constructor takes two bodies that already share a feedback type @s@ and
-- produces the untraced shared body.  The surrounding 'SigYank' closes the
-- feedback loop over @s@, yielding a morphism @t a c -> These b d@.
data SigShared (t :: Type -> Type -> Type) arr rec i o where
  SigShared ::
    Schedule s ->
    rec (t s a) (t s b) ->
    rec (t s c) (t s d) ->
    SigShared t arr rec (t s (t a c)) (t s (These b d))

instance (Shared t arr') => Algebra (SigShared t) arr arr' where
  type Ctx (SigShared t) arr arr' = Shared t arr'
  alg :: forall (rec :: * -> * -> *) a b.
Ctx (SigShared t) arr arr' =>
(forall x y. arr x y -> arr' x y)
-> (forall x y. rec x y -> arr' x y)
-> SigShared t arr rec a b
-> arr' a b
alg forall x y. arr x y -> arr' x y
_ forall x y. rec x y -> arr' x y
rec (SigShared Schedule s
sched rec (t s a) (t s b)
f rec (t s c) (t s d)
g) = Schedule s
-> arr' (t s a) (t s b)
-> arr' (t s c) (t s d)
-> arr' (t s (t a c)) (t s (These b d))
forall s a b c d.
Schedule s
-> arr' (t s a) (t s b)
-> arr' (t s c) (t s d)
-> arr' (t s (t a c)) (t s (These b d))
forall (t :: * -> * -> *) (arr :: * -> * -> *) s a b c d.
Shared t arr =>
Schedule s
-> arr (t s a) (t s b)
-> arr (t s c) (t s d)
-> arr (t s (t a c)) (t s (These b d))
sharedBy Schedule s
sched (rec (t s a) (t s b) -> arr' (t s a) (t s b)
forall x y. rec x y -> arr' x y
rec rec (t s a) (t s b)
f) (rec (t s c) (t s d) -> arr' (t s c) (t s d)
forall x y. rec x y -> arr' x y
rec rec (t s c) (t s d)
g)

-- | Free traced category with shared-medium fusion (the ⅋ connective).
type AlgShared t arr = Syntax (SigCompose :+: SigShared t :+: SigYank t) arr