-- |
-- Module: Covenant.Test
-- Copyright: (C) MLabs 2025
-- License: Apache 2.0
-- Maintainer: koz@mlabs.city, sean@mlabs.city
--
-- Utilities designed to help test Covenant itself.
--
-- @since 1.0.0
module Covenant.Test
  ( Concrete (Concrete),
  )
where

import Control.Applicative ((<|>))
import Covenant.Index (count0)
import Covenant.Type
  ( AbstractTy,
    BuiltinFlatT
      ( BLS12_381_G1_ElementT,
        BLS12_381_G2_ElementT,
        BLS12_381_MlResultT,
        BoolT,
        ByteStringT,
        IntegerT,
        StringT,
        UnitT
      ),
    CompT (Comp0, CompN),
    CompTBody (ArgsAndResult),
    ValT (Abstraction, BuiltinFlat, ThunkT),
  )
import Data.Coerce (coerce)
import Data.Vector qualified as Vector
import Test.QuickCheck
  ( Arbitrary (arbitrary, shrink),
    Gen,
    elements,
    liftArbitrary,
    oneof,
    sized,
  )
import Test.QuickCheck.Instances.Vector ()

-- | Wrapper for 'ValT' to provide an 'Arbitrary' instance to generate only
-- value types without any type variables.
--
-- @since 1.0.0
newtype Concrete = Concrete (ValT AbstractTy)
  deriving
    ( -- | @since 1.0.0
      Concrete -> Concrete -> Bool
(Concrete -> Concrete -> Bool)
-> (Concrete -> Concrete -> Bool) -> Eq Concrete
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Concrete -> Concrete -> Bool
== :: Concrete -> Concrete -> Bool
$c/= :: Concrete -> Concrete -> Bool
/= :: Concrete -> Concrete -> Bool
Eq
    )
    via (ValT AbstractTy)
  deriving stock
    ( -- | @since 1.0.0
      Int -> Concrete -> ShowS
[Concrete] -> ShowS
Concrete -> String
(Int -> Concrete -> ShowS)
-> (Concrete -> String) -> ([Concrete] -> ShowS) -> Show Concrete
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Concrete -> ShowS
showsPrec :: Int -> Concrete -> ShowS
$cshow :: Concrete -> String
show :: Concrete -> String
$cshowList :: [Concrete] -> ShowS
showList :: [Concrete] -> ShowS
Show
    )

-- | @since 1.0.0
instance Arbitrary Concrete where
  {-# INLINEABLE arbitrary #-}
  arbitrary :: Gen Concrete
arbitrary = ValT AbstractTy -> Concrete
Concrete (ValT AbstractTy -> Concrete)
-> Gen (ValT AbstractTy) -> Gen Concrete
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> Gen (ValT AbstractTy)) -> Gen (ValT AbstractTy)
forall a. (Int -> Gen a) -> Gen a
sized Int -> Gen (ValT AbstractTy)
go
    where
      go :: Int -> Gen (ValT AbstractTy)
      go :: Int -> Gen (ValT AbstractTy)
go Int
size
        | Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 =
            BuiltinFlatT -> ValT AbstractTy
forall a. BuiltinFlatT -> ValT a
BuiltinFlat
              (BuiltinFlatT -> ValT AbstractTy)
-> Gen BuiltinFlatT -> Gen (ValT AbstractTy)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> [BuiltinFlatT] -> Gen BuiltinFlatT
forall a. HasCallStack => [a] -> Gen a
elements
                [ BuiltinFlatT
UnitT,
                  BuiltinFlatT
BoolT,
                  BuiltinFlatT
IntegerT,
                  BuiltinFlatT
StringT,
                  BuiltinFlatT
ByteStringT,
                  BuiltinFlatT
BLS12_381_G1_ElementT,
                  BuiltinFlatT
BLS12_381_G2_ElementT,
                  BuiltinFlatT
BLS12_381_MlResultT
                ]
        | Bool
