Free.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Props.Free where

import Control.Applicative (Alternative (..))
import Data.Foldable (for_)
import Data.Kind (Type)
import Data.Type.Equality ((:~:) (..))
import Data.Type.Nat (Nat2)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.Falsify (testProperty)
import Prelude hiding (Monoid, curry, fst, id, mempty, snd, (**), (.))
import Prelude qualified as P

import Proarrow.Category.Instance.FinRel (FINREL (..))
import Proarrow.Category.Instance.Free (FREE (..), Free (..), Lower, retract, widen)
import Proarrow.Category.Instance.Opposite (OPPOSITE (..))
import Proarrow.Category.Instance.Unit (Unit (..))
import Proarrow.Category.Monoidal (Monoidal, MonoidalProfunctor (..), SymMonoidal, UnitF, withOb2, type (**!))
import Proarrow.Category.Monoidal.Cartesian (Cartesian, prodToTensor, tensorToProd, termToUnit, unitToTerm)
import Proarrow.Category.Monoidal.Closed (Closed, apply, curry, withObExp, type (-->))
import Proarrow.Category.Monoidal.CompactClosed (CompactClosed)
import Proarrow.Category.Monoidal.Dialogue (Dialogue, DualF)
import Proarrow.Category.Monoidal.Distributive (Distributive)
import Proarrow.Category.Monoidal.IsoMix (IsoMix)
import Proarrow.Category.Monoidal.StarAutonomous (StarAutonomous)
import Proarrow.Category.Sheaf (Cover (..), Leg (..), Sheaf (..), Summands, Sums, legArrow)
import Proarrow.Colimit.BinaryCoproduct (HasBinaryCoproducts (..), type (+))
import Proarrow.Colimit.Initial (HasInitialObject (..), InitF)
import Proarrow.Core (CAT, CategoryOf (..), Promonad (..), lmap, obj, type (+->))
import Proarrow.Functor (FunctorForRep (..), type (@))
import Proarrow.Limit.BinaryProduct (HasBinaryProducts (..), type (*!))
import Proarrow.Limit.Terminal (HasTerminalObject (..), TermF)
import Proarrow.Monoid (Comonoid (..), Monoid (..), Supplies)
import Proarrow.Profunctor.Instance.Initial (InitialProfunctor)
import Proarrow.Profunctor.Instance.Yoneda (Yo (..))
import Proarrow.Profunctor.Representable (Rep (..))

import Proarrow.Testing
  ( GenTotal (..)
  , MkSomeList (..)
  , Some (..)
  , Testable (..)
  , TestableProfunctor
  , TestableType (..)
  , TestingEqShow (..)
  , check
  , expect
  , genNamed
  , genSomeDef
  , oneOfTotal
  , testEq
  )
import Proarrow.Testing.Laws
import Props.Hask ()

type FREECS =
  '[ HasInitialObject
   , HasTerminalObject
   , HasBinaryProducts
   , HasBinaryCoproducts
   , Monoidal
   , SymMonoidal
   , Closed
   , Distributive
   , Dialogue
   , StarAutonomous
   , IsoMix
   , CompactClosed
   , Supplies Monoid
   , Supplies Comonoid
   ]
type FREEKIND = FREE FREECS (InitialProfunctor :: CAT ())

-- | The free category has no generating morphisms to interpret (@'InitialProfunctor'@ is
-- uninhabited), so any object at all works as the interpretation of @'()@. A small finite set
-- gives 'retract' below plenty to sample from. It is in 'FINREL', not 'Type', since
-- 'StarAutonomous'\/'CompactClosed' need a target that has dual objects.
data family Interp :: () +-> FINREL

instance FunctorForRep Interp where
  type Interp @ '() = FR Nat2
  fmap Unit = obj @(FR Nat2)

-- | The object @a@ interpreted in 'FINREL', by folding its structure through 'Interp'.
type LowerT a = Lower (Rep Interp) a

