Optimise how we deal with either specs

bench/Constrained/Bench.hs

@@ -10,6 +10,11 @@
 module Constrained.Bench where
 
 import Constrained.API
+import Constrained.Generation
+import Constrained.Examples.Set
+import Constrained.Examples.Map
+import Constrained.Examples.Basic
+
 import Control.DeepSeq
 import Criterion
 import Data.Map (Map)
@@ -30,6 +35,10 @@ benchmarks =
         (giveHint (Nothing, 30) <> trueSpec :: Specification (Tree Int))
     , benchSpec 10 30 "roseTreeMaybe" roseTreeMaybe
     , benchSpec 10 30 "listSumPair" listSumPair
+    , benchSpec 10 30 "maybeJustSetSpec" maybeJustSetSpec
+    , benchSpec 10 40 "eitherKeys" eitherKeys
+    , benchSimplifySpec "eitherKeys" eitherKeys
+    , benchSpec 10 40 "chooseBackwards" chooseBackwards'
     ]
 
 roseTreeMaybe :: Specification (Tree (Maybe (Int, Int)))
@@ -49,6 +58,11 @@ listSumPair = constrained $ \xs ->
   , forAll' xs $ \x y -> [20 <. x, x <. 30, y <. 100]
   ]
 
+benchSimplifySpec :: HasSpec a => String -> Specification a -> Benchmark
+benchSimplifySpec nm spec =
+  bench ("simplify/" ++ nm) $
+    nf (show . simplifySpec) spec
+
 benchSpec :: (HasSpec a, NFData a) => Int -> Int -> String -> Specification a -> Benchmark
 benchSpec seed size nm spec =
   bench (unlines [nm, show (genFromSpecWithSeed seed size spec)]) $

constrained-generators.cabal

@@ -163,6 +163,7 @@ benchmark bench
   build-depends:
     base,
     constrained-generators,
+    constrained-generators:examples,
     containers,
     criterion,
     deepseq,

examples/Constrained/Examples/Map.hs

@@ -138,3 +138,14 @@ eitherKeys = constrained' $ \ [var| as |] [var| bs |] [var| m |] ->
           k `elem_` ls ++. rs
       ]
   ]
+
+keysExample :: Specification (Either Int Int)
+keysExample = constrained $ \ k ->
+  [ caseOn k
+      (branch $ \ a -> a `elem_` as)
+      (branch $ \ b -> b `elem_` bs)
+  , reify as (map Left) $ \ ls ->
+    reify bs (map Right) $ \ rs ->
+      k `elem_` ls ++. rs
+  ] where as = lit [ 1 .. 10]
+          bs = lit [ 11 .. 20 ]

src/Constrained/Generation.hs

@@ -101,6 +101,8 @@ import Prelude hiding (cycle, pred)
 -- generators are not flexible enough.
 genFromSpecT ::
   forall a m. (HasCallStack, HasSpec a, MonadGenError m) => Specification a -> GenT m a
+genFromSpecT (ExplainSpec [] s) = genFromSpecT s
+genFromSpecT (ExplainSpec es s) = push es (genFromSpecT s)
 genFromSpecT (simplifySpec -> spec) = case spec of
   ExplainSpec [] s -> genFromSpecT s
   ExplainSpec es s -> push es (genFromSpecT s)
@@ -947,7 +949,6 @@ genFromPreds env0 (optimisePred . optimisePred -> preds) = do
           (env', plan') <- stepPlan origPlan env plan
           go env' plan'
     go env0 origPlan
-  where
 
 -- | Push as much information we can backwards through the plan.
 backPropagation :: Set Name -> SolverPlan -> SolverPlan
