module Plutarch.Enum (
  -- * Type classes
  PCountable (..),
  PEnumerable (..),
) where

import Data.Kind (Type)
import Plutarch.Builtin.Bool (pif)
import Plutarch.Integer (PInteger)
import Plutarch.Internal.Eq ((#==))
import Plutarch.Internal.Ord (POrd)
import Plutarch.Internal.Other (pfix, pto)
import Plutarch.Internal.PLam (plam)
import Plutarch.Internal.Term (
  S,
  Term,
  phoistAcyclic,
  (#),
  (#$),
  (:-->),
 )
import Plutarch.Positive (PPositive)

{- | A notion of \'next\' value. More formally, instances of this type class are
discrete linear orders with no maximal element.

= Laws

1. @x /= psuccessor x@
2. @y < x@ @=@ @psuccessor y <= x@
3. @x < psuccessor y@ @=@ @x <= y@

If you define 'psuccessorN', you must also ensure the following hold; the
default implementation ensures this.

4. @psuccessorN 1@ @=@ @psuccessor@
5. @psuccessorN n . psuccessorN m@ @=@ @psuccessorN (n + m)@

Law 1 ensures no value is its own successor. Laws 2 and 3 ensure that there
are no \'gaps\': every value is \'reachable\' from any lower value by a
finite number of applications of 'successor'.

@since WIP
-}
class POrd a => PCountable (a :: S -> Type) where
  -- | @since WIP
  psuccessor :: forall (s :: S). Term s (a :--> a)

  -- | The default implementation of this function is inefficient: if at all
  -- possible, give instances an optimized version that doesn't require
  -- recursion.
  --
  -- @since WIP
  {-# INLINEABLE psuccessorN #-}
  psuccessorN :: forall (s :: S). Term s (PPositive :--> a :--> a)
  psuccessorN = (forall (s :: S). Term s (PPositive :--> (a :--> a)))
-> Term s (PPositive :--> (a :--> a))
forall (a :: S -> Type) (s :: S).
HasCallStack =>
ClosedTerm a -> Term s a
phoistAcyclic ((forall (s :: S). Term s (PPositive :--> (a :--> a)))
 -> Term s (PPositive :--> (a :--> a)))
-> (forall (s :: S). Term s (PPositive :--> (a :--> a)))
-> Term s (PPositive :--> (a :--> a))
forall a b. (a -> b) -> a -> b
$ (Term s PPositive -> Term s a -> Term s a)
-> Term s (PPositive :--> (a :--> a))
forall a (b :: S -> Type) (s :: S) (c :: S -> Type).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: S -> Type).
HasCallStack =>
(Term s c -> Term s a -> Term s a) -> Term s (c :--> (a :--> a))
plam ((Term s PPositive -> Term s a -> Term s a)
 -> Term s (PPositive :--> (a :--> a)))
-> (Term s PPositive -> Term s a -> Term s a)
-> Term s (PPositive :--> (a :--> a))
forall a b. (a -> b) -> a -> b
$ \Term s PPositive
n Term s a
x -> Term s PPositive -> Term s (a :--> (PPositive :--> a))
forall (s' :: S).
Term s' PPositive -> Term s' (a :--> (PPositive :--> a))
go Term s PPositive
n Term s (a :--> (PPositive :--> a))
-> Term s a -> Term s (PPositive :--> a)
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# (Term s (a :--> a)
forall (s :: S). Term s (a :--> a)
forall (a :: S -> Type) (s :: S). PCountable a => Term s (a :--> a)
psuccessor Term s (a :--> a) -> Term s a -> Term s a
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# Term s a
x) Term s (PPositive :--> a) -> Term s PPositive -> Term s a
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# Term s PPositive
1
    where
      go ::
        forall (s' :: S).
        Term s' PPositive ->
        Term s' (a :--> PPositive :--> a)
      go :: forall (s' :: S).
Term s' PPositive -> Term s' (a :--> (PPositive :--> a))
go Term s' PPositive
limit = Term
  s'
  (((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a)))
   :--> (a :--> (PPositive :--> a)))
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (((a :--> b) :--> (a :--> b)) :--> (a :--> b))
pfix Term
  s'
  (((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a)))
   :--> (a :--> (PPositive :--> a)))
-> Term
     s' ((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a)))
-> Term s' (a :--> (PPositive :--> a))
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
#$ (Term s' (a :--> (PPositive :--> a))
 -> Term s' a -> Term s' PPositive -> Term s' a)
-> Term
     s' ((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a)))
forall a (b :: S -> Type) (s :: S) (c :: S -> Type).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: S -> Type).
HasCallStack =>
(Term s' c -> Term s' a -> Term s' PPositive -> Term s' a)
-> Term s' (c :--> (a :--> (PPositive :--> a)))
plam ((Term s' (a :--> (PPositive :--> a))
  -> Term s' a -> Term s' PPositive -> Term s' a)
 -> Term
      s' ((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a))))
-> (Term s' (a :--> (PPositive :--> a))
    -> Term s' a -> Term s' PPositive -> Term s' a)
-> Term
     s' ((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a)))
forall a b. (a -> b) -> a -> b
$ \Term s' (a :--> (PPositive :--> a))
self Term s' a
acc Term s' PPositive
count ->
        Term s' PBool -> Term s' a -> Term s' a -> Term s' a
forall (a :: S -> Type) (s :: S).
Term s PBool -> Term s a -> Term s a -> Term s a
pif
          (Term s' PPositive
count Term s' PPositive -> Term s' PPositive -> Term s' PBool
forall (s :: S).
Term s PPositive -> Term s PPositive -> Term s PBool
forall (t :: S -> Type) (s :: S).
PEq t =>
Term s t -> Term s t -> Term s PBool
#== Term s' PPositive
limit)
          Term s' a
acc
          (Term s' (a :--> (PPositive :--> a))
self Term s' (a :--> (PPositive :--> a))
-> Term s' a -> Term s' (PPositive :--> a)
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# (Term s' (a :--> a)
forall (s :: S). Term s (a :--> a)
forall (a :: S -> Type) (s :: S). PCountable a => Term s (a :--> a)
psuccessor Term s' (a :--> a) -> Term s' a -> Term s' a
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# Term s' a
acc) Term s' (PPositive :--> a) -> Term s' PPositive -> Term s' a
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# (Term s' PPositive
count Term s' PPositive -> Term s' PPositive -> Term s' PPositive
forall a. Num a => a -> a -> a
+ Term s' PPositive
1))

-- | @since WIP
instance PCountable PInteger where
  {-# INLINEABLE psuccessor #-}
  psuccessor :: forall (s :: S). Term s (PInteger :--> PInteger)
psuccessor = (forall (s :: S). Term s (PInteger :--> PInteger))
-> Term s (PInteger :--> PInteger)
forall (a :: S -> Type) (s :: S).
HasCallStack =>
ClosedTerm a -> Term s a
phoistAcyclic ((forall (s :: S). Term s (PInteger :--> PInteger))
 -> Term s (PInteger :--> PInteger))
-> (forall (s :: S). Term s (PInteger :--> PInteger))
-> Term s (PInteger :--> PInteger)
forall a b. (a -> b) -> a -> b
$ (Term s PInteger -> Term s PInteger)
-> Term s (PInteger :--> PInteger)
forall a (b :: S -> Type) (s :: S) (c :: S -> Type).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: S -> Type).
HasCallStack =>
(Term s c -> Term s PInteger) -> Term s (c :--> PInteger)
plam (Term s PInteger -> Term s PInteger -> Term s PInteger
forall a. Num a => a -> a -> a
+ Term s PInteger
1)
  {-# INLINEABLE psuccessorN #-}
  psuccessorN :: forall (s :: S). Term s (PPositive :--> (PInteger :--> PInteger))
psuccessorN = (forall (s :: S). Term s (PPositive :--> (PInteger :--> PInteger)))
-> Term s (PPositive :--> (PInteger :--> PInteger))
forall (a :: S -> Type) (s :: S).
HasCallStack =>
ClosedTerm a -> Term s a
phoistAcyclic ((forall (s :: S).
  Term s (PPositive :--> (PInteger :--> PInteger)))
 -> Term s (PPositive :--> (PInteger :--> PInteger)))
-> (forall (s :: S).
    Term s (PPositive :--> (PInteger :--> PInteger)))
-> Term s (PPositive :--> (PInteger :--> PInteger))
forall a b. (a -> b) -> a -> b
$ (Term s PPositive -> Term s PInteger -> Term s PInteger)
-> Term s (PPositive :--> (PInteger :--> PInteger))
forall a (b :: S -> Type) (s :: S) (c :: S -> Type).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: S -> Type).
HasCallStack =>
(Term s c -> Term s PInteger -> Term s PInteger)
-> Term s (c :--> (PInteger :--> PInteger))
plam ((Term s PPositive -> Term s PInteger -> Term s PInteger)
 -> Term s (PPositive :--> (PInteger :--> PInteger)))
-> (Term s PPositive -> Term s PInteger -> Term s PInteger)
-> Term s (PPositive :--> (PInteger :--> PInteger))
forall a b. (a -> b) -> a -> b
$ \Term s PPositive
p Term s PInteger
i -> Term s PPositive -> Term s (PInner PPositive)
forall (s :: S) (a :: S -> Type). Term s a -> Term s (PInner a)
pto Term s PPositive
p Term s PInteger -> Term s PInteger -> Term s PInteger
forall a. Num a => a -> a -> a
+ Term s PInteger
i

-- | @since WIP
instance PCountable PPositive where
  {-# INLINEABLE psuccessor #-}
  psuccessor :: forall (s :: S). Term s (PPositive :--> PPositive)
psuccessor = (forall (s :: S). Term s (PPositive :--> PPositive))
-> Term s (PPositive :--> PPositive)
forall (a :: S -> Type) (s :: S).
HasCallStack =>
ClosedTerm a -> Term s a
phoistAcyclic ((forall (s :: S). Term s (PPositive :--> PPositive))
 -> Term s (PPositive :--> PPositive))
-> (forall (s :: S). Term s (PPositive :--> PPositive))
-> Term s (PPositive :--> PPositive)
forall a b. (a -> b) -> a -> b
$ (Term s PPositive -> Term s PPositive)
-> Term s (PPositive :--> PPositive)
forall a (b :: S -> Type) (s :: S) (c :: S -> Type).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: S -> Type).
HasCallStack =>
(Term s c -> Term s PPositive) -> Term s (c :--> PPositive)
plam (Term s PPositive -> Term s PPositive -> Term s PPositive
forall a. Num a => a -> a -> a
+ Term s PPositive
1)
  {-# INLINEABLE psuccessorN #-}
  psuccessorN :: forall (s :: S). Term s (PPositive :--> (PPositive :--> PPositive))
psuccessorN = (forall (s :: S).
 Term s (PPositive :--> (PPositive :--> PPositive)))
-> Term s (PPositive :--> (PPositive :--> PPositive))
forall (a :: S -> Type) (s :: S).
HasCallStack =>
ClosedTerm a -> Term s a
phoistAcyclic ((forall (s :: S).
  Term s (PPositive :--> (PPositive :--> PPositive)))
 -> Term s (PPositive :--> (PPositive :--> PPositive)))
-> (forall (s :: S).
    Term s (PPositive :--> (PPositive :--> PPositive)))
-> Term s (PPositive :--> (PPositive :--> PPositive))
forall a b. (a -> b) -> a -> b
$ (Term s PPositive -> Term s PPositive -> Term s PPositive)
-> Term s (PPositive :--> (PPositive :--> PPositive))
forall a (b :: S -> Type) (s :: S) (c :: S -> Type).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: S -> Type).
HasCallStack =>
(Term s c -> Term s PPositive -> Term s PPositive)
-> Term s (c :--> (PPositive :--> PPositive))
plam Term s PPositive -> Term s PPositive -> Term s PPositive
forall a. Num a => a -> a -> a
(+)

{- | Similar to 'PCountable', but has the ability to get a \'previous\' value as
well. More formally, instances of this type class are discrete linear orders
with no maximal or minimal element.

= Laws

1. @ppredecessor . psuccessor@ @=@ @psuccessor . ppredecessor@ @=@ @id@

If you define 'ppredecessorN', you must also ensure the following hold; the
default implementation ensures this.

2. @ppredecessorN 1@ @=@ @ppredecessor@
3. @ppredecessorN n . ppredecessorN m@ @=@ @ppredecessorN (n + m)@

From Law 1, we obtain the following theorem:

* @x /= predecessor x@

@since WIP
-}
class PCountable a => PEnumerable (a :: S -> Type) where
  -- | @since WIP
  ppredecessor :: forall (s :: S). Term s (a :--> a)

  -- | The default implementation of this function is inefficient: if at all
  -- possible, give instances an optimized version that doesn't require
  -- recursion.
  --
  -- @since WIP
  {-# INLINEABLE ppredecessorN #-}
  ppredecessorN :: forall (s :: S). Term s (PPositive :--> a :--> a)
  ppredecessorN = (forall (s :: S). Term s (PPositive :--> (a :--> a)))
-> Term s (PPositive :--> (a :--> a))
forall (a :: S -> Type) (s :: S).
HasCallStack =>
ClosedTerm a -> Term s a
phoistAcyclic ((forall (s :: S). Term s (PPositive :--> (a :--> a)))
 -> Term s (PPositive :--> (a :--> a)))
-> (forall (s :: S). Term s (PPositive :--> (a :--> a)))
-> Term s (PPositive :--> (a :--> a))
forall a b. (a -> b) -> a -> b
$ (Term s PPositive -> Term s a -> Term s a)
-> Term s (PPositive :--> (a :--> a))
forall a (b :: S -> Type) (s :: S) (c :: S -> Type).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: S -> Type).
HasCallStack =>
(Term s c -> Term s a -> Term s a) -> Term s (c :--> (a :--> a))
plam ((Term s PPositive -> Term s a -> Term s a)
 -> Term s (PPositive :--> (a :--> a)))
-> (Term s PPositive -> Term s a -> Term s a)
-> Term s (PPositive :--> (a :--> a))
forall a b. (a -> b) -> a -> b
$ \Term s PPositive
n Term s a
x -> Term s PPositive -> Term s (a :--> (PPositive :--> a))
forall (s' :: S).
Term s' PPositive -> Term s' (a :--> (PPositive :--> a))
go Term s PPositive
n Term s (a :--> (PPositive :--> a))
-> Term s a -> Term s (PPositive :--> a)
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# (Term s (a :--> a)
forall (s :: S). Term s (a :--> a)
forall (a :: S -> Type) (s :: S).
PEnumerable a =>
Term s (a :--> a)
ppredecessor Term s (a :--> a) -> Term s a -> Term s a
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# Term s a
x) Term s (PPositive :--> a) -> Term s PPositive -> Term s a
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# Term s PPositive
1
    where
      go ::
        forall (s' :: S).
        Term s' PPositive ->
        Term s' (a :--> PPositive :--> a)
      go :: forall (s' :: S).
Term s' PPositive -> Term s' (a :--> (PPositive :--> a))
go Term s' PPositive
limit = Term
  s'
  (((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a)))
   :--> (a :--> (PPositive :--> a)))
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (((a :--> b) :--> (a :--> b)) :--> (a :--> b))
pfix Term
  s'
  (((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a)))
   :--> (a :--> (PPositive :--> a)))
-> Term
     s' ((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a)))
-> Term s' (a :--> (PPositive :--> a))
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
#$ (Term s' (a :--> (PPositive :--> a))
 -> Term s' a -> Term s' PPositive -> Term s' a)
-> Term
     s' ((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a)))
forall a (b :: S -> Type) (s :: S) (c :: S -> Type).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: S -> Type).
HasCallStack =>
(Term s' c -> Term s' a -> Term s' PPositive -> Term s' a)
-> Term s' (c :--> (a :--> (PPositive :--> a)))
plam ((Term s' (a :--> (PPositive :--> a))
  -> Term s' a -> Term s' PPositive -> Term s' a)
 -> Term
      s' ((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a))))
-> (Term s' (a :--> (PPositive :--> a))
    -> Term s' a -> Term s' PPositive -> Term s' a)
-> Term
     s' ((a :--> (PPositive :--> a)) :--> (a :--> (PPositive :--> a)))
forall a b. (a -> b) -> a -> b
$ \Term s' (a :--> (PPositive :--> a))
self Term s' a
acc Term s' PPositive
count ->
        Term s' PBool -> Term s' a -> Term s' a -> Term s' a
forall (a :: S -> Type) (s :: S).
Term s PBool -> Term s a -> Term s a -> Term s a
pif
          (Term s' PPositive
count Term s' PPositive -> Term s' PPositive -> Term s' PBool
forall (s :: S).
Term s PPositive -> Term s PPositive -> Term s PBool
forall (t :: S -> Type) (s :: S).
PEq t =>
Term s t -> Term s t -> Term s PBool
#== Term s' PPositive
limit)
          Term s' a
acc
          (Term s' (a :--> (PPositive :--> a))
self Term s' (a :--> (PPositive :--> a))
-> Term s' a -> Term s' (PPositive :--> a)
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# (Term s' (a :--> a)
forall (s :: S). Term s (a :--> a)
forall (a :: S -> Type) (s :: S).
PEnumerable a =>
Term s (a :--> a)
ppredecessor Term s' (a :--> a) -> Term s' a -> Term s' a
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# Term s' a
acc) Term s' (PPositive :--> a) -> Term s' PPositive -> Term s' a
forall (s :: S) (a :: S -> Type) (b :: S -> Type).
Term s (a :--> b) -> Term s a -> Term s b
# (Term s' PPositive
count Term s' PPositive -> Term s' PPositive -> Term s' PPositive
forall a. Num a => a -> a -> a
+ Term s' PPositive
1))

-- | @since WIP
instance PEnumerable PInteger where
  {-# INLINEABLE ppredecessor #-}
  ppredecessor :: forall (s :: S). Term s (PInteger :--> PInteger)
ppredecessor = (forall (s :: S). Term s (PInteger :--> PInteger))
-> Term s (PInteger :--> PInteger)
forall (a :: S -> Type) (s :: S).
HasCallStack =>
ClosedTerm a -> Term s a
phoistAcyclic ((forall (s :: S). Term s (PInteger :--> PInteger))
 -> Term s (PInteger :--> PInteger))
-> (forall (s :: S). Term s (PInteger :--> PInteger))
-> Term s (PInteger :--> PInteger)
forall a b. (a -> b) -> a -> b
$ (Term s PInteger -> Term s PInteger)
-> Term s (PInteger :--> PInteger)
forall a (b :: S -> Type) (s :: S) (c :: S -> Type).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: S -> Type).
HasCallStack =>
(Term s c -> Term s PInteger) -> Term s (c :--> PInteger)
plam (Term s PInteger -> Term s PInteger -> Term s PInteger
forall a. Num a => a -> a -> a
- Term s PInteger
1)
  {-# INLINEABLE ppredecessorN #-}
  ppredecessorN :: forall (s :: S). Term s (PPositive :--> (PInteger :--> PInteger))
ppredecessorN = (forall (s :: S). Term s (PPositive :--> (PInteger :--> PInteger)))
-> Term s (PPositive :--> (PInteger :--> PInteger))
forall (a :: S -> Type) (s :: S).
HasCallStack =>
ClosedTerm a -> Term s a
phoistAcyclic ((forall (s :: S).
  Term s (PPositive :--> (PInteger :--> PInteger)))
 -> Term s (PPositive :--> (PInteger :--> PInteger)))
-> (forall (s :: S).
    Term s (PPositive :--> (PInteger :--> PInteger)))
-> Term s (PPositive :--> (PInteger :--> PInteger))
forall a b. (a -> b) -> a -> b
$ (Term s PPositive -> Term s PInteger -> Term s PInteger)
-> Term s (PPositive :--> (PInteger :--> PInteger))
forall a (b :: S -> Type) (s :: S) (c :: S -> Type).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: S -> Type).
HasCallStack =>
(Term s c -> Term s PInteger -> Term s PInteger)
-> Term s (c :--> (PInteger :--> PInteger))
plam ((Term s PPositive -> Term s PInteger -> Term s PInteger)
 -> Term s (PPositive :--> (PInteger :--> PInteger)))
-> (Term s PPositive -> Term s PInteger -> Term s PInteger)
-> Term s (PPositive :--> (PInteger :--> PInteger))
forall a b. (a -> b) -> a -> b
$ \Term s PPositive
p Term s PInteger
i -> Term s PInteger
i Term s PInteger -> Term s PInteger -> Term s PInteger
forall a. Num a => a -> a -> a
- Term s PPositive -> Term s (PInner PPositive)
forall (s :: S) (a :: S -> Type). Term s a -> Term s (PInner a)
pto Term s PPositive
p