test :: TestTree
test =
  testGroup
    "Free"
    [ testCategory @FREEKIND
    , testTerminalObject @FREEKIND
    , testInitialObject @FREEKIND
    , testBinaryProducts @FREEKIND (\ @a @b r -> withObProd @FINREL @(LowerT a) @(LowerT b) r)
    , testBinaryCoproducts @FREEKIND (\ @a @b r -> withObCoprod @FINREL @(LowerT a) @(LowerT b) r)
    , testClosed @FREEKIND
        (\ @a @b r -> withOb2 @FINREL @(LowerT a) @(LowerT b) r)
        (\ @a @b r -> withObExp @FINREL @(LowerT a) @(LowerT b) r)
    , testMonoidal @FREEKIND (\ @a @b r -> withOb2 @FINREL @(LowerT a) @(LowerT b) r)
    , testSymMonoidal @FREEKIND (\ @a @b r -> withOb2 @FINREL @(LowerT a) @(LowerT b) r)
    , testDistributive @FREEKIND
        (\ @a @b r -> withOb2 @FINREL @(LowerT a) @(LowerT b) r)
        (\ @a @b r -> withObCoprod @FINREL @(LowerT a) @(LowerT b) r)
    , -- 'testStarAutonomous' isn't wired in here. Its naturality checks need e.g. an arbitrary
      -- @a ** b ~> Dual c@ for independently-drawn a,b,c, but in a free category that hom-set is
      -- empty for most palette triples (no unitor/associator-driven bridge connects a plain
      -- tensor shape to an unrelated dualized one). 'genTerm' can't conjure a morphism that
      -- doesn't exist, so every sample gets discarded. 'testCompactClosed' avoids this, since none
      -- of its checks need to generate a random Dual-involving morphism, only compose the fixed
      -- ones 'CompactClosed' already provides.
      testCompactClosed @FREEKIND
        (\ @a @b r -> withOb2 @FINREL @(LowerT a) @(LowerT b) r)
        (\ @a @b r -> withObExp @FINREL @(LowerT a) @(LowerT b) r)
        (\r -> r)
    , testIsoMix @FREEKIND (\ @a @b r -> withOb2 @FINREL @(LowerT a) @(LowerT b) r) (\r -> r)
    , testHypergraph @FREEKIND (\ @a @b r -> withOb2 @FINREL @(LowerT a) @(LowerT b) r)
    , sheafTests
    , testProperty "cartesian coercions interpret to identities" P.$ do
        let roundTrip = retract @CARTCS @(Rep InterpT) (tensorToProd @(EMB '()) @(EMB '()) . prodToTensor @(EMB '()) @(EMB '()))
            unitTrip = retract @CARTCS @(Rep InterpT) (unitToTerm . termToUnit)
        check "cartesian coercions" (roundTrip (True, False) P.== (True, False) P.&& unitTrip () P.== ())
    , testProperty "retract . widen = retract" P.$ do
        let l = retract @NARROWCS @(Rep Interp) narrowTerm
            r = retract @FREECS @(Rep Interp) (widen @FREECS narrowTerm)
        expect "retract . widen" r l
    ]

-- * The cartesian coercions

-- | A free category with the 'Cartesian' marker, interpreted into @Type@ (where the coercions
-- are identities), with the one generator object standing for 'Bool'.
type CARTCS = '[Cartesian, HasTerminalObject, HasBinaryProducts, Monoidal]

data family InterpT :: () +-> Type

instance FunctorForRep InterpT where
  type InterpT @ '() = Bool
  fmap Unit = obj @Bool

-- * Widening

type NARROWCS = '[HasInitialObject, HasTerminalObject, HasBinaryProducts]

-- | A term using all three structures of the narrow list, for the widening test above.
narrowTerm
  :: Free
       ((InitF *! TermF) :: FREE NARROWCS (InitialProfunctor :: CAT ()))
       ((TermF *! InitF) *! TermF)
narrowTerm = (terminate &&& (initiate @_ @InitF . fst @_ @InitF @TermF)) &&& snd @_ @InitF @TermF

-- | A singleton witnessing the shape of an object expression, so 'genTerm' can pattern-match
-- on source and target shapes directly instead of needing a type class per shape.
data SFree (a :: FREEKIND) where
  SInit :: SFree InitF
  STerm :: SFree TermF
  SProd :: (Ob a, Ob b) => SFree a -> SFree b -> SFree (a *! b)
  SSum :: (Ob a, Ob b) => SFree a -> SFree b -> SFree (a + b)
  SUnit :: SFree UnitF
  STen :: (Ob a, Ob b) => SFree a -> SFree b -> SFree (a **! b)
  SExp :: (Ob a, Ob b) => SFree a -> SFree b -> SFree (a --> b)
  SDual :: (Ob a) => SFree a -> SFree (DualF a)

class (Ob a) => KnownFree (a :: FREEKIND) where
  theFree :: SFree a
instance KnownFree InitF where
  theFree = SInit
instance KnownFree TermF where
  theFree = STerm
instance (KnownFree a, KnownFree b) => KnownFree (a *! b) where
  theFree = SProd theFree theFree
instance (KnownFree a, KnownFree b) => KnownFree (a + b) where
  theFree = SSum theFree theFree
