optimise how we deal with either specs

Author: MaximilianAlgehed

bench/Constrained/Bench.hs

@@ -10,6 +10,10 @@
 module Constrained.Bench where
 
 import Constrained.API
+import Constrained.Generation
+import Constrained.Examples.Set
+import Constrained.Examples.Map
+
 import Control.DeepSeq
 import Criterion
 import Data.Map (Map)
@@ -30,6 +34,9 @@ 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
     ]
 
 roseTreeMaybe :: Specification (Tree (Maybe (Int, Int)))
@@ -49,6 +56,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/Map.hs

@@ -179,12 +179,13 @@ 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
                 ]
         n <- genFromSpecT size'
         let go 0 _ m = pure m
+            -- go n' ((\ s -> trace (show s) s) -> kvs') m = do
             go n' kvs' m = do
               mkv <- inspect $ genFromSpecT kvs'
               case mkv of

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 ->