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

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

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')
    ]