instance KnownFree UnitF where
  theFree = SUnit
instance (KnownFree a, KnownFree b) => KnownFree (a **! b) where
  theFree = STen theFree theFree
instance (KnownFree a, KnownFree b) => KnownFree (a --> b) where
  theFree = SExp theFree theFree
instance (KnownFree a) => KnownFree (DualF a) where
  theFree = SDual theFree

-- | Decides whether two object shapes are the same, structurally. 'genTerm' uses this to
-- check whether a generator branch's source\/target lines up with the shape it wants
-- to produce, without needing runtime type reflection.
eqSFree :: SFree a -> SFree b -> Maybe (a :~: b)
eqSFree SInit SInit = Just Refl
eqSFree STerm STerm = Just Refl
eqSFree (SProd a1 a2) (SProd b1 b2) = case (eqSFree a1 b1, eqSFree a2 b2) of
  (Just Refl, Just Refl) -> Just Refl
  _ -> Nothing
eqSFree (SSum a1 a2) (SSum b1 b2) = case (eqSFree a1 b1, eqSFree a2 b2) of
  (Just Refl, Just Refl) -> Just Refl
  _ -> Nothing
eqSFree SUnit SUnit = Just Refl
eqSFree (STen a1 a2) (STen b1 b2) = case (eqSFree a1 b1, eqSFree a2 b2) of
  (Just Refl, Just Refl) -> Just Refl
  _ -> Nothing
eqSFree (SExp a1 a2) (SExp b1 b2) = case (eqSFree a1 b1, eqSFree a2 b2) of
  (Just Refl, Just Refl) -> Just Refl
  _ -> Nothing
eqSFree (SDual a1) (SDual b1) = case eqSFree a1 b1 of
  Just Refl -> Just Refl
  Nothing -> Nothing
eqSFree _ _ = Nothing

-- | Render an object shape for test failure output.
showSFree :: SFree a -> String
showSFree SInit = "InitF"
showSFree STerm = "TermF"
showSFree (SProd a b) = "(" ++ showSFree a ++ " *! " ++ showSFree b ++ ")"
showSFree (SSum a b) = "(" ++ showSFree a ++ " + " ++ showSFree b ++ ")"
showSFree SUnit = "UnitF"
showSFree (STen a b) = "(" ++ showSFree a ++ " **! " ++ showSFree b ++ ")"
showSFree (SExp a b) = "(" ++ showSFree a ++ " --> " ++ showSFree b ++ ")"
showSFree (SDual a) = "(Dual " ++ showSFree a ++ ")"

-- | The finite palette of shapes 'Testable' picks 'Some' objects from as the /endpoints/ of a
-- generated term. @composeB@ routes through 'Intermediates'.
--
-- Every shape is built from 'UnitF', and none from 'TermF' or 'InitF'. Terms are compared by
-- interpretation into 'FINREL', where the terminal and initial objects are both the empty set, so
-- at any object built from them alone every hom-set has one element and every comparison is
-- vacuous. The unit interprets to a one-element set, and the shapes over it do not collapse.
type Palette = '[UnitF, UnitF *! UnitF, UnitF + UnitF, UnitF **! UnitF, UnitF --> UnitF]

-- | The shapes @composeB@ routes intermediates through. It is 'Palette' plus 'TermF' and 'InitF'.
-- As /endpoints/ those two are useless (every hom-set at them is a singleton in 'FINREL', so a
-- comparison there cannot fail), but as /waypoints/ they are not. Going through 'TermF' builds
-- @'counit' '.' 'terminate' :: 'UnitF' '~>' 'UnitF'@, which interprets to the empty relation and is
-- the only non-identity endomorphism of 'UnitF' the generator can reach. Without it every law
-- comparison landing at @'UnitF' '~>' 'UnitF'@ is a single fixed instance.
type Intermediates = TermF ': InitF ': Palette

intermediates :: [Some FREEKIND]
intermediates = mkSomeList @FREEKIND @Intermediates

