Make ShR a newtype over ShX

1 files changed, 18 insertions, 19 deletions

diff --git a/test/Tests/C.hs b/test/Tests/C.hs
index e26c3dd..4d372e9 100644
--- a/test/Tests/C.hs
+++ b/test/Tests/C.hs
@@ -15,7 +15,6 @@ module Tests.C where
 
 import Control.Monad
 import Data.Array.RankedS qualified as OR
-import Data.Foldable (toList)
 import Data.Functor.Const
 import Data.Type.Equality
 import Foreign
@@ -50,10 +49,10 @@ gen_red_nonempty f = property $ genRank $ \outrank@(SNat @n) -> do
   -- Test nonempty _results_. The first dimension of the input is allowed to be 0, because then OR.rerank doesn't fail yet.
   let inrank = SNat @(n + 1)
   sh <- forAll $ genShR inrank
-  -- traceM ("sh: " ++ show sh ++ " -> " ++ show (product sh))
-  guard (all (> 0) (shrTail sh))  -- only constrain the tail
-  arr <- forAllT $ OR.fromVector @Double @(n + 1) (toList sh) <$>
-           genStorables (Range.singleton (product sh))
+  -- traceM ("sh: " ++ show sh ++ " -> " ++ show (shrSize sh))
+  guard (all (> 0) (shrToList $ shrTail sh))  -- only constrain the tail
+  arr <- forAllT $ OR.fromVector @Double @(n + 1) (shrToList sh) <$>
+           genStorables (Range.singleton (shrSize sh))
                         (\w -> fromIntegral w / fromIntegral (maxBound :: Word64))
   f inrank outrank arr
 
@@ -67,19 +66,19 @@ gen_red_empty f = property $ genRank $ \outrankm1@(SNat @nm1) -> do
     sht <- shuffleShR (0 :$: shtt)  -- n
     n <- Gen.int (Range.linear 0 20)
     return (n :$: sht)  -- n + 1
-  guard (0 `elem` shrTail sh)
+  guard (0 `elem` (shrToList $ shrTail sh))
   -- traceM ("sh: " ++ show sh ++ " -> " ++ show (product sh))
-  let arr = OR.fromList @(n + 1) @Double (toList sh) []
+  let arr = OR.fromList @(n + 1) @Double (shrToList sh) []
   f inrank arr
 
 gen_red_lasteq1 :: (forall n. SNat (n + 1) -> SNat n -> OR.Array (n + 1) Double -> PropertyT IO ()) -> Property
 gen_red_lasteq1 f = property $ genRank $ \outrank@(SNat @n) -> do
   let inrank = SNat @(n + 1)
   outsh <- forAll $ genShR outrank
-  guard (all (> 0) outsh)
+  guard (all (> 0) $ shrToList outsh)
   let insh = shrAppend outsh (1 :$: ZSR)
-  arr <- forAllT $ OR.fromVector @Double @(n + 1) (toList insh) <$>
-           genStorables (Range.singleton (product insh))
+  arr <- forAllT $ OR.fromVector @Double @(n + 1) (shrToList insh) <$>
+           genStorables (Range.singleton (shrSize insh))
                         (\w -> fromIntegral w / fromIntegral (maxBound :: Word64))
   f inrank outrank arr
 
@@ -96,12 +95,12 @@ gen_red_replicated doTranspose f = property $
       label (LabelName ("rankdiff " ++ show (fromSNat' inrank2 - fromSNat' inrank1)))
       label (LabelName ("size sh1 10^" ++ show (floor (logBase 10 (fromIntegral (shrSize sh1) :: Double)) :: Int)))
       label (LabelName ("size sh3 10^" ++ show (floor (logBase 10 (fromIntegral (shrSize sh3) :: Double)) :: Int)))
-    guard (all (> 0) sh3)
+    guard (all (> 0) $ shrToList sh3)
     arr <- forAllT $
-      OR.stretch (toList sh3)
-      . OR.reshape (toList sh2)
-      . OR.fromVector @Double @m (toList sh1) <$>
-               genStorables (Range.singleton (product sh1))
+      OR.stretch (shrToList sh3)
+      . OR.reshape (shrToList sh2)
+      . OR.fromVector @Double @m (shrToList sh1) <$>
+               genStorables (Range.singleton (shrSize sh1))
                             (\w -> fromIntegral w / fromIntegral (maxBound :: Word64))
     arrTrans <-
       if doTranspose then do perm <- forAll $ genPermR (fromSNat' inrank2)
@@ -160,9 +159,9 @@ prop_negate_with :: forall f b. Show b
 prop_negate_with genRank' genB preproc = property $
   genRank' $ \extra rank@(SNat @n) -> do
     sh <- forAll $ genShR rank
-    guard (all (> 0) sh)
-    arr <- forAllT $ OR.fromVector @Double @n (toList sh) <$>
-             genStorables (Range.singleton (product sh))
+    guard (all (> 0) $ shrToList sh)
+    arr <- forAllT $ OR.fromVector @Double @n (shrToList sh) <$>
+             genStorables (Range.singleton (shrSize sh))
                           (\w -> fromIntegral w / fromIntegral (maxBound :: Word64))
     bval <- forAll $ genB extra sh
     let arr' = preproc extra bval arr
@@ -202,7 +201,7 @@ tests = testGroup "C"
         (\_ sh -> do let genPair n = do lo <- Gen.integral (Range.constant 0 (n-1))
                                         len <- Gen.integral (Range.constant 0 (n-lo-1))
                                         return (lo, len)
-                     pairs <- mapM genPair (toList sh)
+                     pairs <- mapM genPair (shrToList sh)
                      return pairs)
         (\_ -> OR.slice)
     ]