From 03c1bdff08fff2f9bcdeceb7e59113e4b79c9b1f Mon Sep 17 00:00:00 2001
From: Maximilian Algehed <m.algehed@gmail.com>
Date: Thu, 14 Nov 2024 10:44:50 +0100
Subject: [PATCH] `constrained-generators`: Fix test failure with narrowing

---
 .../src/Constrained/Base.hs                   | 33 ++++++++++---------
 .../test/Constrained/Test.hs                  | 13 ++++++++
 2 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/libs/constrained-generators/src/Constrained/Base.hs b/libs/constrained-generators/src/Constrained/Base.hs
index 1d0f8b5168c..6df2ce57b9d 100644
--- a/libs/constrained-generators/src/Constrained/Base.hs
+++ b/libs/constrained-generators/src/Constrained/Base.hs
@@ -3209,7 +3209,7 @@ narrowByFuelAndSize ::
   (Specification fn a, Specification fn a) ->
   (Specification fn a, Specification fn a)
 narrowByFuelAndSize fuel size specs =
-  loop (1000 :: Int) (narrowFoldSpecs specs)
+  loop (100 :: Int) (onlyOnceTransformations $ narrowFoldSpecs specs)
   where
     loop 0 specs =
       error $
@@ -3225,6 +3225,13 @@ narrowByFuelAndSize fuel size specs =
       Nothing -> specs
       Just specs' -> loop (n - 1) (narrowFoldSpecs specs')
 
+    -- Transformations only applied once. It's annoying to check if you're
+    -- going to change the spec with these so easier to just make sure you only apply
+    -- these once
+    onlyOnceTransformations (elemS, foldS)
+      | fuel == 1 = (elemS <> foldS, foldS)
+      | otherwise = (elemS, foldS)
+
     canReach _ 0 s = s == 0
     canReach e fuel s
       -- You can reach it in one step
@@ -3246,9 +3253,11 @@ narrowByFuelAndSize fuel size specs =
       where
         d = a `div` b
 
-    go (elemS, foldS)
+    go (simplifySpec -> elemS, simplifySpec -> foldS)
       -- There is nothing we can do
       | fuel == 0 = Nothing
+      | ErrorSpec {} <- elemS = Nothing
+      | ErrorSpec {} <- foldS = Nothing
       -- Give up as early as possible
       | Just 0 <- knownUpperBound elemS
       , Just 0 <- knownLowerBound elemS
@@ -3285,14 +3294,15 @@ narrowByFuelAndSize fuel size specs =
       , Just s <- knownUpperBound foldS
       , s < 0
       , let c = divCeil (safeNegate s) fuel
-      , negate c < e =
+      , negate c < e
+      , maybe True (c <) (knownUpperBound elemS) =
           Just (elemS <> leqSpec @fn c, foldS)
       -- It's time to stop generating negative numbers
       | Just s <- knownLowerBound foldS
       , s > 0
       , Just e <- knownUpperBound elemS
       , e > 0
-      , not $ canReach e (fuel `div` 2) s
+      , not $ canReach e (fuel `div` 2 + 1) s
       , maybe True (<= 0) (knownLowerBound elemS) =
           Just (elemS <> gtSpec @fn 0, foldS)
       -- It's time to stop generating positive numbers
@@ -3300,16 +3310,9 @@ narrowByFuelAndSize fuel size specs =
       , s < 0
       , Just e <- knownLowerBound elemS
       , e < 0
-      , not $ canReach (safeNegate e) (fuel `div` 2) (safeNegate s)
-      , maybe True (<= 0) (knownLowerBound elemS) =
-          Just (elemS <> gtSpec @fn 0, foldS)
-      -- We HAVE to set the lower bound to the lower
-      -- bound on the sum
-      | Just s <- knownLowerBound foldS
-      , fuel == 1
-      , s `conformsToSpec` elemS
-      , maybe True (< s) (knownLowerBound elemS) =
-          Just (elemS <> geqSpec @fn s, foldS)
+      , not $ canReach (safeNegate e) (fuel `div` 2 + 1) (safeNegate s)
+      , maybe True (0 <=) (knownUpperBound elemS) =
+          Just (elemS <> ltSpec @fn 0, foldS)
       -- There is nothing we need to do
       | otherwise = Nothing
 
@@ -3329,7 +3332,7 @@ narrowFoldSpecs ::
 narrowFoldSpecs specs = maybe specs narrowFoldSpecs (go specs)
   where
     -- Note: make sure there is some progress when returning Just or this will loop forever
-    go (elemS, foldS) = case (elemS, foldS) of
+    go (simplifySpec -> elemS, simplifySpec -> foldS) = case (elemS, foldS) of
       -- Empty foldSpec
       (_, ErrorSpec {}) -> Nothing
       _ | isEmptyNumSpec foldS -> Just (elemS, ErrorSpec (NE.fromList ["Empty foldSpec:", show foldS]))
diff --git a/libs/constrained-generators/test/Constrained/Test.hs b/libs/constrained-generators/test/Constrained/Test.hs
index ea0e9da572a..b93f76dc75b 100644
--- a/libs/constrained-generators/test/Constrained/Test.hs
+++ b/libs/constrained-generators/test/Constrained/Test.hs
@@ -182,6 +182,7 @@ tests nightly =
       prop "[(Set Int, Set Bool)]" $ prop_gen_sound @BaseFn @[(Set Int, Set Bool)]
       prop "Set (Set Bool)" $ prop_gen_sound @BaseFn @(Set (Set Bool))
     negativeTests
+    prop "prop_noNarrowLoop" $ withMaxSuccess 1000 prop_noNarrowLoop
 
 negativeTests :: Spec
 negativeTests =
@@ -369,3 +370,15 @@ hasSizeSet = hasSize (rangeSize 1 3)
 
 hasSizeMap :: Specification BaseFn (Map Int Int)
 hasSizeMap = hasSize (rangeSize 1 3)
+
+------------------------------------------------------------------------
+-- Tests for narrowing
+------------------------------------------------------------------------
+
+prop_noNarrowLoop :: Int -> Int -> Specification BaseFn Int -> Specification BaseFn Int -> Property
+prop_noNarrowLoop f s eSpec fSpec =
+  -- Make sure the fuel is non-negative
+  f >= 0 ==>
+    discardAfter 100_000 $
+      narrowByFuelAndSize f s (eSpec, fSpec) `seq`
+        property True