diff options
Diffstat (limited to 'test/Gen.hs')
| -rw-r--r-- | test/Gen.hs | 29 |
1 files changed, 17 insertions, 12 deletions
diff --git a/test/Gen.hs b/test/Gen.hs index 50c671f..4f5fe96 100644 --- a/test/Gen.hs +++ b/test/Gen.hs @@ -4,7 +4,6 @@ {-# LANGUAGE NumericUnderscores #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeAbstractions #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeOperators #-} {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-} @@ -20,10 +19,10 @@ import Foreign import GHC.TypeLits import GHC.TypeNats qualified as TN -import Data.Array.Mixed.Permutation -import Data.Array.Mixed.Types import Data.Array.Nested +import Data.Array.Nested.Permutation import Data.Array.Nested.Ranked.Shape +import Data.Array.Nested.Types import Hedgehog import Hedgehog.Gen qualified as Gen @@ -59,9 +58,12 @@ shuffleShR = \sh -> go (length sh) (toList sh) sh (dim :$:) <$> go (nbag - 1) bag' sh genShR :: SNat n -> Gen (IShR n) -genShR sn = do +genShR = genShRwithTarget 100_000 + +genShRwithTarget :: Int -> SNat n -> Gen (IShR n) +genShRwithTarget targetMax sn = do let n = fromSNat' sn - targetSize <- Gen.int (Range.linear 0 100_000) + targetSize <- Gen.int (Range.linear 0 targetMax) let genDims :: SNat m -> Int -> Gen (IShR m) genDims SZ _ = return ZSR genDims (SS m) 0 = do @@ -76,9 +78,8 @@ genShR sn = do dims <- genDims m (if dim == 0 then 0 else tgt `div` dim) return (dim :$: dims) dims <- genDims sn targetSize - let dimsL = toList dims - maxdim = maximum dimsL - cap = binarySearch (`div` 2) 1 maxdim (\cap' -> product (min cap' <$> dimsL) <= targetSize) + let maxdim = maximum dims + cap = binarySearch (`div` 2) 1 maxdim (\cap' -> shrSize (min cap' <$> dims) <= targetSize) shuffleShR (min cap <$> dims) -- | Example: given 3 and 7, might return: @@ -93,10 +94,14 @@ genShR sn = do -- other dimensions might be zero. genReplicatedShR :: m <= n => SNat m -> SNat n -> Gen (IShR m, IShR n, IShR n) genReplicatedShR = \m n -> do - sh1 <- genShR m + let expectedSizeIncrease = round (repvalavg ^ (fromSNat' n - fromSNat' m)) + sh1 <- genShRwithTarget (1_000_000 `div` expectedSizeIncrease) m (sh2, sh3) <- injectOnes n sh1 sh1 return (sh1, sh2, sh3) where + repvalrange = (1::Int, 5) + repvalavg = let (lo, hi) = repvalrange in fromIntegral (lo + hi) / 2 :: Double + injectOnes :: m <= n => SNat n -> IShR m -> IShR m -> Gen (IShR n, IShR n) injectOnes n@SNat shOnes sh | m@SNat <- shrRank sh @@ -105,7 +110,7 @@ genReplicatedShR = \m n -> do EQI -> return (shOnes, sh) GTI -> do index <- Gen.int (Range.linear 0 (fromSNat' m)) - value <- Gen.int (Range.linear 1 5) + value <- Gen.int (uncurry Range.linear repvalrange) Refl <- return (lem n m) injectOnes n (inject index 1 shOnes) (inject index value sh) @@ -115,7 +120,7 @@ genReplicatedShR = \m n -> do inject :: Int -> Int -> IShR m -> IShR (m + 1) inject 0 v sh = v :$: sh inject i v (w :$: sh) = w :$: inject (i - 1) v sh - inject _ v ZSR = v :$: ZSR -- invalid input, but meh + inject _ _ ZSR = error "unreachable" genStorables :: forall a. Storable a => Range Int -> (Word64 -> a) -> GenT IO (VS.Vector a) genStorables rng f = do @@ -157,7 +162,7 @@ genPermR n = Gen.shuffle [0 .. n-1] genPerm :: Monad m => SNat n -> (forall p. (IsPermutation p, Rank p ~ n) => Perm p -> PropertyT m r) -> PropertyT m r genPerm n@SNat k = do list <- forAll $ genPermR (fromSNat' n) - permFromList list $ \perm -> do + permFromListCont list $ \perm -> do case permCheckPermutation perm $ case sameNat' (permRank perm) n of Just Refl -> Just (k perm) |