{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Test.Cardano.Data.Arbitrary (genOSet) where

import Data.Map.NonEmpty as NEM
import Data.Map.Strict as Map
import Data.Maybe (fromJust)
import Data.OMap.Strict qualified as OMap
import Data.OSet.Strict qualified as OSet
import Data.Set qualified as Set
import Data.Set.NonEmpty (NonEmptySet)
import Data.Set.NonEmpty qualified as NES
import Test.Cardano.Ledger.Binary.Arbitrary ()
import Test.QuickCheck

instance (Arbitrary a, Ord a) => Arbitrary (OSet.OSet a) where
  arbitrary :: Gen (OSet a)
arbitrary = Gen a -> Gen (OSet a)
forall a. Ord a => Gen a -> Gen (OSet a)
genOSet Gen a
forall a. Arbitrary a => Gen a
arbitrary

genOSet :: Ord a => Gen a -> Gen (OSet.OSet a)
genOSet :: forall a. Ord a => Gen a -> Gen (OSet a)
genOSet = ([a] -> OSet a) -> Gen [a] -> Gen (OSet a)
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [a] -> OSet a
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> OSet a
OSet.fromFoldable (Gen [a] -> Gen (OSet a))
-> (Gen a -> Gen [a]) -> Gen a -> Gen (OSet a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Gen a -> Gen [a]
forall a. Gen a -> Gen [a]
listOf

instance (Arbitrary v, OMap.HasOKey k v, Arbitrary k) => Arbitrary (OMap.OMap k v) where
  arbitrary :: Gen (OMap k v)
arbitrary = forall (f :: * -> *) k v.
(Foldable f, HasOKey k v) =>
f v -> OMap k v
OMap.fromFoldable @[] ([v] -> OMap k v) -> Gen [v] -> Gen (OMap k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen [v]
forall a. Arbitrary a => Gen a
arbitrary

instance (Arbitrary a, Ord a) => Arbitrary (NonEmptySet a) where
  arbitrary :: Gen (NonEmptySet a)
arbitrary = do
    el <- Gen a
forall a. Arbitrary a => Gen a
arbitrary
    fromJust . NES.fromSet . Set.insert el <$> arbitrary

  shrink :: NonEmptySet a -> [NonEmptySet a]
shrink NonEmptySet a
nes =
    [ Maybe (NonEmptySet a) -> NonEmptySet a
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (NonEmptySet a) -> NonEmptySet a)
-> Maybe (NonEmptySet a) -> NonEmptySet a
forall a b. (a -> b) -> a -> b
$ Set a -> Maybe (NonEmptySet a)
forall a. Set a -> Maybe (NonEmptySet a)
NES.fromSet Set a
xs'
    | Set a
xs' <- Set a -> [Set a]
forall a. Arbitrary a => a -> [a]
shrink (Set a -> [Set a]) -> Set a -> [Set a]
forall a b. (a -> b) -> a -> b
$ NonEmptySet a -> Set a
forall a. NonEmptySet a -> Set a
NES.toSet NonEmptySet a
nes
    , Bool -> Bool
not (Set a -> Bool
forall a. Set a -> Bool
Set.null Set a
xs')
    ]

instance (Arbitrary k, Arbitrary v, Ord k) => Arbitrary (NonEmptyMap k v) where
  arbitrary :: Gen (NonEmptyMap k v)
arbitrary = do
    k <- Gen k
forall a. Arbitrary a => Gen a
arbitrary
    v <- arbitrary
    fromJust . NEM.fromMap . Map.insert k v <$> arbitrary

  shrink :: NonEmptyMap k v -> [NonEmptyMap k v]
shrink NonEmptyMap k v
nem =
    [ Maybe (NonEmptyMap k v) -> NonEmptyMap k v
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (NonEmptyMap k v) -> NonEmptyMap k v)
-> Maybe (NonEmptyMap k v) -> NonEmptyMap k v
forall a b. (a -> b) -> a -> b
$ Map k v -> Maybe (NonEmptyMap k v)
forall k v. Map k v -> Maybe (NonEmptyMap k v)
NEM.fromMap Map k v
xs'
    | Map k v
xs' <- Map k v -> [Map k v]
forall a. Arbitrary a => a -> [a]
shrink (Map k v -> [Map k v]) -> Map k v -> [Map k v]
forall a b. (a -> b) -> a -> b
$ NonEmptyMap k v -> Map k v
forall k v. NonEmptyMap k v -> Map k v
NEM.toMap NonEmptyMap k v
nem
    , Bool -> Bool
not (Map k v -> Bool
forall k a. Map k a -> Bool
Map.null Map k v
xs')
    ]