{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | Probability as a double-dual continuation category.
--
-- A morphism @Prob arr r a b@ is an expectation transformer: it turns a
-- continuation @arr (x, b) r@ (a "test" on the output) into a continuation
-- @arr (x, a) r@ (a test on the input).  The rank-2 quantification over @x@
-- is the cost of arrow-polymorphism — the same move used by 'Circuit.Poles'.
--
-- This is the categorical substrate for probability, conditioning, and
-- verification: choosing the dualizing object @r@ picks the semantics.
--
-- * @r = Log Double@ over @(->)@ gives expectation transformers / measures.
-- * @r = Bool@ over @(->)@ gives Dijkstra's weakest-precondition semantics.
-- * @r = Min Double@ (tropical) gives Bellman / Viterbi / MAP semantics.
--
-- This module currently provides structural instances for the function arrow
-- @(->)@.  Effectful variants (e.g. @Kleisli m@) follow the same pattern but
-- need scalar-lifting plumbing; the function case is where the design is
-- easiest to validate.
--
-- The tensor action on @Prob@ is /premonoidal/ in general: two valid nestings
-- ('parFG' and 'parGF') agree only on the linear (commutative) fragment.  We
-- therefore do not provide a canonical 'Circuit.Tensor.Tensor' instance; use
-- the explicit nesting you mean.
module Circuit.Prob
  ( -- * Double-dual probability arrow
    Prob (..),

    -- * Primitive constructors
    embed,
    fromWeighted,
    score,
    mass,

    -- * Cartesian copy/discard (deterministic)
    copyP,
    discardP,

    -- * Choice combined by a scalar operation
    choiceBy,
    orP,

    -- * Parallel nestings (Fubini on the linear fragment)
    parFG,
    parGF,

    -- * Traced Either (explicit, computability varies by scalar)
    traceE,
    traceEN,

    -- * Semiring scalars
    Semiring (..),
    Tropical (..),
  )
where

import Circuit.Category (Category (..))
import Circuit.Channel (Channel (..), Strength (..))
import Circuit.Tensor (Unit, Unital (..))
import Data.Bifunctor (second)
import Prelude hiding (id, (.))
import Prelude qualified

-- $setup
-- >>> import Circuit.Prob
-- >>> import Prelude hiding (id, (.))

-- | A semiring: an additive monoid and a multiplicative monoid, with
-- multiplication distributing over addition.
--
-- This class is intentionally minimal. It captures the scalar structure
-- needed by 'Circuit.Prob' without pulling in a full numeric prelude.
class Semiring r where
  sAdd :: r -> r -> r
  sMul :: r -> r -> r
  sZero :: r
  sOne :: r

-- | Min-plus tropical semiring over 'Double'.
--
-- Addition is 'min', multiplication is ordinary addition, the additive unit
-- is positive infinity, and the multiplicative unit is zero.
newtype Tropical = Tropical {Tropical -> Double
getTropical :: Double}
  deriving (Tropical -> Tropical -> Bool
(Tropical -> Tropical -> Bool)
-> (Tropical -> Tropical -> Bool) -> Eq Tropical
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Tropical -> Tropical -> Bool
== :: Tropical -> Tropical -> Bool
$c/= :: Tropical -> Tropical -> Bool
/= :: Tropical -> Tropical -> Bool
Eq, Eq Tropical
Eq Tropical =>
(Tropical -> Tropical -> Ordering)
-> (Tropical -> Tropical -> Bool)
-> (Tropical -> Tropical -> Bool)
-> (Tropical -> Tropical -> Bool)
-> (Tropical -> Tropical -> Bool)
-> (Tropical -> Tropical -> Tropical)
-> (Tropical -> Tropical -> Tropical)
-> Ord Tropical
Tropical -> Tropical -> Bool
Tropical -> Tropical -> Ordering
Tropical -> Tropical -> Tropical
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Tropical -> Tropical -> Ordering
compare :: Tropical -> Tropical -> Ordering
$c< :: Tropical -> Tropical -> Bool
< :: Tropical -> Tropical -> Bool
$c<= :: Tropical -> Tropical -> Bool
<= :: Tropical -> Tropical -> Bool
$c> :: Tropical -> Tropical -> Bool
> :: Tropical -> Tropical -> Bool
$c>= :: Tropical -> Tropical -> Bool
>= :: Tropical -> Tropical -> Bool
$cmax :: Tropical -> Tropical -> Tropical
max :: Tropical -> Tropical -> Tropical
$cmin :: Tropical -> Tropical -> Tropical
min :: Tropical -> Tropical -> Tropical
Ord, Int -> Tropical -> ShowS
[Tropical] -> ShowS
Tropical -> String
(Int -> Tropical -> ShowS)
-> (Tropical -> String) -> ([Tropical] -> ShowS) -> Show Tropical
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Tropical -> ShowS
showsPrec :: Int -> Tropical -> ShowS
$cshow :: Tropical -> String
show :: Tropical -> String
$cshowList :: [Tropical] -> ShowS
showList :: [Tropical] -> ShowS
Show)

instance Semiring Tropical where
  sAdd :: Tropical -> Tropical -> Tropical
sAdd (Tropical Double
a) (Tropical Double
b) = Double -> Tropical
Tropical (Double -> Double -> Double
forall a. Ord a => a -> a -> a
Prelude.min Double
a Double
b)
  sMul :: Tropical -> Tropical -> Tropical
sMul (Tropical Double
a) (Tropical Double
b) = Double -> Tropical
Tropical (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b)
  sZero :: Tropical
sZero = Double -> Tropical
Tropical (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
  sOne :: Tropical
sOne = Double -> Tropical
Tropical Double
0

-- | 'Double' is the usual probability semiring.
instance Semiring Double where
  sAdd :: Double -> Double -> Double
sAdd = Double -> Double -> Double
forall a. Num a => a -> a -> a
(+)
  sMul :: Double -> Double -> Double
sMul = Double -> Double -> Double
forall a. Num a => a -> a -> a
(*)
  sZero :: Double
sZero = Double
0
  sOne :: Double
sOne = Double
1

-- | 'Bool' is the reachability / model-checking semiring.
instance Semiring Bool where
  sAdd :: Bool -> Bool -> Bool
sAdd = Bool -> Bool -> Bool
(||)
  sMul :: Bool -> Bool -> Bool
sMul = Bool -> Bool -> Bool
(&&)
  sZero :: Bool
sZero = Bool
False
  sOne :: Bool
sOne = Bool
True

-- | Double-dual embedding of @arr@ with respect to dualizing object @r@.
--
-- A value @Prob arr r a b@ reads an output-continuation @arr (x, b) r@ and
-- produces an input-continuation @arr (x, a) r@.  Composition is continuation
-- composition (contravariant in the middle type).
newtype Prob arr r a b = Prob
  { forall {k} (arr :: * -> k -> *) (r :: k) a b.
Prob arr r a b -> forall x. arr (x, b) r -> arr (x, a) r
runProb :: forall x. arr (x, b) r -> arr (x, a) r
  }

-- ---------------------------------------------------------------------------
-- Category
-- ---------------------------------------------------------------------------

-- | Identity and composition are arrow-polymorphic: they only manipulate the
-- continuation function, never the base arrow.  This is why 'Category' costs
-- nothing from @arr@.
instance (Category arr) => Category (Prob arr r) where
  id :: Prob arr r a a
  id :: forall a. Prob arr r a a
id = (forall x. arr (x, a) r -> arr (x, a) r) -> Prob arr r a a
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob arr (x, a) r -> arr (x, a) r
forall a. a -> a
forall x. arr (x, a) r -> arr (x, a) r
Prelude.id
  {-# INLINE id #-}

  (.) ::
    Prob arr r b c ->
    Prob arr r a b ->
    Prob arr r a c
  Prob forall x. arr (x, c) r -> arr (x, b) r
f . :: forall b c a. Prob arr r b c -> Prob arr r a b -> Prob arr r a c
. Prob forall x. arr (x, b) r -> arr (x, a) r
g = (forall x. arr (x, c) r -> arr (x, a) r) -> Prob arr r a c
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x. arr (x, c) r -> arr (x, a) r) -> Prob arr r a c)
-> (forall x. arr (x, c) r -> arr (x, a) r) -> Prob arr r a c
forall a b. (a -> b) -> a -> b
$ \arr (x, c) r
k -> arr (x, b) r -> arr (x, a) r
forall x. arr (x, b) r -> arr (x, a) r
g (arr (x, c) r -> arr (x, b) r
forall x. arr (x, c) r -> arr (x, b) r
f arr (x, c) r
k)
  {-# INLINE (.) #-}

-- ---------------------------------------------------------------------------
-- Primitives (function arrow)
-- ---------------------------------------------------------------------------

-- | Embed a deterministic function as a probability morphism.
--
-- The continuation is applied to the transformed output, with the context
-- wire carried along unchanged.
embed :: (a -> b) -> Prob (->) r a b
embed :: forall a b r. (a -> b) -> Prob (->) r a b
embed a -> b
h = (forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b)
-> (forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b
forall a b. (a -> b) -> a -> b
$ \(x, b) -> r
k -> (x, b) -> r
k ((x, b) -> r) -> ((x, a) -> (x, b)) -> (x, a) -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. (a -> b) -> (x, a) -> (x, b)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second a -> b
h
{-# INLINE embed #-}

-- | Build a probability morphism from a finite weighted table.
--
-- This is the bridge to 'Circuit.Parser.Weighted' and the entry point for
-- genuine measures in the linear fragment: every entry contributes linearly
-- to the expectation.
fromWeighted :: (Semiring r) => [(b, r)] -> Prob (->) r () b
fromWeighted :: forall r b. Semiring r => [(b, r)] -> Prob (->) r () b
fromWeighted [(b, r)]
xs = (forall x. ((x, b) -> r) -> (x, ()) -> r) -> Prob (->) r () b
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x. ((x, b) -> r) -> (x, ()) -> r) -> Prob (->) r () b)
-> (forall x. ((x, b) -> r) -> (x, ()) -> r) -> Prob (->) r () b
forall a b. (a -> b) -> a -> b
$ \(x, b) -> r
k (x
x, ()) -> (r -> r -> r) -> r -> [r] -> r
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr r -> r -> r
forall r. Semiring r => r -> r -> r
sAdd r
forall r. Semiring r => r
sZero [r -> r -> r
forall r. Semiring r => r -> r -> r
sMul r
w ((x, b) -> r
k (x
x, b
b)) | (b
b, r
w) <- [(b, r)]
xs]
{-# INLINE fromWeighted #-}

-- | Scale the result of a continuation.
--
-- With endomorphisms @r -> r@ this is a /modality/, not necessarily a scalar
-- multiplication.  The definitional law is the anti-homomorphism
-- @score w . score v = score (v . w)@; commutativity holds only when the
-- endos commute.  For the probabilistic sub-case @score (w *)@, the usual
-- multiplicative law is recovered.
score :: (r -> r) -> Prob (->) r a a
score :: forall r a. (r -> r) -> Prob (->) r a a
score r -> r
scale = (forall x. ((x, a) -> r) -> (x, a) -> r) -> Prob (->) r a a
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x. ((x, a) -> r) -> (x, a) -> r) -> Prob (->) r a a)
-> (forall x. ((x, a) -> r) -> (x, a) -> r) -> Prob (->) r a a
forall a b. (a -> b) -> a -> b
$ \(x, a) -> r
k (x
x, a
a) -> r -> r
scale ((x, a) -> r
k (x
x, a
a))
{-# INLINE score #-}

-- | Compute the total mass of an unnormalised morphism against the unit
-- continuation.
mass :: (Semiring r) => Prob (->) r a b -> a -> r
mass :: forall r a b. Semiring r => Prob (->) r a b -> a -> r
mass (Prob forall x. ((x, b) -> r) -> (x, a) -> r
f) a
a = (((), b) -> r) -> ((), a) -> r
forall x. ((x, b) -> r) -> (x, a) -> r
f (r -> ((), b) -> r
forall a b. a -> b -> a
const r
forall r. Semiring r => r
sOne) ((), a
a)
{-# INLINE mass #-}

-- | Deterministic copy.  Naturality of this morphism characterises the
-- deterministic fragment: @copyP . embed h == parFG (embed h) (embed h) . copyP@.
copyP :: Prob (->) r a (a, a)
copyP :: forall r a. Prob (->) r a (a, a)
copyP = (a -> (a, a)) -> Prob (->) r a (a, a)
forall a b r. (a -> b) -> Prob (->) r a b
embed (\a
a -> (a
a, a
a))
{-# INLINE copyP #-}

-- | Deterministic discard.  On the mass-1 fragment @f . discardP == discardP@;
-- unnormalised morphisms fail this equation.
discardP :: Prob (->) r a ()
discardP :: forall r a. Prob (->) r a ()
discardP = (a -> ()) -> Prob (->) r a ()
forall a b r. (a -> b) -> Prob (->) r a b
embed (() -> a -> ()
forall a b. a -> b -> a
const ())
{-# INLINE discardP #-}

-- | Binary choice combined by a scalar operation.  This one combinator covers
-- several rows of the instance table:
--
-- * @choiceBy (||)@ — angelic / reachability (Bool).
-- * @choiceBy (&&)@ — demonic / refutation (Bool).
-- * @choiceBy (+)@  — sum of weighted alternatives (Num r).
-- * @choiceBy min@  — tropical / Viterbi choice (Ord r).
choiceBy :: (r -> r -> r) -> Prob (->) r a b -> Prob (->) r a b -> Prob (->) r a b
choiceBy :: forall r a b.
(r -> r -> r)
-> Prob (->) r a b -> Prob (->) r a b -> Prob (->) r a b
choiceBy r -> r -> r
(<+>) (Prob forall x. ((x, b) -> r) -> (x, a) -> r
f) (Prob forall x. ((x, b) -> r) -> (x, a) -> r
g) = (forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b)
-> (forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b
forall a b. (a -> b) -> a -> b
$ \(x, b) -> r
k (x, a)
p -> ((x, b) -> r) -> (x, a) -> r
forall x. ((x, b) -> r) -> (x, a) -> r
f (x, b) -> r
k (x, a)
p r -> r -> r
<+> ((x, b) -> r) -> (x, a) -> r
forall x. ((x, b) -> r) -> (x, a) -> r
g (x, b) -> r
k (x, a)
p
{-# INLINE choiceBy #-}

-- | Angelic choice for @r = Bool@ (weakest-precondition / reachability
-- semantics).  Succeeds if either branch can; short-circuiting of @(||)@
-- gives the trace on this scalar for free.
orP :: Prob (->) Bool a b -> Prob (->) Bool a b -> Prob (->) Bool a b
orP :: forall a b.
Prob (->) Bool a b -> Prob (->) Bool a b -> Prob (->) Bool a b
orP = (Bool -> Bool -> Bool)
-> Prob (->) Bool a b -> Prob (->) Bool a b -> Prob (->) Bool a b
forall r a b.
(r -> r -> r)
-> Prob (->) r a b -> Prob (->) r a b -> Prob (->) r a b
choiceBy Bool -> Bool -> Bool
(||)
{-# INLINE orP #-}

-- ---------------------------------------------------------------------------
-- Structural instances (cartesian tensor, function arrow)
-- ---------------------------------------------------------------------------

-- | The cartesian structural morphisms are deterministic, so they are just
-- 'embed's of the base-arrow associators and braiding.  'strength' is the
-- non-trivial one: it instantiates the rank-2 context @x@ at @(x, s)@,
-- which is exactly why the universally quantified context is the honest cost
-- of the tensor.
instance Channel (,) (Prob (->) r) where
  assoc :: forall a b c. Prob (->) r ((a, b), c) (a, (b, c))
assoc = (((a, b), c) -> (a, (b, c))) -> Prob (->) r ((a, b), c) (a, (b, c))
forall a b r. (a -> b) -> Prob (->) r a b
embed ((a, b), c) -> (a, (b, c))
forall a b c. ((a, b), c) -> (a, (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
  {-# INLINE assoc #-}

  assoc' :: forall a b c. Prob (->) r (a, (b, c)) ((a, b), c)
assoc' = ((a, (b, c)) -> ((a, b), c)) -> Prob (->) r (a, (b, c)) ((a, b), c)
forall a b r. (a -> b) -> Prob (->) r a b
embed (a, (b, c)) -> ((a, b), c)
forall a b c. (a, (b, c)) -> ((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'
  {-# INLINE assoc' #-}

  slide :: forall a b c. Prob (->) r (a, (b, c)) (b, (a, c))
slide = ((a, (b, c)) -> (b, (a, c))) -> Prob (->) r (a, (b, c)) (b, (a, c))
forall a b r. (a -> b) -> Prob (->) r a b
embed (a, (b, c)) -> (b, (a, c))
forall a b c. (a, (b, c)) -> (b, (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
  {-# INLINE slide #-}

instance Strength (,) (Prob (->) r) where
  strength :: forall b c a. Prob (->) r b c -> Prob (->) r (a, b) (a, c)
strength (Prob forall x. ((x, c) -> r) -> (x, b) -> r
f) = (forall x. ((x, (a, c)) -> r) -> (x, (a, b)) -> r)
-> Prob (->) r (a, b) (a, c)
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x. ((x, (a, c)) -> r) -> (x, (a, b)) -> r)
 -> Prob (->) r (a, b) (a, c))
-> (forall x. ((x, (a, c)) -> r) -> (x, (a, b)) -> r)
-> Prob (->) r (a, b) (a, c)
forall a b. (a -> b) -> a -> b
$ \(x, (a, c)) -> r
k -> (((x, a), c) -> r) -> ((x, a), b) -> r
forall x. ((x, c) -> r) -> (x, b) -> r
f ((x, (a, c)) -> r
k ((x, (a, c)) -> r)
-> (((x, a), c) -> (x, (a, c))) -> ((x, a), c) -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. ((x, a), c) -> (x, (a, c))
forall a b c. ((a, b), c) -> (a, (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) (((x, a), b) -> r)
-> ((x, (a, b)) -> ((x, a), b)) -> (x, (a, b)) -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. (x, (a, b)) -> ((x, a), b)
forall a b c. (a, (b, c)) -> ((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'
  {-# INLINE strength #-}

-- | The cartesian unitors are deterministic, so they embed cleanly.  This is
-- why the 'Unital'/'Tensor' split matters: @Prob (->) r@ is premonoidal and
-- refuses a canonical 'Tensor' instance, but it still has honest unitors.
instance Unital (,) (Prob (->) r) where
  unitl :: forall a. Prob (->) r (Unit (,), a) a
unitl = (((), a) -> a) -> Prob (->) r ((), a) a
forall a b r. (a -> b) -> Prob (->) r a b
embed ((), a) -> a
(Unit (,), a) -> a
forall a. (Unit (,), a) -> a
forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k).
Unital t arr =>
arr (t (Unit t) a) a
unitl
  {-# INLINE unitl #-}
  unitl' :: forall a. Prob (->) r a (Unit (,), a)
unitl' = (a -> ((), a)) -> Prob (->) r a ((), a)
forall a b r. (a -> b) -> Prob (->) r a b
embed a -> ((), a)
a -> (Unit (,), a)
forall a. a -> (Unit (,), a)
forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k).
Unital t arr =>
arr a (t (Unit t) a)
unitl'
  {-# INLINE unitl' #-}
  unitr :: forall a. Prob (->) r (a, Unit (,)) a
unitr = ((a, ()) -> a) -> Prob (->) r (a, ()) a
forall a b r. (a -> b) -> Prob (->) r a b
embed (a, ()) -> a
(a, Unit (,)) -> a
forall a. (a, Unit (,)) -> a
forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k).
Unital t arr =>
arr (t a (Unit t)) a
unitr
  {-# INLINE unitr #-}
  unitr' :: forall a. Prob (->) r a (a, Unit (,))
unitr' = (a -> (a, ())) -> Prob (->) r a (a, ())
forall a b r. (a -> b) -> Prob (->) r a b
embed a -> (a, ())
a -> (a, Unit (,))
forall a. a -> (a, Unit (,))
forall {k} (t :: k -> k -> k) (arr :: k -> k -> *) (a :: k).
Unital t arr =>
arr a (t a (Unit t))
unitr'
  {-# INLINE unitr' #-}

-- ---------------------------------------------------------------------------
-- Parallel nestings (Fubini on the linear fragment)
-- ---------------------------------------------------------------------------

-- | Parallel composition: @g@ runs at context @(x, b)@, @f@ runs at context
-- @(x, c)@.  This is one of two lawful nestings; it agrees with 'parGF' on
-- the linear/commutative fragment.
parFG ::
  Prob (->) r a b ->
  Prob (->) r c d ->
  Prob (->) r (a, c) (b, d)
parFG :: forall r a b c d.
Prob (->) r a b -> Prob (->) r c d -> Prob (->) r (a, c) (b, d)
parFG (Prob forall x. ((x, b) -> r) -> (x, a) -> r
f) (Prob forall x. ((x, d) -> r) -> (x, c) -> r
g) = (forall x. ((x, (b, d)) -> r) -> (x, (a, c)) -> r)
-> Prob (->) r (a, c) (b, d)
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x. ((x, (b, d)) -> r) -> (x, (a, c)) -> r)
 -> Prob (->) r (a, c) (b, d))
-> (forall x. ((x, (b, d)) -> r) -> (x, (a, c)) -> r)
-> Prob (->) r (a, c) (b, d)
forall a b. (a -> b) -> a -> b
$ \(x, (b, d)) -> r
k ->
  let kg :: ((x, b), d) -> r
kg ((x
ctx, b
b), d
d) = (x, (b, d)) -> r
k (x
ctx, (b
b, d
d))
      gc :: ((x, b), c) -> r
gc = (((x, b), d) -> r) -> ((x, b), c) -> r
forall x. ((x, d) -> r) -> (x, c) -> r
g ((x, b), d) -> r
kg
      kf :: ((x, c), b) -> r
kf ((x
ctx, c
c), b
b) = ((x, b), c) -> r
gc ((x
ctx, b
b), c
c)
      fa :: ((x, c), a) -> r
fa = (((x, c), b) -> r) -> ((x, c), a) -> r
forall x. ((x, b) -> r) -> (x, a) -> r
f ((x, c), b) -> r
kf
   in \(x
ctx, (a
a, c
c)) -> ((x, c), a) -> r
fa ((x
ctx, c
c), a
a)
{-# INLINE parFG #-}

-- | Parallel composition: @f@ runs at context @(x, d)@, @g@ runs at context
-- @(x, a)@.  The other nesting; agrees with 'parFG' on the linear fragment.
parGF ::
  Prob (->) r a b ->
  Prob (->) r c d ->
  Prob (->) r (a, c) (b, d)
parGF :: forall r a b c d.
Prob (->) r a b -> Prob (->) r c d -> Prob (->) r (a, c) (b, d)
parGF (Prob forall x. ((x, b) -> r) -> (x, a) -> r
f) (Prob forall x. ((x, d) -> r) -> (x, c) -> r
g) = (forall x. ((x, (b, d)) -> r) -> (x, (a, c)) -> r)
-> Prob (->) r (a, c) (b, d)
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x. ((x, (b, d)) -> r) -> (x, (a, c)) -> r)
 -> Prob (->) r (a, c) (b, d))
-> (forall x. ((x, (b, d)) -> r) -> (x, (a, c)) -> r)
-> Prob (->) r (a, c) (b, d)
forall a b. (a -> b) -> a -> b
$ \(x, (b, d)) -> r
k ->
  let kf :: ((x, d), b) -> r
kf ((x
ctx, d
d), b
b) = (x, (b, d)) -> r
k (x
ctx, (b
b, d
d))
      fa :: ((x, d), a) -> r
fa = (((x, d), b) -> r) -> ((x, d), a) -> r
forall x. ((x, b) -> r) -> (x, a) -> r
f ((x, d), b) -> r
kf
      kg :: ((x, a), d) -> r
kg ((x
ctx, a
a), d
d) = ((x, d), a) -> r
fa ((x
ctx, d
d), a
a)
      gb :: ((x, a), c) -> r
gb = (((x, a), d) -> r) -> ((x, a), c) -> r
forall x. ((x, d) -> r) -> (x, c) -> r
g ((x, a), d) -> r
kg
   in \(x
ctx, (a
a, c
c)) -> ((x, a), c) -> r
gb ((x
ctx, a
a), c
c)
{-# INLINE parGF #-}

-- ---------------------------------------------------------------------------
-- Traced Either (explicit, not a canonical instance)
-- ---------------------------------------------------------------------------

-- | Least-fixpoint trace over the 'Either' tensor.
--
-- This is the denotationally correct definition: @Right@ values feed back into
-- the body, @Left@ values escape.  For genuinely cyclic bodies and strict
-- numeric scalars (e.g. @r = Double@) it diverges — the geometric series
-- exists but strict @(+)@ never reaches it.  Use 'traceEN' for a computable
-- approximation, or switch to a scalar whose lattice structure supplies the
-- fixpoint (e.g. @r = Bool@, where @(||)@ short-circuits) or to an effectful
-- base arrow where sampling terminates almost surely.
--
-- We do not provide a @Traced Either (Prob (->) r)@ instance because the
-- canonical trace is only available on a fragment; 'traceE' and 'traceEN' are
-- exported as explicit choices.
traceE ::
  Prob (->) r (Either a s) (Either b s) ->
  Prob (->) r a b
traceE :: forall r a s b.
Prob (->) r (Either a s) (Either b s) -> Prob (->) r a b
traceE (Prob forall x. ((x, Either b s) -> r) -> (x, Either a s) -> r
f) = (forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b)
-> (forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b
forall a b. (a -> b) -> a -> b
$ \(x, b) -> r
k ->
  let step :: (x, Either b s) -> r
step (x
x, Left b
b) = (x, b) -> r
k (x
x, b
b)
      step (x
x, Right s
s) = ((x, Either b s) -> r) -> (x, Either a s) -> r
forall x. ((x, Either b s) -> r) -> (x, Either a s) -> r
f (x, Either b s) -> r
step (x
x, s -> Either a s
forall a b. b -> Either a b
Right s
s)
   in \(x
x, a
a) -> ((x, Either b s) -> r) -> (x, Either a s) -> r
forall x. ((x, Either b s) -> r) -> (x, Either a s) -> r
f (x, Either b s) -> r
step (x
x, a -> Either a s
forall a b. a -> Either a b
Left a
a)
{-# INLINE traceE #-}

-- | Fuel-bounded variant of 'traceE'.  After the fuel is exhausted, re-entries
-- contribute the supplied @zero@ value.  This converges to the least fixpoint
-- with error proportional to the probability of not having terminated by the
-- fuel limit.
traceEN ::
  r ->
  Int ->
  Prob (->) r (Either a s) (Either b s) ->
  Prob (->) r a b
traceEN :: forall r a s b.
r
-> Int -> Prob (->) r (Either a s) (Either b s) -> Prob (->) r a b
traceEN r
zero Int
n0 (Prob forall x. ((x, Either b s) -> r) -> (x, Either a s) -> r
f) = (forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b)
-> (forall x. ((x, b) -> r) -> (x, a) -> r) -> Prob (->) r a b
forall a b. (a -> b) -> a -> b
$ \(x, b) -> r
k ->
  let step :: Int -> (x, Either b s) -> r
step Int
_ (x
x, Left b
b) = (x, b) -> r
k (x
x, b
b)
      step Int
n (x
x, Right s
s)
        | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = r
zero
        | Bool
otherwise = ((x, Either b s) -> r) -> (x, Either a s) -> r
forall x. ((x, Either b s) -> r) -> (x, Either a s) -> r
f (Int -> (x, Either b s) -> r
step (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) (x
x, s -> Either a s
forall a b. b -> Either a b
Right s
s)
   in \(x
x, a
a) -> ((x, Either b s) -> r) -> (x, Either a s) -> r
forall x. ((x, Either b s) -> r) -> (x, Either a s) -> r
f (Int -> (x, Either b s) -> r
step Int
n0) (x
x, a -> Either a s
forall a b. a -> Either a b
Left a
a)
{-# INLINE traceEN #-}