{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}

module Plutarch.Test.QuickCheck.Function (
  PFun (..),
  pattern PFn,
  applyPFun,
  plamTable,
  plamFinite,
) where

import Control.Arrow (Arrow ((&&&)))
import Data.Kind (Type)
import Data.List (intercalate, nubBy)
import Data.Universe (Finite (universeF))
import Plutarch (S, Term, plam, (#), (#$), type (:-->))
import Plutarch.Extra.Maybe (pfromMaybe)
import Plutarch.Lift (PLift, PUnsafeLiftDecl (PLifted), pconstant)
import Plutarch.Maybe (pfromJust)
import Plutarch.Prelude (
  PBuiltinList,
  PBuiltinPair,
  PEq,
  PIsListLike,
  PMaybe (PJust, PNothing),
  pcon,
  pfind,
  pfstBuiltin,
  phoistAcyclic,
  pmatch,
  psndBuiltin,
  (#==),
 )
import Plutarch.Test.QuickCheck.Internal (
  TestableTerm (TestableTerm),
  unTestableTerm,
 )
import Test.QuickCheck (
  Arbitrary (arbitrary, shrink),
  CoArbitrary (coarbitrary),
  Gen,
  sized,
  vectorOf,
 )

-- TODO: Generate function with more arguments
-- TODO: ^ would this be possible want script size is limited?
data PFun (a :: S -> Type) (b :: S -> Type) where
  PFun ::
    (PLift a, PLift b) =>
    [(PLifted a, PLifted b)] ->
    PLifted b ->
    (TestableTerm (a :--> b)) ->
    PFun a b

{-# COMPLETE PFn #-}
pattern PFn ::
  forall
    {a :: S -> Type}
    {b :: S -> Type}.
  (PUnsafeLiftDecl a, PUnsafeLiftDecl b) =>
  (forall (s :: S). Term s (a :--> b)) ->
  PFun a b
pattern $mPFn :: forall {r} {a :: S -> *} {b :: S -> *}.
(PUnsafeLiftDecl a, PUnsafeLiftDecl b) =>
PFun a b
-> ((forall (s :: S). Term s (a :--> b)) -> r) -> ((# #) -> r) -> r
PFn f <- (unTestableTerm . applyPFun -> f)

applyPFun ::
  forall {a :: S -> Type} {b :: S -> Type}.
  PFun a b ->
  TestableTerm (a :--> b)
applyPFun :: forall {a :: S -> *} {b :: S -> *}.
PFun a b -> TestableTerm (a :--> b)
applyPFun (PFun [(PLifted a, PLifted b)]
_ PLifted b
_ TestableTerm (a :--> b)
f) = TestableTerm (a :--> b)
f

mkPFun ::
  forall (a :: S -> Type) (b :: S -> Type).
  ( PLift a
  , PLift b
  , PEq a
  ) =>
  [(PLifted a, PLifted b)] ->
  PLifted b ->
  PFun a b
mkPFun :: forall (a :: S -> *) (b :: S -> *).
(PLift a, PLift b, PEq a) =>
[(PLifted a, PLifted b)] -> PLifted b -> PFun a b
mkPFun [(PLifted a, PLifted b)]
t PLifted b
d = forall (a :: S -> *) (b :: S -> *).
(PLift a, PLift b) =>
[(PLifted a, PLifted b)]
-> PLifted b -> TestableTerm (a :--> b) -> PFun a b
PFun [(PLifted a, PLifted b)]
t PLifted b
d forall a b. (a -> b) -> a -> b
$ forall (a :: S -> *). (forall (s :: S). Term s a) -> TestableTerm a
TestableTerm forall a b. (a -> b) -> a -> b
$ forall (a :: S -> *) (b :: S -> *) (s :: S).
(PLift a, PLift b, PEq a) =>
[(PLifted a, PLifted b)] -> PLifted b -> Term s (a :--> b)
plamTable [(PLifted a, PLifted b)]
t PLifted b
d

instance
  forall (a :: S -> Type) (b :: S -> Type).
  ( PLift a
  , PLift b
  , Arbitrary (PLifted a)
  , Arbitrary (PLifted b)
  , CoArbitrary (PLifted a)
  , Eq (PLifted a)
  , PEq a
  ) =>
  Arbitrary (PFun a b)
  where
  arbitrary :: Gen (PFun a b)
arbitrary = forall a. (Int -> Gen a) -> Gen a
sized forall a b. (a -> b) -> a -> b
$ \Int
r -> do
    [PLifted a]
xs <- forall a. Int -> Gen a -> Gen [a]
vectorOf Int
r (forall a. Arbitrary a => Gen a
arbitrary :: Gen (PLifted a))
    [PLifted b]
ys <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((forall a b. (a -> b) -> a -> b
$ (forall a. Arbitrary a => Gen a
arbitrary :: Gen (PLifted b))) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. CoArbitrary a => a -> Gen b -> Gen b
coarbitrary) [PLifted a]
xs
    let table :: [(PLifted a, PLifted b)]
table = forall a b. [a] -> [b] -> [(a, b)]
zip [PLifted a]
xs [PLifted b]
ys

    PLifted b
d <- forall a. Arbitrary a => Gen a
arbitrary :: Gen (PLifted b)
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (a :: S -> *) (b :: S -> *).
(PLift a, PLift b, PEq a) =>
[(PLifted a, PLifted b)] -> PLifted b -> PFun a b
mkPFun (forall a. (a -> a -> Bool) -> [a] -> [a]
nubBy (\(PLifted a, PLifted b)
x (PLifted a, PLifted b)
y -> forall a b. (a, b) -> a
fst (PLifted a, PLifted b)
x forall a. Eq a => a -> a -> Bool
== forall a b. (a, b) -> a
fst (PLifted a, PLifted b)
y) [(PLifted a, PLifted b)]
table) PLifted b
d

  shrink :: PFun a b -> [PFun a b]
shrink (PFun [(PLifted a, PLifted b)]
t PLifted b
d TestableTerm (a :--> b)
_) =
    [forall (a :: S -> *) (b :: S -> *).
(PLift a, PLift b, PEq a) =>
[(PLifted a, PLifted b)] -> PLifted b -> PFun a b
mkPFun [(PLifted a, PLifted b)]
t' PLifted b
d' | ([(PLifted a, PLifted b)]
t', PLifted b
d') <- forall a. Arbitrary a => a -> [a]
shrink ([(PLifted a, PLifted b)]
t, PLifted b
d)]

instance
  forall (a :: S -> Type) (b :: S -> Type).
  ( PLift a
  , PLift b
  , Show (PLifted a)
  , Show (PLifted b)
  ) =>
  Show (PFun a b)
  where
  show :: PFun a b -> String
show = forall (a :: S -> *) (b :: S -> *).
(Show (PLifted a), Show (PLifted b)) =>
PFun a b -> String
showPFun

showPFun ::
  forall (a :: S -> Type) (b :: S -> Type).
  ( Show (PLifted a)
  , Show (PLifted b)
  ) =>
  PFun a b ->
  String
showPFun :: forall (a :: S -> *) (b :: S -> *).
(Show (PLifted a), Show (PLifted b)) =>
PFun a b -> String
showPFun (PFun [(PLifted a, PLifted b)]
t PLifted b
d TestableTerm (a :--> b)
_) =
  String
"{\n"
    forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [[a]] -> [a]
intercalate
      String
", \n"
      ( [ forall a. Show a => a -> String
show PLifted a
x forall a. [a] -> [a] -> [a]
++ String
" -> " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show PLifted b
c
        | (PLifted a
x, PLifted b
c) <- [(PLifted a, PLifted b)]
t
        ]
          forall a. [a] -> [a] -> [a]
++ [String
"_ -> " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show PLifted b
d]
      )
    forall a. [a] -> [a] -> [a]
++ String
"\n}"

plamTable ::
  forall (a :: S -> Type) (b :: S -> Type) (s :: S).
  ( PLift a
  , PLift b
  , PEq a
  ) =>
  [(PLifted a, PLifted b)] ->
  PLifted b ->
  Term s (a :--> b)
plamTable :: forall (a :: S -> *) (b :: S -> *) (s :: S).
(PLift a, PLift b, PEq a) =>
[(PLifted a, PLifted b)] -> PLifted b -> Term s (a :--> b)
plamTable [(PLifted a, PLifted b)]
t PLifted b
d = forall a (b :: S -> *) (s :: S) (c :: S -> *).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
plam forall a b. (a -> b) -> a -> b
$ \Term s a
x -> forall (a :: S -> *) (s :: S). Term s (a :--> (PMaybe a :--> a))
pfromMaybe forall (s :: S) (a :: S -> *) (b :: S -> *).
HasCallStack =>
Term s (a :--> b) -> Term s a -> Term s b
# forall (p :: S -> *) (s :: S). PLift p => PLifted p -> Term s p
pconstant PLifted b
d forall (s :: S) (a :: S -> *) (b :: S -> *).
HasCallStack =>
Term s (a :--> b) -> Term s a -> Term s b
# (forall (a :: S -> *) (b :: S -> *) (s :: S)
       (list :: (S -> *) -> S -> *).
(PEq a, PIsListLike list (PBuiltinPair a b)) =>
Term s (a :--> (list (PBuiltinPair a b) :--> PMaybe b))
plookup forall (s :: S) (a :: S -> *) (b :: S -> *).
HasCallStack =>
Term s (a :--> b) -> Term s a -> Term s b
# Term s a
x forall (s :: S) (a :: S -> *) (b :: S -> *).
HasCallStack =>
Term s (a :--> b) -> Term s a -> Term s b
# forall (p :: S -> *) (s :: S). PLift p => PLifted p -> Term s p
pconstant [(PLifted a, PLifted b)]
t)

plamFinite ::
  forall (a :: S -> Type) (b :: S -> Type) (s :: S).
  ( Finite (PLifted a)
  , PLift a
  , PLift b
  , PEq a
  ) =>
  (PLifted a -> PLifted b) ->
  Term s (a :--> b)
plamFinite :: forall (a :: S -> *) (b :: S -> *) (s :: S).
(Finite (PLifted a), PLift a, PLift b, PEq a) =>
(PLifted a -> PLifted b) -> Term s (a :--> b)
plamFinite PLifted a -> PLifted b
f = forall a (b :: S -> *) (s :: S) (c :: S -> *).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
plam forall a b. (a -> b) -> a -> b
$ \Term s a
x -> forall (s :: S) (a :: S -> *). Term s (PMaybe a :--> a)
pfromJust forall (s :: S) (a :: S -> *) (b :: S -> *).
HasCallStack =>
Term s (a :--> b) -> Term s a -> Term s b
#$ forall (a :: S -> *) (b :: S -> *) (s :: S)
       (list :: (S -> *) -> S -> *).
(PEq a, PIsListLike list (PBuiltinPair a b)) =>
Term s (a :--> (list (PBuiltinPair a b) :--> PMaybe b))
plookup forall (s :: S) (a :: S -> *) (b :: S -> *).
HasCallStack =>
Term s (a :--> b) -> Term s a -> Term s b
# Term s a
x forall (s :: S) (a :: S -> *) (b :: S -> *).
HasCallStack =>
Term s (a :--> b) -> Term s a -> Term s b
# Term s (PBuiltinList (PBuiltinPair a b))
table
  where
    table :: Term s (PBuiltinList (PBuiltinPair a b))
    table :: Term s (PBuiltinList (PBuiltinPair a b))
table = forall (p :: S -> *) (s :: S). PLift p => PLifted p -> Term s p
pconstant forall a b. (a -> b) -> a -> b
$ (forall a. a -> a
id forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& PLifted a -> PLifted b
f) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Finite a => [a]
universeF

plookup ::
  forall (a :: S -> Type) (b :: S -> Type) (s :: S) list.
  (PEq a, PIsListLike list (PBuiltinPair a b)) =>
  Term s (a :--> list (PBuiltinPair a b) :--> PMaybe b)
plookup :: forall (a :: S -> *) (b :: S -> *) (s :: S)
       (list :: (S -> *) -> S -> *).
(PEq a, PIsListLike list (PBuiltinPair a b)) =>
Term s (a :--> (list (PBuiltinPair a b) :--> PMaybe b))
plookup =
  forall (a :: S -> *) (s :: S).
HasCallStack =>
ClosedTerm a -> Term s a
phoistAcyclic forall a b. (a -> b) -> a -> b
$
    forall a (b :: S -> *) (s :: S) (c :: S -> *).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
plam forall a b. (a -> b) -> a -> b
$ \Term s a
k Term s (list (PBuiltinPair a b))
xs ->
      forall (a :: S -> *) (s :: S) (b :: S -> *).
PlutusType a =>
Term s a -> (a s -> Term s b) -> Term s b
pmatch (forall (l :: (S -> *) -> S -> *) (a :: S -> *) (s :: S).
PIsListLike l a =>
Term s ((a :--> PBool) :--> (l a :--> PMaybe a))
pfind forall (s :: S) (a :: S -> *) (b :: S -> *).
HasCallStack =>
Term s (a :--> b) -> Term s a -> Term s b
# forall a (b :: S -> *) (s :: S) (c :: S -> *).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
plam (\Term s (PBuiltinPair a b)
p -> forall (s :: S) (a :: S -> *) (b :: S -> *).
Term s (PBuiltinPair a b :--> a)
pfstBuiltin forall (s :: S) (a :: S -> *) (b :: S -> *).
HasCallStack =>
Term s (a :--> b) -> Term s a -> Term s b
# Term s (PBuiltinPair a b)
p forall (t :: S -> *) (s :: S).
PEq t =>
Term s t -> Term s t -> Term s PBool
#== Term s a
k) forall (s :: S) (a :: S -> *) (b :: S -> *).
HasCallStack =>
Term s (a :--> b) -> Term s a -> Term s b
# Term s (list (PBuiltinPair a b))
xs) forall a b. (a -> b) -> a -> b
$ \case
        PMaybe (PBuiltinPair a b) s
PNothing -> forall (a :: S -> *) (s :: S). PlutusType a => a s -> Term s a
pcon forall (a :: S -> *) (s :: S). PMaybe a s
PNothing
        PJust Term s (PBuiltinPair a b)
p -> forall (a :: S -> *) (s :: S). PlutusType a => a s -> Term s a
pcon (forall (a :: S -> *) (s :: S). Term s a -> PMaybe a s
PJust (forall (s :: S) (a :: S -> *) (b :: S -> *).
Term s (PBuiltinPair a b :--> b)
psndBuiltin forall (s :: S) (a :: S -> *) (b :: S -> *).
HasCallStack =>
Term s (a :--> b) -> Term s a -> Term s b
# Term s (PBuiltinPair a b)
p))