-----------------------------------------------------------------------------
-- |
-- Module    : Documentation.SBV.Examples.Queries.AllSat
-- Copyright : (c) Levent Erkok
-- License   : BSD3
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- When we would like to find all solutions to a problem, we can query the
-- solver repeatedly, telling it to give us a new model each time. SBV already
-- provides 'Data.SBV.allSat' that precisely does this. However, this example demonstrates
-- how the query mode can be used to achieve the same, and can also incorporate
-- extra conditions with easy as we walk through solutions.
-----------------------------------------------------------------------------

{-# OPTIONS_GHC -Wall -Werror #-}

module Documentation.SBV.Examples.Queries.AllSat where

import Data.SBV
import Data.SBV.Control

-- | Find all solutions to @x + y .== 10@ for positive @x@ and @y@.
-- This is rather silly to do in the query mode as `allSat` can do
-- this automatically, but it demonstrates how we can dynamically
-- query the result and put in new constraints based on those.
goodSum :: Symbolic [(Integer, Integer)]
goodSum :: Symbolic [(Integer, Integer)]
goodSum = do SInteger
x <- String -> Symbolic SInteger
sInteger String
"x"
             SInteger
y <- String -> Symbolic SInteger
sInteger String
"y"

             -- constrain positive and sum:
             SBool -> SymbolicT IO ()
forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain (SBool -> SymbolicT IO ()) -> SBool -> SymbolicT IO ()
forall a b. (a -> b) -> a -> b
$ SInteger
x SInteger -> SInteger -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.>= SInteger
0
             SBool -> SymbolicT IO ()
forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain (SBool -> SymbolicT IO ()) -> SBool -> SymbolicT IO ()
forall a b. (a -> b) -> a -> b
$ SInteger
y SInteger -> SInteger -> SBool
forall a. OrdSymbolic a => a -> a -> SBool
.>= SInteger
0
             SBool -> SymbolicT IO ()
forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain (SBool -> SymbolicT IO ()) -> SBool -> SymbolicT IO ()
forall a b. (a -> b) -> a -> b
$ SInteger
x SInteger -> SInteger -> SInteger
forall a. Num a => a -> a -> a
+ SInteger
y SInteger -> SInteger -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== SInteger
10

             -- Capture the "next" solution function:
             let next :: Integer -> [(Integer, Integer)] -> QueryT IO [(Integer, Integer)]
next Integer
i [(Integer, Integer)]
sofar = do
                    IO () -> Query ()
forall a. IO a -> Query a
io (IO () -> Query ()) -> IO () -> Query ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Iteration: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show (Integer
i :: Integer)

                    -- Using a check-sat assuming, we force the solver to walk through
                    -- the entire range of x's
                    CheckSatResult
cs <- [SBool] -> Query CheckSatResult
checkSatAssuming [SInteger
x SInteger -> SInteger -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== Integer -> SInteger
forall a. SymVal a => a -> SBV a
literal (Integer
iInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1)]

                    case CheckSatResult
cs of
                      CheckSatResult
Unk    -> String -> QueryT IO [(Integer, Integer)]
forall a. HasCallStack => String -> a
error String
"Too bad, solver said unknown.." -- Won't happen
                      DSat{} -> String -> QueryT IO [(Integer, Integer)]
forall a. HasCallStack => String -> a
error String
"Unexpected dsat result.."       -- Won't happen
                      CheckSatResult
Unsat  -> do IO () -> Query ()
forall a. IO a -> Query a
io (IO () -> Query ()) -> IO () -> Query ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStrLn String
"No other solution!"
                                   [(Integer, Integer)] -> QueryT IO [(Integer, Integer)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Integer, Integer)] -> QueryT IO [(Integer, Integer)])
-> [(Integer, Integer)] -> QueryT IO [(Integer, Integer)]
forall a b. (a -> b) -> a -> b
$ [(Integer, Integer)] -> [(Integer, Integer)]
forall a. [a] -> [a]
reverse [(Integer, Integer)]
sofar

                      CheckSatResult
