arith: Some negate tests (to check stride handling)

1 files changed, 32 insertions, 6 deletions

diff --git a/test/Tests/C.hs b/test/Tests/C.hs
index 97b425f..bc8e0de 100644
--- a/test/Tests/C.hs
+++ b/test/Tests/C.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeAbstractions #-}
 {-# LANGUAGE TypeApplications #-}
@@ -12,6 +13,7 @@ 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
 import GHC.TypeLits
@@ -97,16 +99,23 @@ prop_sum_replicated doTranspose = property $
     let rarr = rfromOrthotope inrank2 arrTrans
     almostEq 1e-8 (rtoOrthotope (rsumOuter1 rarr)) (orSumOuter1 outrank arrTrans)
 
-prop_negate_normalised :: Property
-prop_negate_normalised = property $
-  genRank $ \rank@(SNat @n) -> do
+prop_negate_with :: forall f b. Show b
+                 => ((forall n. f n -> SNat n -> PropertyT IO ()) -> PropertyT IO ())
+                 -> (forall n. f n -> IShR n -> Gen b)
+                 -> (forall n. f n -> b -> OR.Array n Double -> OR.Array n Double)
+                 -> Property
+prop_negate_with genRank' genB preproc = property $
+  genRank' $ \extra rank@(SNat @n) -> do
     sh <- forAll $ genShR rank
     guard (all (> 0) (toList sh))
     arr <- forAllT $ OR.fromVector @Double @n (toList sh) <$>
              genStorables (Range.singleton (product sh))
                           (\w -> fromIntegral w / fromIntegral (maxBound :: Word64))
-    let rarr = rfromOrthotope rank arr
-    rtoOrthotope (negate rarr) === OR.mapA negate arr
+    bval <- forAll $ genB extra sh
+    let arr' = preproc extra bval arr
+    annotate (show (OR.shapeL arr'))
+    let rarr = rfromOrthotope rank arr'
+    rtoOrthotope (negate rarr) === OR.mapA negate arr'
 
 tests :: TestTree
 tests = testGroup "C"
@@ -118,6 +127,23 @@ tests = testGroup "C"
     ,testProperty "replicated_transposed" (prop_sum_replicated True)
     ]
   ,testGroup "negate"
-    [testProperty "normalised" prop_negate_normalised
+    [testProperty "normalised" $ prop_negate_with
+        (\k -> genRank (k (Const ())))
+        (\_ _ -> pure ())
+        (\_ _ -> id)
+    ,testProperty "slice 1D" $ prop_negate_with @((:~:) 1)
+        (\k -> k Refl (SNat @1))
+        (\Refl (n :$: _) -> do lo <- Gen.integral (Range.constant 0 (n-1))
+                               len <- Gen.integral (Range.constant 0 (n-lo))
+                               return [(lo, len)])
+        (\_ -> OR.slice)
+    ,testProperty "slice nD" $ prop_negate_with
+        (\k -> genRank (k (Const ())))
+        (\_ 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)
+                     return pairs)
+        (\_ -> OR.slice)
     ]
   ]