otherwise =
            [Gen (ValT AbstractTy)] -> Gen (ValT AbstractTy)
forall a. HasCallStack => [Gen a] -> Gen a
oneof
              [ ValT AbstractTy -> Gen (ValT AbstractTy)
forall a. a -> Gen a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ValT AbstractTy -> Gen (ValT AbstractTy))
-> (BuiltinFlatT -> ValT AbstractTy)
-> BuiltinFlatT
-> Gen (ValT AbstractTy)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BuiltinFlatT -> ValT AbstractTy
forall a. BuiltinFlatT -> ValT a
BuiltinFlat (BuiltinFlatT -> Gen (ValT AbstractTy))
-> BuiltinFlatT -> Gen (ValT AbstractTy)
forall a b. (a -> b) -> a -> b
$ BuiltinFlatT
UnitT,
                ValT AbstractTy -> Gen (ValT AbstractTy)
forall a. a -> Gen a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ValT AbstractTy -> Gen (ValT AbstractTy))
-> (BuiltinFlatT -> ValT AbstractTy)
-> BuiltinFlatT
-> Gen (ValT AbstractTy)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BuiltinFlatT -> ValT AbstractTy
forall a. BuiltinFlatT -> ValT a
BuiltinFlat (BuiltinFlatT -> Gen (ValT AbstractTy))
-> BuiltinFlatT -> Gen (ValT AbstractTy)
forall a b. (a -> b) -> a -> b
$ BuiltinFlatT
BoolT,
                ValT AbstractTy -> Gen (ValT AbstractTy)
forall a. a -> Gen a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ValT AbstractTy -> Gen (ValT AbstractTy))
-> (BuiltinFlatT -> ValT AbstractTy)
-> BuiltinFlatT
-> Gen (ValT AbstractTy)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BuiltinFlatT -> ValT AbstractTy
forall a. BuiltinFlatT -> ValT a
BuiltinFlat (BuiltinFlatT -> Gen (ValT AbstractTy))
-> BuiltinFlatT -> Gen (ValT AbstractTy)
forall a b. (a -> b) -> a -> b
$ BuiltinFlatT
IntegerT,
                ValT AbstractTy -> Gen (ValT AbstractTy)
forall a. a -> Gen a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ValT AbstractTy -> Gen (ValT AbstractTy))
-> (BuiltinFlatT -> ValT AbstractTy)
-> BuiltinFlatT
-> Gen (ValT AbstractTy)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BuiltinFlatT -> ValT AbstractTy
forall a. BuiltinFlatT -> ValT a
BuiltinFlat (BuiltinFlatT -> Gen (ValT AbstractTy))
-> BuiltinFlatT -> Gen (ValT AbstractTy)
forall a b. (a -> b) -> a -> b
$ BuiltinFlatT
StringT,
                ValT AbstractTy -> Gen (ValT AbstractTy)
forall a. a -> Gen a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ValT AbstractTy -> Gen (ValT AbstractTy))
-> (BuiltinFlatT -> ValT AbstractTy)
-> BuiltinFlatT
-> Gen (ValT AbstractTy)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BuiltinFlatT -> ValT AbstractTy
forall a. BuiltinFlatT -> ValT a
BuiltinFlat (BuiltinFlatT -> Gen (ValT AbstractTy))
-> BuiltinFlatT -> Gen (ValT AbstractTy)
forall a b. (a -> b) -> a -> b
$ BuiltinFlatT
ByteStringT,
                ValT AbstractTy -> Gen (ValT AbstractTy)