Sat    -> do Integer
xv <- SInteger -> Query Integer
forall a. SymVal a => SBV a -> Query a
getValue SInteger
x
                                   Integer
yv <- SInteger -> Query Integer
forall a. SymVal a => SBV a -> Query a
getValue SInteger
y

                                   IO () -> Query ()
forall a. IO a -> Query a
io (IO () -> Query ()) -> IO () -> Query ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Current solution is: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Integer, Integer) -> String
forall a. Show a => a -> String
show (Integer
xv, Integer
yv)

                                   -- For next iteration: Put in constraints outlawing the current one:
                                   -- Note that we do *not* put these separately, as we do want
                                   -- to allow repetition on one value if the other is different!
                                   SBool -> Query ()
forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain (SBool -> Query ()) -> SBool -> Query ()
forall a b. (a -> b) -> a -> b
$   SInteger
x SInteger -> SInteger -> SBool
forall a. EqSymbolic a => a -> a -> SBool
./= Integer -> SInteger
forall a. SymVal a => a -> SBV a
literal Integer
xv
                                             SBool -> SBool -> SBool
.|| SInteger
y SInteger -> SInteger -> SBool
forall a. EqSymbolic a => a -> a -> SBool
./= Integer -> SInteger
forall a. SymVal a => a -> SBV a
literal Integer
yv

                                   -- loop around!
                                   Integer -> [(Integer, Integer)] -> QueryT IO [(Integer, Integer)]
next (Integer
iInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1) ((Integer
xv, Integer
yv) (Integer, Integer) -> [(Integer, Integer)] -> [(Integer, Integer)]
forall a. a -> [a] -> [a]
: [(Integer, Integer)]
sofar)

             -- Go into query mode and execute the loop:
             QueryT IO [(Integer, Integer)] -> Symbolic [(Integer, Integer)]
forall a. Query a -> Symbolic a
query (QueryT IO [(Integer, Integer)] -> Symbolic [(Integer, Integer)])
-> QueryT IO [(Integer, Integer)] -> Symbolic [(Integer, Integer)]
forall a b. (a -> b) -> a -> b
$ do IO () -> Query ()
forall a. IO a -> Query a
io (IO () -> Query ()) -> IO () -> Query ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStrLn String
"Starting the all-sat engine!"
                        Integer -> [(Integer, Integer)] -> QueryT IO [(Integer, Integer)]
next Integer
1 []

-- | Run the query. We have:
--
-- >>> demo
-- Starting the all-sat engine!
-- Iteration: 1
-- Current solution is: (0,10)
-- Iteration: 2
-- Current solution is: (1,9)
-- Iteration: 3
-- Current solution is: (2,8)
-- Iteration: 4
-- Current solution is: (3,7)
-- Iteration: 5
-- Current solution is: (4,6)
-- Iteration: 6
-- Current solution is: (5,5)
-- Iteration: 7
-- Current solution is: (6,4)
-- Iteration: 8
-- Current solution is: (7,3)
-- Iteration: 9
-- Current solution is: (8,2)
-- Iteration: 10
-- Current solution is: (9,1)
-- Iteration: 11
-- Current solution is: (10,0)
-- Iteration: 12
-- No other solution!
-- [(0,10),(1,9),(2,8),(3,7),(4,6),(5,5),(6,4),(7,3),(8,2),(9,1),(10,0)]
demo :: IO ()
demo :: IO ()
demo = [(Integer, Integer)] -> IO ()
forall a. Show a => a -> IO ()
print ([(Integer, Integer)] -> IO ()) -> IO [(Integer, Integer)] -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Symbolic [(Integer, Integer)] -> IO [(Integer, Integer)]
forall a. Symbolic a -> IO a
runSMT Symbolic [(Integer, Integer)]
goodSum