Optimise numspec

bench/Constrained/Bench.hs

@@ -10,6 +10,7 @@
 module Constrained.Bench where
 
 import Constrained.API
+import Constrained.Generation
 import Control.DeepSeq
 import Criterion
 import Data.Map (Map)
@@ -30,8 +31,12 @@ benchmarks =
         (giveHint (Nothing, 30) <> trueSpec :: Specification (Tree Int))
     , benchSpec 10 30 "roseTreeMaybe" roseTreeMaybe
     , benchSpec 10 30 "listSumPair" listSumPair
+    , benchSpec 10 30 "intSpec" (simplifySpec intSpecSimple)
     ]
 
+intSpecSimple :: Specification Int
+intSpecSimple = constrained $ \ x -> 0 <. x
+
 roseTreeMaybe :: Specification (Tree (Maybe (Int, Int)))
 roseTreeMaybe = constrained $ \t ->
   [ forAll' t $ \mp ts ->

src/Constrained/GenT.hs

@@ -28,6 +28,7 @@ module Constrained.GenT (
   tryGenT,
   chooseT,
   sizeT,
+  sizedT,
   withMode,
   frequencyT,
   oneofT,
@@ -109,7 +110,9 @@ data GenMode
 -- | A `Gen` monad wrapper that allows different generation modes and different
 -- failure types.
 newtype GenT m a = GenT {runGenT :: GenMode -> [NonEmpty String] -> Gen (m a)}
-  deriving (Functor)
+
+instance Functor f => Functor (GenT f) where
+  fmap f (GenT k) = GenT $ \ mode msgs -> fmap (fmap f) (k mode msgs)
 
 instance Monad m => Applicative (GenT m) where
   pure a = GenT (\_ _ -> pure @Gen (pure @m a))
@@ -388,7 +391,11 @@ chooseT (a, b)
 
 -- | Get the size provided to the generator
 sizeT :: Monad m => GenT m Int
-sizeT = GenT $ \mode msgs -> sized $ \n -> runGenT (pure n) mode msgs
+sizeT = sizedT pure
+
+-- | Get the size provided to the generator
+sizedT :: (Int -> GenT m a) -> GenT m a
+sizedT c = GenT $ \mode msgs -> sized $ \n -> runGenT (c n) mode msgs
 
 -- ==================================================================
 -- Reflective analysis of the internal GE structure of (GenT GE x)
@@ -397,27 +404,23 @@ sizeT = GenT $ \mode msgs -> sized $ \n -> runGenT (pure n) mode msgs
 
 -- | Always succeeds, but returns the internal GE structure for analysis
 inspect :: forall m x. MonadGenError m => GenT GE x -> GenT m (GE x)
-inspect (GenT f) = GenT g
-  where
-    g mode msgs = do geThing <- f mode msgs; pure @Gen (pure @m geThing)
+inspect (GenT f) = GenT (\ mode msgs -> pure <$> f mode msgs)
 
 -- | Ignore all kinds of Errors, by squashing them into Nothing
 tryGenT :: MonadGenError m => GenT GE a -> GenT m (Maybe a)
-tryGenT g = do
-  r <- inspect g
-  case r of
-    FatalError _ -> pure Nothing
-    GenError _ -> pure Nothing
-    Result a -> pure $ Just a
+tryGenT g = cont <$> inspect g
+  where cont r = case r of
+          FatalError _ -> Nothing
+          GenError _ -> Nothing
+          Result a -> Just a
 
 -- Pass on the error messages of both kinds of Errors, by squashing and combining both of them into Left constructor
 catchGenT :: MonadGenError m => GenT GE a -> GenT m (Either (NonEmpty (NonEmpty String)) a)
-catchGenT g = do
-  r <- inspect g
-  case r of
-    FatalError es -> pure $ Left es
-    GenError es -> pure $ Left es
-    Result a -> pure $ Right a
+catchGenT g = cont <$> inspect g
+  where cont r = case r of
+          FatalError es -> Left es
+          GenError es -> Left es
+          Result a -> Right a
 
 -- | Pass on the error messages of both kinds of Errors in the Gen (not the GenT) monad
 catchGen :: GenT GE a -> Gen (Either (NonEmpty (NonEmpty String)) a)

src/Constrained/Generation.hs

@@ -119,27 +119,16 @@ genFromSpecT (simplifySpec -> spec) = case spec of
         env <- genFromPreds mempty p
         Env.find env x
   TypeSpec s cant -> do
-    mode <- getMode
-    explainNE
-      ( NE.fromList
-          [ "genFromSpecT on (TypeSpec tspec cant) at type " ++ showType @a
-          , "tspec = "
-          , show s
-          , "cant = " ++ show cant
-          , "with mode " ++ show mode
-          ]
-      )
-      $
-      -- TODO: we could consider giving `cant` as an argument to `genFromTypeSpec` if this
-      -- starts giving us trouble.
-      genFromTypeSpec s `suchThatT` (`notElem` cant)
+    -- TODO: we could consider giving `cant` as an argument to `genFromTypeSpec` if this
+    -- starts giving us trouble.
+    case cant of
+      [] -> genFromTypeSpec s
+      _  -> genFromTypeSpec s `suchThatT` (`notElem` cant)
   ErrorSpec e -> genErrorNE e
 
 -- | A version of `genFromSpecT` that simply errors if the generator fails
 genFromSpec :: forall a. (HasCallStack, HasSpec a) => Specification a -> Gen a
-genFromSpec spec = do
-  res <- catchGen $ genFromSpecT @a @GE spec
-  either (error . ('\n' :) . catMessages) pure res
+genFromSpec spec = genFromGenT (genFromSpecT @a @GE spec)
 
 -- | A version of `genFromSpecT` that takes a seed and a size and gives you a result
 genFromSpecWithSeed ::

src/Constrained/NumOrd.hs

@@ -311,9 +311,11 @@ genFromNumSpec ::
   (MonadGenError m, Show n, Random n, Ord n, Num n, MaybeBounded n) =>
   NumSpec n ->
   GenT m n
-genFromNumSpec (NumSpecInterval ml mu) = do
-  n <- sizeT
-  pureGen . choose =<< constrainInterval (ml <|> lowerBound) (mu <|> upperBound) (fromIntegral n)
+genFromNumSpec (NumSpecInterval ml mu) =
+  sizedT $ \ n ->
+    case constrainInterval (ml <|> lowerBound) (mu <|> upperBound) (fromIntegral n) of
+      Just interval -> pureGen $ choose interval
+      Nothing -> genError $ "bad interval: " ++ show ml ++ " " ++ show mu
 
 -- TODO: fixme
 
@@ -326,8 +328,9 @@ shrinkWithNumSpec _ = shrink
 fixupWithNumSpec :: Arbitrary n => NumSpec n -> n -> Maybe n
 fixupWithNumSpec _ = listToMaybe . shrink
 
+{-# SCC constrainInterval #-}
 constrainInterval ::
-  (MonadGenError m, Ord a, Num a, Show a) => Maybe a -> Maybe a -> Integer -> m (a, a)
+  (Ord a, Num a) => Maybe a -> Maybe a -> Integer -> Maybe (a, a)
 constrainInterval ml mu r =
   case (ml, mu) of
     (Nothing, Nothing) -> pure (-r', r')
@@ -336,21 +339,20 @@ constrainInterval ml mu r =
       | otherwise -> pure (l, l + 2 * r')
     (Nothing, Just u)
       | u > 0 -> pure (negate r', min u r')
-      | otherwise -> pure (u - r' - r', u)
+      | otherwise -> pure (u - 2 * r', u)
     (Just l, Just u)
-      | l > u -> genError ("bad interval: " ++ show l ++ " " ++ show u)
+      | l > u -> Nothing
       | u < 0 -> pure (safeSub l (safeSub l u r') r', u)
-      | l >= 0 -> pure (l, safeAdd u (safeAdd u l r') r')
+      | l >= 0 -> pure (l, safeAdd u l r')
       -- TODO: this is a bit suspect if the bounds are lopsided
       | otherwise -> pure (max l (-r'), min u r')
   where
     r' = abs $ fromInteger r
     safeSub l a b
       | a - b > a = l
       | otherwise = max l (a - b)
-    safeAdd u a b
-      | a + b < a = u
-      | otherwise = min u (a + b)
+    safeAdd u a b =
+      let ab = a + b in if ab < a then u else min u ab
 
 -- | Check that a value is in the spec
 conformsToNumSpec :: Ord n => n -> NumSpec n -> Bool