forall a. a -> Gen a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ValT AbstractTy -> Gen (ValT AbstractTy))
-> (BuiltinFlatT -> ValT AbstractTy)
-> BuiltinFlatT
-> Gen (ValT AbstractTy)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BuiltinFlatT -> ValT AbstractTy
forall a. BuiltinFlatT -> ValT a
BuiltinFlat (BuiltinFlatT -> Gen (ValT AbstractTy))
-> BuiltinFlatT -> Gen (ValT AbstractTy)
forall a b. (a -> b) -> a -> b
$ BuiltinFlatT
BLS12_381_G1_ElementT,
                ValT AbstractTy -> Gen (ValT AbstractTy)
forall a. a -> Gen a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ValT AbstractTy -> Gen (ValT AbstractTy))
-> (BuiltinFlatT -> ValT AbstractTy)
-> BuiltinFlatT
-> Gen (ValT AbstractTy)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BuiltinFlatT -> ValT AbstractTy
forall a. BuiltinFlatT -> ValT a
BuiltinFlat (BuiltinFlatT -> Gen (ValT AbstractTy))
-> BuiltinFlatT -> Gen (ValT AbstractTy)
forall a b. (a -> b) -> a -> b
$ BuiltinFlatT
BLS12_381_G2_ElementT,
                ValT AbstractTy -> Gen (ValT AbstractTy)
forall a. a -> Gen a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ValT AbstractTy -> Gen (ValT AbstractTy))
-> (BuiltinFlatT -> ValT AbstractTy)
-> BuiltinFlatT
-> Gen (ValT AbstractTy)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BuiltinFlatT -> ValT AbstractTy
forall a. BuiltinFlatT -> ValT a
BuiltinFlat (BuiltinFlatT -> Gen (ValT AbstractTy))
-> BuiltinFlatT -> Gen (ValT AbstractTy)
forall a b. (a -> b) -> a -> b
$ BuiltinFlatT
BLS12_381_MlResultT,
                CompT AbstractTy -> ValT AbstractTy
