-----------------------------------------------------------------------------
-- |
-- Module    : Documentation.SBV.Examples.Puzzles.Jugs
-- Copyright : (c) Levent Erkok
-- License   : BSD3
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- Solves the classic water jug puzzle: We have 3 jugs. The capacity of the jugs are 8, 5,
-- and 3 gallons. We begin with the 8 gallon jug full, the other two empty. We can transfer
-- from any jug to any other, completely topping off the latter. We want to end with
-- 4 gallons of water in the first and second jugs, and with an empty third jug. What
-- moves should we execute in order to do so?
-----------------------------------------------------------------------------

{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric  #-}

{-# OPTIONS_GHC -Wall -Werror #-}

module Documentation.SBV.Examples.Puzzles.Jugs where

import Data.SBV
import Data.SBV.Control

import GHC.Generics(Generic)

-- | A Jug has a capacity (i.e., maximum amount of water it can hold), and content, showing how much
-- it currently has. The invariant is that content is always non-negative and is at most the capacity.
data Jug = Jug { Jug -> Integer
capacity :: Integer
               , Jug -> SInteger
content  :: SInteger
               } deriving ((forall x. Jug -> Rep Jug x)
-> (forall x. Rep Jug x -> Jug) -> Generic Jug
forall x. Rep Jug x -> Jug
forall x. Jug -> Rep Jug x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Jug x -> Jug
$cfrom :: forall x. Jug -> Rep Jug x
Generic, Bool -> SBool -> Jug -> Jug -> Jug
(Bool -> SBool -> Jug -> Jug -> Jug)
-> (forall b.
    (Ord b, SymVal b, Num b) =>
    [Jug] -> Jug -> SBV b -> Jug)
-> Mergeable Jug
forall b. (Ord b, SymVal b, Num b) => [Jug] -> Jug -> SBV b -> Jug
forall a.
(Bool -> SBool -> a -> a -> a)
-> (forall b. (Ord b, SymVal b, Num b) => [a] -> a -> SBV b -> a)
-> Mergeable a
select :: [Jug] -> Jug -> SBV b -> Jug
$cselect :: forall b. (Ord b, SymVal b, Num b) => [Jug] -> Jug -> SBV b -> Jug
symbolicMerge :: Bool -> SBool -> Jug -> Jug -> Jug
$csymbolicMerge :: Bool -> SBool -> Jug -> Jug -> Jug
Mergeable)

-- | Transfer from one jug to another. By definition,
-- we transfer to fill the second jug, which may end up
-- filling it fully, or leaving some in the first jug.
transfer :: Jug -> Jug -> (Jug, Jug)
transfer :: Jug -> Jug -> (Jug, Jug)
transfer Jug
j1 Jug
j2 = (Jug
j1', Jug
j2')
  where empty :: SInteger
empty         = Integer -> SInteger
forall a. SymVal a => a -> SBV a
literal (Jug -> Integer
capacity Jug
j2) SInteger -> SInteger -> SInteger
forall a. Num a => a -> a -> a
- Jug -> SInteger
content Jug
j2
        transferrable :: SInteger
transferrable = SInteger
empty SInteger -> SInteger -> SInteger
forall a. OrdSymbolic a => a -> a -> a
`smin` Jug -> SInteger
content Jug
j1
        j1' :: Jug
j1'           = Jug
j1 { content :: SInteger
content = Jug -> SInteger
content Jug
j1 SInteger -> SInteger -> SInteger
forall a. Num a => a -> a -> a
- SInteger
transferrable }
        j2' :: Jug
j2'           = Jug
j2 { content :: SInteger
content = Jug -> SInteger
content Jug
j2 SInteger -> SInteger -> SInteger
forall a. Num a => a -> a -> a
+ SInteger
transferrable }

-- | At the beginning, we have an full 8-gallon jug, and two empty jugs, 5 and 3 gallons each.
initJugs :: (Jug, Jug, Jug)
initJugs :: (Jug, Jug, Jug)
initJugs = (Jug
j1, Jug
j2, Jug
j3)
  where j1 :: Jug
j1 = Integer -> SInteger -> Jug
Jug Integer
8 SInteger
8
        j2 :: Jug
j2 = Integer -> SInteger -> Jug
Jug Integer
5 SInteger
0
        j3 :: Jug
j3 = Integer -> SInteger -> Jug
Jug Integer
3 SInteger
0

-- | We've solved the puzzle if 8 and 5 gallon jugs have 4 gallons each, and the third one is empty.
solved :: (Jug, Jug, Jug) -> SBool
solved :: (Jug, Jug, Jug) -> SBool
solved (Jug
j1, Jug
j2, Jug
j3) = Jug -> SInteger
content Jug
j1 SInteger -> SInteger -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== SInteger
4 SBool -> SBool -> SBool
.&& Jug -> SInteger
content Jug
j2 SInteger -> SInteger -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== SInteger
4 SBool -> SBool -> SBool
.&& Jug -> SInteger
content Jug
j3 SInteger -> SInteger -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== SInteger
0

-- | Execute a bunch of moves.
moves :: [(SInteger, SInteger)] -> (Jug, Jug, Jug)
moves :: [(SInteger, SInteger)] -> (Jug, Jug, Jug)
moves = ((Jug, Jug, Jug) -> (SInteger, SInteger) -> (Jug, Jug, Jug))
-> (Jug, Jug, Jug) -> [(SInteger, SInteger)] -> (Jug, Jug, Jug)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Jug, Jug, Jug) -> (SInteger, SInteger) -> (Jug, Jug, Jug)
move (Jug, Jug, Jug)
initJugs
  where move :: (Jug, Jug, Jug) -> (SInteger, SInteger) -> (Jug, Jug, Jug)
        move :: (Jug, Jug, Jug) -> (SInteger, SInteger) -> (Jug, Jug, Jug)
move (Jug
j0, Jug
j1, Jug
j2) (SInteger
from, SInteger
to) =
              SBool -> (Jug, Jug, Jug) -> (Jug, Jug, Jug) -> (Jug, Jug, Jug)
forall a. Mergeable a => SBool -> a -> a -> a
ite ((SInteger
from, SInteger
to) (SInteger, SInteger) -> (SInteger, SInteger) -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== (SInteger
1, SInteger
2)) (let (Jug
j0', Jug
j1') = Jug -> Jug -> (Jug, Jug)
transfer Jug
j0 Jug
j1 in (Jug
j0', Jug
j1', Jug
j2))
            ((Jug, Jug, Jug) -> (Jug, Jug, Jug))
-> (Jug, Jug, Jug) -> (Jug, Jug, Jug)
forall a b. (a -> b) -> a -> b
$ SBool -> (Jug, Jug, Jug) -> (Jug, Jug, Jug) -> (Jug, Jug, Jug)
forall a. Mergeable a => SBool -> a -> a -> a
ite ((SInteger
from, SInteger
to) (SInteger, SInteger) -> (SInteger, SInteger) -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== (SInteger
2, SInteger
1)) (let (Jug
j1', Jug
j0') = Jug -> Jug -> (Jug, Jug)
transfer Jug
j1 Jug
j0 in (Jug
j0', Jug
j1', Jug
j2))
            ((Jug, Jug, Jug) -> (Jug, Jug, Jug))
-> (Jug, Jug, Jug) -> (Jug, Jug, Jug)
forall a b. (a -> b) -> a -> b
$ SBool -> (Jug, Jug, Jug) -> (Jug, Jug, Jug) -> (Jug, Jug, Jug)
forall a. Mergeable a => SBool -> a -> a -> a
ite ((SInteger
from, SInteger
to) (SInteger, SInteger) -> (SInteger, SInteger) -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== (SInteger
1, SInteger
3)) (let (Jug
j0', Jug
j2') = Jug -> Jug -> (Jug, Jug)
transfer Jug
j0 Jug
j2 in (Jug
j0', Jug
j1,  Jug
j2'))
            ((Jug, Jug, Jug) -> (Jug, Jug, Jug))
-> (Jug, Jug, Jug) -> (Jug, Jug, Jug)
forall a b. (a -> b) -> a -> b
$ SBool -> (Jug, Jug, Jug) -> (Jug, Jug, Jug) -> (Jug, Jug, Jug)
forall a. Mergeable a => SBool -> a -> a -> a
ite ((SInteger
from, SInteger
to) (SInteger, SInteger) -> (SInteger, SInteger) -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== (SInteger
3, SInteger
1)) (let (Jug
j2', Jug
j0') = Jug -> Jug -> (Jug, Jug)
transfer Jug
j2 Jug
j0 in (Jug
j0', Jug
j1,  Jug
j2'))
            ((Jug, Jug, Jug) -> (Jug, Jug, Jug))
-> (Jug, Jug, Jug) -> (Jug, Jug, Jug)
forall a b. (a -> b) -> a -> b
$ SBool -> (Jug, Jug, Jug) -> (Jug, Jug, Jug) -> (Jug, Jug, Jug)
forall a. Mergeable a => SBool -> a -> a -> a
ite ((SInteger
from, SInteger
to) (SInteger, SInteger) -> (SInteger, SInteger) -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== (SInteger
2, SInteger
3)) (let (Jug
j1', Jug
j2') = Jug -> Jug -> (Jug, Jug)
transfer Jug
j1 Jug
j2 in (Jug
j0,  Jug
j1', Jug
j2'))
            ((Jug, Jug, Jug) -> (Jug, Jug, Jug))
-> (Jug, Jug, Jug) -> (Jug, Jug, Jug)
forall a b. (a -> b) -> a -> b
$ SBool -> (Jug, Jug, Jug) -> (Jug, Jug, Jug) -> (Jug, Jug, Jug)
forall a. Mergeable a => SBool -> a -> a -> a
ite ((SInteger
from, SInteger
to) (SInteger, SInteger) -> (SInteger, SInteger) -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== (SInteger
3, SInteger
2)) (let (Jug
j2', Jug
j1') = Jug -> Jug -> (Jug, Jug)
transfer Jug
j2 Jug
j1 in (Jug
j0,  Jug
j1', Jug
j2'))
                                          (Jug
j0, Jug
j1, Jug
j2)

-- | Solve the puzzle. We have:
--
-- >>> puzzle
-- # of moves: 0
-- # of moves: 1
-- # of moves: 2
-- # of moves: 3
-- # of moves: 4
-- # of moves: 5
-- # of moves: 6
-- # of moves: 7
-- 1 --> 2
-- 2 --> 3
-- 3 --> 1
-- 2 --> 3
-- 1 --> 2
-- 2 --> 3
-- 3 --> 1
--
-- Here's the contents in terms of gallons after each move:
-- (8, 0, 0)
-- (3, 5, 0)
-- (3, 2, 3)
-- (6, 2, 0)
-- (6, 0, 2)
-- (1, 5, 2)
-- (1, 4, 3)
-- (4, 4, 0)
--
-- Note that by construction this is the minimum length solution. (Though our construction
-- does not guarantee that it is unique.)
puzzle :: IO ()
puzzle :: IO ()
puzzle = Symbolic () -> IO ()
forall a. Symbolic a -> IO a
runSMT (Symbolic () -> IO ()) -> Symbolic () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            let run :: Int -> QueryT IO ()
run Int
i = do IO () -> QueryT IO ()
forall a. IO a -> Query a
io (IO () -> QueryT IO ()) -> IO () -> QueryT IO ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"# of moves: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Int
i :: Int)
                           Int -> QueryT IO ()
push Int
1
                           [(SInteger, SInteger)]
ms <- (Int -> QueryT IO (SInteger, SInteger))
-> [Int] -> QueryT IO [(SInteger, SInteger)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (QueryT IO (SInteger, SInteger)
-> Int -> QueryT IO (SInteger, SInteger)
forall a b. a -> b -> a
const QueryT IO (SInteger, SInteger)
genMove) [Int
1..Int
i]
                           SBool -> QueryT IO ()
forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain (SBool -> QueryT IO ()) -> SBool -> QueryT IO ()
forall a b. (a -> b) -> a -> b
$ (Jug, Jug, Jug) -> SBool
solved ((Jug, Jug, Jug) -> SBool) -> (Jug, Jug, Jug) -> SBool
forall a b. (a -> b) -> a -> b
$ [(SInteger, SInteger)] -> (Jug, Jug, Jug)
moves [(SInteger, SInteger)]
ms
                           CheckSatResult
cs <- Query CheckSatResult
checkSat
                           case CheckSatResult
cs of
                             CheckSatResult
Unsat -> do Int -> QueryT IO ()
pop Int
1
                                         Int -> QueryT IO ()
run (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                             CheckSatResult
Sat   -> ((SInteger, SInteger) -> QueryT IO ())
-> [(SInteger, SInteger)] -> QueryT IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (SInteger, SInteger) -> QueryT IO ()
forall a a.
(SymVal a, SymVal a, Show a, Show a) =>
(SBV a, SBV a) -> QueryT IO ()
sh [(SInteger, SInteger)]
ms
                             CheckSatResult
_     -> String -> QueryT IO ()
forall a. HasCallStack => String -> a
error (String -> QueryT IO ()) -> String -> QueryT IO ()
forall a b. (a -> b) -> a -> b
$ String
"Unexpected result: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ CheckSatResult -> String
forall a. Show a => a -> String
show CheckSatResult
cs
            QueryT IO () -> Symbolic ()
forall a. Query a -> Symbolic a
query (QueryT IO () -> Symbolic ()) -> QueryT IO () -> Symbolic ()
forall a b. (a -> b) -> a -> b
$ Int -> QueryT IO ()
run Int
0
  where genMove :: QueryT IO (SInteger, SInteger)
genMove = (,) (SInteger -> SInteger -> (SInteger, SInteger))
-> QueryT IO SInteger
-> QueryT IO (SInteger -> (SInteger, SInteger))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> QueryT IO SInteger
forall a. SymVal a => Query (SBV a)
freshVar_ QueryT IO (SInteger -> (SInteger, SInteger))
-> QueryT IO SInteger -> QueryT IO (SInteger, SInteger)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> QueryT IO SInteger
forall a. SymVal a => Query (SBV a)
freshVar_
        sh :: (SBV a, SBV a) -> QueryT IO ()
sh (SBV a
f, SBV a
t) = do a
from <- SBV a -> Query a
forall a. SymVal a => SBV a -> Query a
getValue SBV a
f
                       a
to   <- SBV a -> Query a
forall a. SymVal a => SBV a -> Query a
getValue SBV a
t
                       IO () -> QueryT IO ()
forall a. IO a -> Query a
io (IO () -> QueryT IO ()) -> IO () -> QueryT IO ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ a -> String
forall a. Show a => a -> String
show a
from String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" --> " String -> String -> String
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
to