@@ -1050,7 +1051,10 @@ pinnedBy :: forall a. HasSpec a => Var a -> Pred -> Maybe (Term a)
 pinnedBy x (Assert (Equal t t'))
   | V x' <- t, Just Refl <- eqVar x x' = Just t'
   | V x' <- t', Just Refl <- eqVar x x' = Just t
+pinnedBy x (ElemPred True (V x') (xs NE.:| []))
+  | Just Refl <- eqVar x x' = Just (lit xs)
 pinnedBy x (And ps) = listToMaybe $ catMaybes $ map (pinnedBy x) ps
+pinnedBy x (Explain _ p) = pinnedBy x p
 pinnedBy _ _ = Nothing
 
 -- ==================================================================================================

src/Constrained/NumOrd.hs

@@ -1063,7 +1063,6 @@ caseBoolSpec spec cont = case possibleValues spec of
   [b] -> cont b
   _ -> mempty
   where
-    -- where possibleValues s = filter (flip conformsToSpec (simplifySpec s)) [True, False]
     -- This will always get the same result, and probably faster since running 2
     -- conformsToSpec on True and False takes less time than simplifying the spec.
     -- Since we are in TheKnot, we could keep the simplifySpec. Is there a good reason to?

src/Constrained/Spec/List.hs

@@ -17,7 +17,7 @@
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE UndecidableSuperClasses #-}
 {-# LANGUAGE ViewPatterns #-}
-{-# OPTIONS_GHC -Wno-orphans #-}
+{-# OPTIONS_GHC -Wno-orphans -Wno-name-shadowing #-}
 
 -- | `TypeSpec` definition for `[]` and functions for writing constraints over
 -- lists
@@ -227,7 +227,8 @@ instance HasSpec a => HasSpec [a] where
       Just szHint -> do
         sz <- genFromSizeSpec (leqSpec szHint)
         listOfUntilLenT (genFromSpecT elemS) (fromIntegral sz) (const True)
-    pureGen $ shuffle (must ++ lst)
+    must' <- pureGen $ shuffle must
+    pureGen $ randomInterleaving must' lst
   genFromTypeSpec (ListSpec msz must szSpec elemS NoFold) = do
     sz0 <- genFromSizeSpec (szSpec <> geqSpec (sizeOf must) <> maybe TrueSpec (leqSpec . max 0) msz)
     let sz = fromIntegral (sz0 - sizeOf must)
@@ -236,7 +237,8 @@ instance HasSpec a => HasSpec [a] where
         (genFromSpecT elemS)
         sz
         ((`conformsToSpec` szSpec) . (+ sizeOf must) . fromIntegral)
-    pureGen $ shuffle (must ++ lst)
+    must' <- pureGen $ shuffle must
+    pureGen $ randomInterleaving must' lst
   genFromTypeSpec (ListSpec msz must szSpec elemS (FoldSpec f foldS)) = do
     let szSpec' = szSpec <> maybe TrueSpec (leqSpec . max 0) msz
     genFromFold must szSpec' elemS f foldS
@@ -266,6 +268,21 @@ instance HasSpec a => HasSpec [a] where
       <> satisfies (sizeOf_ x) size
       <> maybe TruePred (flip genHint x) msz
 
+randomInterleaving :: [a] -> [a] -> Gen [a]
+randomInterleaving xs ys = go xs ys (length ys)
+  where
+    go [] ys _ = pure ys
+    go xs [] _ = pure xs
+    go xs ys l = do
+      -- TODO: think about distribution here
+      i <- choose (0, l)
+      go' i xs ys (l - i)
+
+    go' _ xs [] _ = pure xs
+    go' _ [] ys _ = pure ys
+    go' 0 (x:xs) ys l = (x:) <$> go xs ys l
+    go' i xs (y:ys) l = (y:) <$> go' (i-1) xs ys l
+
 instance HasSpec a => HasGenHint [a] where
   type Hint [a] = Integer
   giveHint szHint = typeSpec $ ListSpec (Just szHint) [] mempty mempty NoFold

src/Constrained/Spec/Map.hs

@@ -179,7 +179,7 @@ instance
               fold
                 [ maybe TrueSpec (leqSpec . max 0) mHint
                 , size
-                , maxSpec (cardinality (fstSpec kvs)) -- (mapSpec FstW  (mapSpec ToGenericW kvs)))
+                , maxSpec (cardinality (fstSpec kvs))
                 , maxSpec (cardinalTrueSpec @k)
                 , geqSpec 0
                 ]

src/Constrained/Syntax.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
@@ -721,6 +722,7 @@ mkCase tm cs
   | Semigroup.getAll $ foldMapList isTrueBinder cs = TruePred
   | Semigroup.getAll $ foldMapList (isFalseBinder . thing) cs = FalsePred (pure "mkCase on all False")
   | Lit a <- tm = runCaseOn a (mapList thing cs) (\x val p -> substPred (Env.singleton x val) p)
+  | Just es <- buildElemList cs = ElemPred True tm es
   | otherwise = Case tm cs
   where
     isTrueBinder (Weighted Nothing (_ :-> TruePred)) = Semigroup.All True
@@ -729,6 +731,17 @@ mkCase tm cs
     isFalseBinder (_ :-> FalsePred {}) = Semigroup.All True
     isFalseBinder _ = Semigroup.All False
 
+buildElemList :: List (Weighted Binder) as -> Maybe (NE.NonEmpty (SumOver as))
+buildElemList Nil = Nothing
+buildElemList (Weighted Nothing (x :-> ElemPred True (V x') as) :> xs)
+  | Just Refl <- eqVar x x' =
+    case xs of
+      Nil -> Just as
+      _ :> _ -> do
+        rest <- buildElemList xs
+        return $ fmap SumLeft as <> fmap SumRight rest
+buildElemList _ = Nothing
+
 -- | Run a `caseOn`
 runCaseOn ::
   SumOver as ->

src/Constrained/Test.hs

@@ -75,8 +75,8 @@ prop_constrained_satisfies_sound spec = prop_sound (constrained $ \a -> satisfie
 -- | Check that explanations don't immediately ruin soundness
 prop_constrained_explained :: HasSpec a => Specification a -> QC.Property
 prop_constrained_explained spec =
-  QC.forAll QC.arbitrary $ \es ->
-    prop_sound $ constrained $ \x -> Explain es $ x `satisfies` spec
+  let es = NE.singleton "Dummy explanation"
+  in prop_sound $ constrained $ \x -> Explain es $ x `satisfies` spec
 
 -- | `prop_complete ps` assumes that `ps` is satisfiable and checks that it doesn't crash
 prop_complete :: HasSpec a => Specification a -> QC.Property

test/Constrained/Tests.hs

@@ -205,6 +205,7 @@ tests nightly =
     testSpec "composeEvenSpec" composeEvenSpec
     testSpec "oddSpec" oddSpec
     testSpec "composeOddSpec" composeOddSpec
+    testSpec "keysExample" keysExample
 
 negativeTests :: Spec
 negativeTests =