-- | Generate a random term between two (given) object shapes. Most branches recurse
-- structurally on a strictly smaller sub-shape of the source or target, so they always
-- terminate on their own; @composeB@ is the exception (it can reach into an unrelated object
-- via composition), so it's the one branch bounded by @fuel@, which decreases on every
-- recursive call and cuts it off at zero.
genTerm :: forall a b. (Ob a, Ob b) => Int -> SFree a -> SFree b -> GenTotal (Free a b)
genTerm fuel sa sb =
  oneOfTotal [idB, initiateB, terminateB, unitB, counitB, fstSndB, applyB, recB]
  where
    recB
      | fuel <= 0 = empty
      | otherwise = oneOfTotal [prodB, tensorB, sumSrcB, sumTgtB, curryB, composeB]
    idB = case eqSFree sa sb of
      Just Refl -> pure id
      Nothing -> empty
    initiateB = case sa of
      SInit -> pure initiate
      _ -> empty
    terminateB = case sb of
      STerm -> pure terminate
      _ -> empty
    -- The unit is neither initial nor terminal, but every object is a monoid and a comonoid here,
    -- so there is a canonical arrow from it and one to it all the same.
    unitB = case sa of
      SUnit -> pure mempty
      _ -> empty
    counitB = case sb of
      SUnit -> pure counit
      _ -> empty
    fstSndB = case sa of
      SProd sa1 sa2 ->
        oneOfTotal
          [ case eqSFree sb sa1 of Just Refl -> pure fst; Nothing -> empty
          , case eqSFree sb sa2 of Just Refl -> pure snd; Nothing -> empty
          ]
      _ -> empty
    prodB = case sb of
      SProd b1 b2 -> (&&&) <$> genTerm (fuel - 1) sa b1 <*> genTerm (fuel - 1) sa b2
      _ -> empty
    tensorB = case (sa, sb) of
      (STen a1 a2, STen b1 b2) -> (**) <$> genTerm (fuel - 1) a1 b1 <*> genTerm (fuel - 1) a2 b2
      _ -> empty
    sumSrcB = case sa of
      SSum a1 a2 -> (|||) <$> genTerm (fuel - 1) a1 sb <*> genTerm (fuel - 1) a2 sb
      _ -> empty
    sumTgtB = case sb of
      SSum b1 b2 -> oneOfTotal [(lft .) <$> genTerm (fuel - 1) sa b1, (rgt .) <$> genTerm (fuel - 1) sa b2]
      _ -> empty
    -- a ~ (b --> c) **! b, with c ~ the target.
    applyB = case sa of
      STen sl sr -> case sl of
        SExp sea1 sea2 -> case (eqSFree sr sea1, eqSFree sb sea2) of
          (Just Refl, Just Refl) -> pure apply
          _ -> empty
        _ -> empty
      _ -> empty
    curryB = case sb of
      SExp b1 b2 -> curry <$> genTerm (fuel - 1) (STen sa b1) b2
      _ -> empty
    -- Route through every palette shape as a possible intermediate object. Without this the
    -- generator could never compose two otherwise-unrelated terms.
    composeB =
      oneOfTotal
        [ (.) <$> genTerm (fuel - 1) (theFree @mid) sb <*> genTerm (fuel - 1) sa (theFree @mid)
        | Some @mid <- intermediates
        ]

