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

Random generation of SNat.
-}

{-# LANGUAGE GADTs #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskellQuotes #-}

module Clash.Hedgehog.Promoted.Nat where

import Clash.Prelude
import Hedgehog (Gen)
import qualified Data.String.Interpolate as I
import qualified GHC.TypeNats as TypeNats
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range

-- | Like @SomeSNat@, but with bounds.
data SomeBoundedSNat lower upperInclusive where
  SomeBoundedSNat
    :: forall n lower upperInclusive
     . (lower <= n, n <= upperInclusive)
    => SNat n
    -> SomeBoundedSNat lower upperInclusive

instance Show (SomeBoundedSNat lower upperInclusive) where
  show :: SomeBoundedSNat lower upperInclusive -> String
show (SomeBoundedSNat sn :: SNat n
sn@SNat n
SNat) = Name -> String
forall a. Show a => a -> String
show ''SomeBoundedSNat String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> SNat n -> String
forall a. Show a => a -> String
show SNat n
sn

{- | Generate a v'SomeBoundedSNat' between the given bounds (inclusive).
Uses 'Hedgehog.Range.linear' and shrinks to @lower@.

For example, to bring into scope the constraints @(KnownNat n, 1 <= n, n <= 8)@:

> SomeBoundedSNat (SNat :: SNat n) <- forAll (genSomeBoundedSNat @1 @8)
-}
genSomeBoundedSNat
  :: (KnownNat lower, KnownNat upperInclusive, lower <= upperInclusive)
  => Gen (SomeBoundedSNat lower upperInclusive)
genSomeBoundedSNat :: forall (lower :: Nat) (upperInclusive :: Nat).
(KnownNat lower, KnownNat upperInclusive,
 lower <= upperInclusive) =>
Gen (SomeBoundedSNat lower upperInclusive)
genSomeBoundedSNat = (Nat -> Nat -> Gen Nat)
-> Gen (SomeBoundedSNat lower upperInclusive)
forall (lower :: Nat) (upperInclusive :: Nat).
(KnownNat lower, KnownNat upperInclusive,
 lower <= upperInclusive) =>
(Nat -> Nat -> Gen Nat)
-> Gen (SomeBoundedSNat lower upperInclusive)
genSomeBoundedSNat# (\Nat
l -> Range Nat -> Gen Nat
forall (m :: Type -> Type) a.
(MonadGen m, Integral a) =>
Range a -> m a
Gen.integral (Range Nat -> Gen Nat) -> (Nat -> Range Nat) -> Nat -> Gen Nat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Nat -> Nat -> Range Nat
forall a. Integral a => a -> a -> Range a
Range.linear Nat
l)

-- | Like 'genSomeBoundedSNat' but does /not/ shrink.
genSomeBoundedSNat_
  :: (KnownNat lower, KnownNat upperInclusive, lower <= upperInclusive)
  => Gen (SomeBoundedSNat lower upperInclusive)
genSomeBoundedSNat_ :: forall (lower :: Nat) (upperInclusive :: Nat).
(KnownNat lower, KnownNat upperInclusive,
 lower <= upperInclusive) =>
Gen (SomeBoundedSNat lower upperInclusive)
genSomeBoundedSNat_ = (Nat -> Nat -> Gen Nat)
-> Gen (SomeBoundedSNat lower upperInclusive)
forall (lower :: Nat) (upperInclusive :: Nat).
(KnownNat lower, KnownNat upperInclusive,
 lower <= upperInclusive) =>
(Nat -> Nat -> Gen Nat)
-> Gen (SomeBoundedSNat lower upperInclusive)
genSomeBoundedSNat# (\Nat
l -> Range Nat -> Gen Nat
forall (m :: Type -> Type) a.
(MonadGen m, Integral a) =>
Range a -> m a
Gen.integral_ (Range Nat -> Gen Nat) -> (Nat -> Range Nat) -> Nat -> Gen Nat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Nat -> Nat -> Range Nat
forall a. Integral a => a -> a -> Range a
Range.linear Nat
l)

{- | Generate a v'SomeBoundedSNat' between the given bounds (inclusive).
To do so, uses a given term level generator.

__NB__: Make sure the given generator respects the bounds that it is passed at
the term level.
-}
genSomeBoundedSNat#
  :: forall lower upperInclusive
   . (KnownNat lower, KnownNat upperInclusive, lower <= upperInclusive)
  => (Natural -> Natural -> Gen Natural)
  -> Gen (SomeBoundedSNat lower upperInclusive)
genSomeBoundedSNat# :: forall (lower :: Nat) (upperInclusive :: Nat).
(KnownNat lower, KnownNat upperInclusive,
 lower <= upperInclusive) =>
(Nat -> Nat -> Gen Nat)
-> Gen (SomeBoundedSNat lower upperInclusive)
genSomeBoundedSNat# Nat -> Nat -> Gen Nat
gen = do
  Nat
n <- Nat -> Nat -> Gen Nat
gen (forall (n :: Nat). KnownNat n => Nat
natToNatural @lower) (forall (n :: Nat). KnownNat n => Nat
natToNatural @upperInclusive)
  case Nat -> SomeNat
TypeNats.someNatVal Nat
n of
    SomeNat Proxy n
p ->
      case (forall (n :: Nat). KnownNat n => SNat n
SNat @lower SNat lower -> SNat n -> SNatLE lower n
forall (a :: Nat) (b :: Nat). SNat a -> SNat b -> SNatLE a b
`compareSNat` SNat n
sn, SNat n
sn SNat n -> SNat upperInclusive -> SNatLE n upperInclusive
forall (a :: Nat) (b :: Nat). SNat a -> SNat b -> SNatLE a b
`compareSNat` forall (n :: Nat). KnownNat n => SNat n
SNat @upperInclusive) of
        (SNatLE lower n
SNatLE, SNatLE n upperInclusive
SNatLE) -> SomeBoundedSNat lower upperInclusive
-> Gen (SomeBoundedSNat lower upperInclusive)
forall a. a -> GenT Identity a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (SNat n -> SomeBoundedSNat lower upperInclusive
forall (n :: Nat) (lower :: Nat) (upperInclusive :: Nat).
(lower <= n, n <= upperInclusive) =>
SNat n -> SomeBoundedSNat lower upperInclusive
SomeBoundedSNat SNat n
sn)
        (SNatLE lower n, SNatLE n upperInclusive)
_ ->
          String -> Gen (SomeBoundedSNat lower upperInclusive)
forall a. HasCallStack => String -> a
error
            [I.__i|
            genSomeBoundedSNat: internal error: generated number out of bounds

            lower bound:             #{natToNatural @lower}
            upper bound (inclusive): #{natToNatural @upperInclusive}
            generated number:        #{n}
            |]
       where
        sn :: SNat n
sn = Proxy n -> SNat n
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> SNat n
snatProxy Proxy n
p