forall a. CompT a -> ValT a
ThunkT (CompT AbstractTy -> ValT AbstractTy)
-> (CompTBody AbstractTy -> CompT AbstractTy)
-> CompTBody AbstractTy
-> ValT AbstractTy
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompTBody AbstractTy -> CompT AbstractTy
forall a. CompTBody a -> CompT a
Comp0 (CompTBody AbstractTy -> ValT AbstractTy)
-> Gen (CompTBody AbstractTy) -> Gen (ValT AbstractTy)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (Vector (ValT AbstractTy) -> ValT AbstractTy -> CompTBody AbstractTy
forall a. Vector (ValT a) -> ValT a -> CompTBody a
ArgsAndResult (Vector (ValT AbstractTy)
 -> ValT AbstractTy -> CompTBody AbstractTy)
-> Gen (Vector (ValT AbstractTy))
-> Gen (ValT AbstractTy -> CompTBody AbstractTy)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (ValT AbstractTy) -> Gen (Vector (ValT AbstractTy))
forall a. Gen a -> Gen (Vector a)
forall (f :: Type -> Type) a. Arbitrary1 f => Gen a -> Gen (f a)
liftArbitrary (Int -> Gen (ValT AbstractTy)
go (Int
size Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
4)) Gen (ValT AbstractTy -> CompTBody AbstractTy)
-> Gen (ValT AbstractTy) -> Gen (CompTBody AbstractTy)
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Int -> Gen (ValT AbstractTy)
go (Int
size Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
4))
              ]
  {-# INLINEABLE shrink #-}
  shrink :: Concrete -> [Concrete]
shrink (Concrete ValT AbstractTy
v) =
    ValT AbstractTy -> Concrete
Concrete (ValT AbstractTy -> Concrete) -> [ValT AbstractTy] -> [Concrete]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> case ValT AbstractTy
v of
      -- impossible
      Abstraction AbstractTy
_ -> []
      ThunkT (CompN Count "tyvar"
_ (ArgsAndResult Vector (ValT AbstractTy)
args ValT AbstractTy
result)) ->
        CompT AbstractTy -> ValT AbstractTy
forall a. CompT a -> ValT a
ThunkT (CompT AbstractTy -> ValT AbstractTy)
-> (CompTBody AbstractTy -> CompT AbstractTy)
-> CompTBody AbstractTy
-> ValT AbstractTy
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Count "tyvar" -> CompTBody AbstractTy -> CompT AbstractTy
CompN Count "tyvar"
forall (ofWhat :: Symbol). Count ofWhat
count0 (CompTBody AbstractTy -> ValT AbstractTy)
-> [CompTBody AbstractTy] -> [ValT AbstractTy]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> do
          let argsList :: [ValT AbstractTy]
argsList = Vector (ValT AbstractTy) -> [ValT AbstractTy]
forall a. Vector a -> [a]
Vector.toList Vector (ValT AbstractTy)
args
          [ValT AbstractTy]
argsList' <- ([Concrete] -> [ValT AbstractTy])
-> [[Concrete]] -> [[ValT AbstractTy]]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap [Concrete] -> [ValT AbstractTy]
forall a b. Coercible a b => a -> b
coerce ([[Concrete]] -> [[ValT AbstractTy]])
-> ([ValT AbstractTy] -> [[Concrete]])
-> [ValT AbstractTy]
-> [[ValT AbstractTy]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Concrete] -> [[Concrete]]
forall a. Arbitrary a => a -> [a]
shrink ([Concrete] -> [[Concrete]])
-> ([ValT AbstractTy] -> [Concrete])
-> [ValT AbstractTy]
-> [[Concrete]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ValT AbstractTy -> Concrete) -> [ValT AbstractTy] -> [Concrete]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap ValT AbstractTy -> Concrete
Concrete ([ValT AbstractTy] -> [[ValT AbstractTy]])
-> [ValT AbstractTy] -> [[ValT AbstractTy]]
forall a b. (a -> b) -> a -> b
$ [ValT AbstractTy]
argsList
          ValT AbstractTy
result' <- (Concrete -> ValT AbstractTy) -> [Concrete] -> [ValT AbstractTy]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Concrete -> ValT AbstractTy
forall a b. Coercible a b => a -> b
coerce ([Concrete] -> [ValT AbstractTy])
-> (ValT AbstractTy -> [Concrete])
-> ValT AbstractTy
-> [ValT AbstractTy]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete -> [Concrete]
forall a. Arbitrary a => a -> [a]
shrink (Concrete -> [Concrete])
-> (ValT AbstractTy -> Concrete) -> ValT AbstractTy -> [Concrete]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ValT AbstractTy -> Concrete
Concrete (ValT AbstractTy -> [ValT AbstractTy])
-> ValT AbstractTy -> [ValT AbstractTy]
forall a b. (a -> b) -> a -> b
$ ValT AbstractTy
result
          let args' :: Vector (ValT AbstractTy)
args' = [ValT AbstractTy] -> Vector (ValT AbstractTy)
forall a. [a] -> Vector a
Vector.fromList [ValT AbstractTy]
argsList'
          CompTBody AbstractTy -> [CompTBody AbstractTy]
forall a. a -> [a]
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Vector (ValT AbstractTy) -> ValT AbstractTy -> CompTBody AbstractTy
forall a. Vector (ValT a) -> ValT a -> CompTBody a
ArgsAndResult Vector (ValT AbstractTy)
args' ValT AbstractTy
result) [CompTBody AbstractTy]
-> [CompTBody AbstractTy] -> [CompTBody AbstractTy]
forall a. [a] -> [a] -> [a]
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> CompTBody AbstractTy -> [CompTBody AbstractTy]
forall a. a -> [a]
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Vector (ValT AbstractTy) -> ValT AbstractTy -> CompTBody AbstractTy
forall a. Vector (ValT a) -> ValT a -> CompTBody a
ArgsAndResult Vector (ValT AbstractTy)
args ValT AbstractTy
result')
      -- Can't shrink this
      BuiltinFlat BuiltinFlatT
_ -> []