-- | Bridges straight to 'FINREL'\'s 'Ob' instead of its 'TestOb'. 'Testable FINREL' leaves
-- 'TestOb' at its class default ('type TestOb a = Ob a'), and an unrestated default associated
-- type equation doesn't get unfolded through an abstract type variable the way an explicit
-- instance override (like 'CategoryOf FINREL'\'s own 'Ob' equation) does.
instance Testable FREEKIND where
  type TestOb a = (KnownFree a, Ob (LowerT a))
  obFromTestOb r = r
  showOb @a = showSFree (theFree @a)
  genSome = genSomeDef @Palette

-- | Two terms are equal iff they denote the same relation once interpreted into 'FINREL' via
-- 'retract', decided by 'FinRel'\'s own 'Eq'. Structural equality on 'Free' terms would be too
-- strict for testing categorical laws: e.g. @'terminate' . f@ and @'terminate'@ are built from
-- different 'Free' constructors even though uniqueness of the terminal object makes them denote
-- the same morphism.
instance (TestOb a, TestOb b) => TestingEqShow (Free (a :: FREEKIND) b) where
  eqP l r = pure (retract @FREECS @(Rep Interp) l == retract @FREECS @(Rep Interp) r)

instance (TestOb a, TestOb b) => TestableType (Free (a :: FREEKIND) b) where
  gen = genTerm 3 (theFree @a) (theFree @b)
instance TestableProfunctor (Free :: CAT FREEKIND)

-- | The cover of the booleans by their two points. 'BySummands' is polymorphic in the summands, so
-- the cover it is used at has to be pinned. The summands are 'UnitF' and not @TermF@ for the reason
-- 'Palette' gives.
boolCover :: Cover Sums FREEKIND (UnitF + UnitF) (Summands (UnitF :: FREEKIND) UnitF)
boolCover = BySummands

-- | The sum coverage on the free category, at the cover of the booleans by their two points: the
-- representables glue by @'|||'@. So restriction along either injection gives the branch back, and
-- an element is the gluing of its restrictions.
sheafTests :: TestTree
sheafTests =
  testGroup
    "Sums"
    [ testGluesBackAt @Sums @(Yo (UnitF + UnitF) (OP '())) "BySummands, sum representable" boolCover
    , -- the two-sided representable's own profunctor laws. These exercise 'Yo'\'s action on the
      -- covariant component, which the presheaf cases above leave at the identity
      testProfunctor @(Yo (UnitF + UnitF) (OP Bool) :: Type +-> FREEKIND)
    , testGluesBackAt
        @Sums
        @(Yo (UnitF + UnitF) (OP Bool) :: Type +-> FREEKIND)
        "BySummands, representable over Hask"
        boolCover
    , -- The gluing keeps one covariant component where the family supplies two, so it is only well
      -- The gluing keeps one covariant component of the two the family supplies, which is well
      -- defined because the two agree on the overlap. After restricting along the overlap the
      -- contravariant components are always equal ('InitF' is initial), so the covariant ones
      -- decide. The un-restricted pair below is the other way round, so both halves of 'eqP' on
      -- 'Yo' are exercised.
      testProperty "the overlap decides the covariant component" do
        let overlap = initiate @FREEKIND @UnitF
            at
              :: Free (UnitF :: FREEKIND) (UnitF + UnitF) -> (Bool -> Bool) -> Yo (UnitF + UnitF) (OP Bool) (UnitF :: FREEKIND) Bool
            at inj h = Yo inj h
        for_ [(h, h') | h <- [P.id, P.not], h' <- [P.id, P.not]] \(h, h') -> do
          agree <- eqP h h'
          same <- eqP (lmap overlap (at lft h)) (lmap overlap (at rgt h'))
          expect "restricted to the overlap: equal exactly when the covariant halves agree" agree same
          apart <- eqP (at lft h) (at rgt h)
          expect "un-restricted: the differing contravariant halves separate them" False apart
    , -- The commuting conversion, in sheaf vocabulary: a map of sheaves carries a gluing to the
      -- gluing of the mapped family. At the representable, gluing is @'|||'@ and the map is
      -- post-composition, so this is @h '.' (t '|||' e) = (h '.' t) '|||' (h '.' e)@, the
      -- equation that makes @f (if b then x else y)@ and @if b then f x else f y@ the same
      -- program. It follows from restriction and uniqueness together with @h@\'s naturality, so it
      -- is not a new law but a demonstration.
      testProperty "a map of sheaves commutes with the gluing" do
        t <- genNamed @(Free (UnitF :: FREEKIND) (UnitF + UnitF)) "t"
        e <- genNamed @(Free (UnitF :: FREEKIND) (UnitF + UnitF)) "e"
        -- @h@ has to land somewhere it can be injective. Every term @(UnitF + UnitF) ~> UnitF@ the
        -- generator can build collapses the two summands, and then @h . t = h . e@ for almost any
        -- branches and the equation holds for the wrong reason.
        h <- genNamed @(Free ((UnitF :: FREEKIND) + UnitF) (UnitF + UnitF)) "h"
        let after
              :: Yo ((UnitF :: FREEKIND) + UnitF) (OP '()) z '()
              -> Yo ((UnitF :: FREEKIND) + UnitF) (OP '()) z '()
            after (Yo f g) = Yo (h . f) g
            fam
              :: forall z
               . Leg Sums FREEKIND (UnitF + UnitF) (Summands (UnitF :: FREEKIND) UnitF) z
              -> Yo ((UnitF :: FREEKIND) + UnitF) (OP '()) z '()
            fam AtLeft = Yo t Unit
            fam AtRight = Yo e Unit
        testEq
          "commuting conversion"
          "h . glue [t, e]"
          (after (glue @Sums boolCover fam))
          "glue [h . t, h . e]"
          (glue @Sums boolCover (\g -> after (fam g)))
    , testProperty "restriction at BySummands" do
        t <- genNamed @(Free UnitF (UnitF + UnitF)) "t"
        e <- genNamed @(Free UnitF (UnitF + UnitF)) "e"
        let ite = glue @Sums @(Yo (UnitF + UnitF) (OP '())) boolCover \case
              AtLeft -> Yo t Unit
              AtRight -> Yo e Unit
        testEq "then" "lmap lft (glue [t, e])" (lmap (legArrow AtLeft) ite) "t" (Yo t Unit)
        testEq "else" "lmap rgt (glue [t, e])" (lmap (legArrow AtRight) ite) "e" (Yo e Unit)
    ]