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

Random generation of vectors.
-}

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

{-# LANGUAGE GADTs #-}
{-# LANGUAGE MagicHash #-}

module Clash.Hedgehog.Sized.Vector
  ( genVec
  , genNonEmptyVec
  , SomeVec(..)
  , genSomeVec
  ) where

import Prelude hiding (repeat)

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

import Clash.Promoted.Nat
import Clash.Sized.Vector

-- | Generate a potentially empty vector, where each element is produced
-- using the supplied generator. For a non-empty vector, see 'genNonEmptyVec'.
--
genVec :: (MonadGen m, KnownNat n) => m a -> m (Vec n a)
genVec :: m a -> m (Vec n a)
genVec m a
genElem = (m a -> m a) -> Vec n (m a) -> m (Vec n a)
forall a (f :: Type -> Type) b (n :: Nat).
Applicative f =>
(a -> f b) -> Vec n a -> f (Vec n b)
traverse# m a -> m a
forall a. a -> a
id (m a -> Vec n (m a)
forall (n :: Nat) a. KnownNat n => a -> Vec n a
repeat m a
genElem)

-- | Generate a non-empty vector, where each element is produced using the
-- supplied generator. For a potentially empty vector, see 'genVec'.
--
genNonEmptyVec :: (MonadGen m, KnownNat n, 1 <= n) => m a -> m (Vec n a)
genNonEmptyVec :: m a -> m (Vec n a)
genNonEmptyVec = m a -> m (Vec n a)
forall (m :: Type -> Type) (n :: Nat) a.
(MonadGen m, KnownNat n) =>
m a -> m (Vec n a)
genVec

data SomeVec atLeast a where
  SomeVec :: SNat n -> Vec (atLeast + n) a -> SomeVec atLeast a

genSomeVec
  :: (MonadGen m, KnownNat atLeast)
  => Range Natural
  -> m a
  -> m (SomeVec atLeast a)
genSomeVec :: Range Natural -> m a -> m (SomeVec atLeast a)
genSomeVec 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 -> Vec (atLeast + n) a -> SomeVec atLeast a
forall (n :: Nat) (atLeast :: Nat) a.
SNat n -> Vec (atLeast + n) a -> SomeVec atLeast a
SomeVec (Proxy n -> SNat n
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> SNat n
snatProxy Proxy n
proxy) (Vec (atLeast + n) a -> SomeVec atLeast a)
-> m (Vec (atLeast + n) a) -> m (SomeVec atLeast a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> m a -> m (Vec (atLeast + n) a)
forall (m :: Type -> Type) (n :: Nat) a.
(MonadGen m, KnownNat n) =>
m a -> m (Vec n a)
genVec m a
genElem