{-|
Copyright   : (C) 2021, QBayLogic B.V.
License     : BSD2 (see the file LICENSE)
Maintainer  : QBayLogic B.V. <devops@qbaylogic.com>

Random generation of RTree.
-}

{-# OPTIONS_GHC -fplugin=GHC.TypeLits.KnownNat.Solver #-}

{-# LANGUAGE GADTs #-}

module Clash.Hedgehog.Sized.RTree
  ( genRTree
  , genNonEmptyRTree
  , SomeRTree(..)
  , genSomeRTree
  ) where

import GHC.Natural (Natural)
import GHC.TypeNats
import Hedgehog (MonadGen, Range)
import qualified Hedgehog.Gen as Gen

import Clash.Promoted.Nat
import Clash.Sized.RTree

genRTree :: (MonadGen m, KnownNat n) => m a -> m (RTree n a)
genRTree :: m a -> m (RTree n a)
genRTree m a
genElem = RTree n (m a) -> m (RTree n a)
forall (t :: Type -> Type) (f :: Type -> Type) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (m a -> RTree n (m a)
forall (d :: Nat) a. KnownNat d => a -> RTree d a
trepeat m a
genElem)

genNonEmptyRTree :: (MonadGen m, KnownNat n, 1 <= n) => m a -> m (RTree n a)
genNonEmptyRTree :: m a -> m (RTree n a)
genNonEmptyRTree = m a -> m (RTree n a)
forall (m :: Type -> Type) (n :: Nat) a.
(MonadGen m, KnownNat n) =>
m a -> m (RTree n a)
genRTree

data SomeRTree atLeast a where
  SomeRTree :: SNat n -> RTree (atLeast + n) a -> SomeRTree atLeast a

genSomeRTree
  :: (MonadGen m, KnownNat atLeast)
  => Range Natural
  -> m a
  -> m (SomeRTree atLeast a)
genSomeRTree :: Range Natural -> m a -> m (SomeRTree atLeast a)
genSomeRTree Range Natural
rangeElems m a
genElem = do
  Natural
numExtra <- Range Natural -> m Natural
forall (m :: Type -> Type) a.
(MonadGen m, Integral a) =>
Range a -> m a
Gen.integral Range Natural
rangeElems

  case Natural -> SomeNat
someNatVal Natural
numExtra of
    SomeNat Proxy n
proxy -> SNat n -> RTree (atLeast + n) a -> SomeRTree atLeast a
forall (n :: Nat) (atLeast :: Nat) a.
SNat n -> RTree (atLeast + n) a -> SomeRTree atLeast a
SomeRTree (Proxy n -> SNat n
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> SNat n
snatProxy Proxy n
proxy) (RTree (atLeast + n) a -> SomeRTree atLeast a)
-> m (RTree (atLeast + n) a) -> m (SomeRTree atLeast a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> m a -> m (RTree (atLeast + n) a)
forall (m :: Type -> Type) (n :: Nat) a.
(MonadGen m, KnownNat n) =>
m a -> m (RTree n a)
genRTree m a
genElem