-----------------------------------------------------------------------------
-- |
-- Module    : Documentation.SBV.Examples.Queries.FourFours
-- Copyright : (c) Levent Erkok
-- License   : BSD3
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- A query based solution to the four-fours puzzle.
-- Inspired by <http://www.gigamonkeys.com/trees/>
--
-- @
-- Try to make every number between 0 and 20 using only four 4s and any
-- mathematical operation, with all four 4s being used each time.
-- @
--
-- We pretty much follow the structure of <http://www.gigamonkeys.com/trees/>,
-- with the exception that we generate the trees filled with symbolic operators
-- and ask the SMT solver to find the appropriate fillings.
-----------------------------------------------------------------------------

{-# LANGUAGE DeriveAnyClass      #-}
{-# LANGUAGE DeriveDataTypeable  #-}
{-# LANGUAGE FlexibleInstances  #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE TemplateHaskell     #-}

{-# OPTIONS_GHC -Wall -Werror #-}

module Documentation.SBV.Examples.Queries.FourFours where

import Data.SBV
import Data.SBV.Control

import Data.List (inits, tails)
import Data.Maybe

-- | Supported binary operators. To keep the search-space small, we will only allow division by @2@ or @4@,
-- and exponentiation will only be to the power @0@. This does restrict the search space, but is sufficient to
-- solve all the instances.
data BinOp = Plus | Minus | Times | Divide | Expt

-- | Make 'BinOp' a symbolic value.
mkSymbolicEnumeration ''BinOp

-- | Supported unary operators. Similar to 'BinOp' case, we will restrict square-root and factorial to
-- be only applied to the value @4.
data UnOp  = Negate | Sqrt | Factorial

-- | Make 'UnOp' a symbolic value.
mkSymbolicEnumeration ''UnOp

-- | The shape of a tree, either a binary node, or a unary node, or the number @4@, represented hear by
-- the constructor @F@. We parameterize by the operator type: When doing symbolic computations, we'll fill
-- those with 'SBinOp' and 'SUnOp'. When finding the shapes, we will simply put unit values, i.e., holes.
data T b u = B b (T b u) (T b u)
           | U u (T b u)
           | F

-- | A rudimentary 'Show' instance for trees, nothing fancy.
instance Show (T BinOp UnOp) where
   show :: T BinOp UnOp -> String
show T BinOp UnOp
F       = String
"4"
   show (U UnOp
u T BinOp UnOp
t) = case UnOp
u of
                    UnOp
Negate    -> String
"-" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show T BinOp UnOp
t
                    UnOp
Sqrt      -> String
"sqrt(" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show T BinOp UnOp
t forall a. [a] -> [a] -> [a]
++ String
")"
                    UnOp
Factorial -> forall a. Show a => a -> String
show T BinOp UnOp
t forall a. [a] -> [a] -> [a]
++ String
"!"
   show (B BinOp
o T BinOp UnOp
l T BinOp UnOp
r) = String
"(" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show T BinOp UnOp
l forall a. [a] -> [a] -> [a]
++ String
" " forall a. [a] -> [a] -> [a]
++ String
so forall a. [a] -> [a] -> [a]
++ String
" " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show T BinOp UnOp
r forall a. [a] -> [a] -> [a]
++ String
")"
     where so :: String
so = forall a. a -> Maybe a -> a
fromMaybe (forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"Unexpected operator: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show BinOp
o)
                        forall a b. (a -> b) -> a -> b
$ BinOp
o forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(BinOp
Plus, String
"+"), (BinOp
Minus, String
"-"), (BinOp
Times, String
"*"), (BinOp
Divide, String
"/"), (BinOp
Expt, String
"^")]

-- | Construct all possible tree shapes. The argument here follows the logic in <http://www.gigamonkeys.com/trees/>:
-- We simply construct all possible shapes and extend with the operators. The number of such trees is:
--
-- >>> length allPossibleTrees
-- 640
--
-- Note that this is a /lot/ smaller than what is generated by <http://www.gigamonkeys.com/trees/>. (There, the
-- number of trees is 10240000: 16000 times more than what we have to consider!)
allPossibleTrees :: [T () ()]
allPossibleTrees :: [T () ()]
allPossibleTrees = [T () ()] -> [T () ()]
trees forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> [a]
replicate Int
4 forall b u. T b u
F
  where trees :: [T () ()] -> [T () ()]
        trees :: [T () ()] -> [T () ()]
trees [T () ()
x] = [T () ()
x, forall b u. u -> T b u -> T b u
U () T () ()
x]
        trees [T () ()]
xs  = do ([T () ()]
left, [T () ()]
right) <- [([T () ()], [T () ()])]
splits
                       T () ()
t1            <- [T () ()] -> [T () ()]
trees [T () ()]
left
                       T () ()
t2            <- [T () ()] -> [T () ()]
trees [T () ()]
right
                       [T () ()] -> [T () ()]
trees [forall b u. b -> T b u -> T b u -> T b u
B () T () ()
t1 T () ()
t2]
          where splits :: [([T () ()], [T () ()])]
splits = forall a. [a] -> [a]
init forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [a]
tail forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall a. [a] -> [[a]]
inits [T () ()]
xs) (forall a. [a] -> [[a]]
tails [T () ()]
xs)

-- | Given a tree with hols, fill it with symbolic operators. This is the /trick/ that allows
-- us to consider only 640 trees as opposed to over 10 million.
fill :: T () () -> Symbolic (T SBinOp SUnOp)
fill :: T () () -> Symbolic (T SBinOp SUnOp)
fill (B ()
_ T () ()
l T () ()
r) = forall b u. b -> T b u -> T b u -> T b u
B forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. SymVal a => Symbolic (SBV a)
free_ forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> T () () -> Symbolic (T SBinOp SUnOp)
fill T () ()
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> T () () -> Symbolic (T SBinOp SUnOp)
fill T () ()
r
fill (U ()
_ T () ()
t)   = forall b u. u -> T b u -> T b u
U forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. SymVal a => Symbolic (SBV a)
free_ forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> T () () -> Symbolic (T SBinOp SUnOp)
fill T () ()
t
fill T () ()
F         = forall (m :: * -> *) a. Monad m => a -> m a
return forall b u. T b u
F

-- | Minor helper for writing "symbolic" case statements. Simply walks down a list
-- of values to match against a symbolic version of the key.
sCase :: (Eq a, SymVal a, Mergeable v) => SBV a -> [(a, v)] -> v
sCase :: forall a v. (Eq a, SymVal a, Mergeable v) => SBV a -> [(a, v)] -> v
sCase SBV a
k = forall {a}. Mergeable a => [(a, a)] -> a
walk
  where walk :: [(a, a)] -> a
walk []              = forall a. HasCallStack => String -> a
error String
"sCase: Expected a non-empty list of cases!"
        walk [(a
_, a
v)]        = a
v
        walk ((a
k1, a
v1):[(a, a)]
rest) = forall a. Mergeable a => SBool -> a -> a -> a
ite (SBV a
k forall a. EqSymbolic a => a -> a -> SBool
.== forall a. SymVal a => a -> SBV a
literal a
k1) a
v1 ([(a, a)] -> a
walk [(a, a)]
rest)

-- | Evaluate a symbolic tree, obtaining a symbolic value. Note how we structure
-- this evaluation so we impose extra constraints on what values square-root, divide
-- etc. can take. This is the power of the symbolic approach: We can put arbitrary
-- symbolic constraints as we evaluate the tree.
eval :: T SBinOp SUnOp -> Symbolic SInteger
eval :: T SBinOp SUnOp -> Symbolic SInteger
eval T SBinOp SUnOp
tree = case T SBinOp SUnOp
tree of
              B SBinOp
b T SBinOp SUnOp
l T SBinOp SUnOp
r -> T SBinOp SUnOp -> Symbolic SInteger
eval T SBinOp SUnOp
l forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \SInteger
l' -> T SBinOp SUnOp -> Symbolic SInteger
eval T SBinOp SUnOp
r forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \SInteger
r' -> SBinOp -> SInteger -> SInteger -> Symbolic SInteger
binOp SBinOp
b SInteger
l' SInteger
r'
              U SUnOp
u T SBinOp SUnOp
t   -> T SBinOp SUnOp -> Symbolic SInteger
eval T SBinOp SUnOp
t forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SUnOp -> SInteger -> Symbolic SInteger
uOp SUnOp
u
              T SBinOp SUnOp
F       -> forall (m :: * -> *) a. Monad m => a -> m a
return SInteger
4

  where binOp :: SBinOp -> SInteger -> SInteger -> Symbolic SInteger
        binOp :: SBinOp -> SInteger -> SInteger -> Symbolic SInteger
binOp SBinOp
o SInteger
l SInteger
r = do forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain forall a b. (a -> b) -> a -> b
$ SBinOp
o forall a. EqSymbolic a => a -> a -> SBool
.== SBinOp
sDivide SBool -> SBool -> SBool
.=> SInteger
r forall a. EqSymbolic a => a -> a -> SBool
.== SInteger
4 SBool -> SBool -> SBool
.|| SInteger
r forall a. EqSymbolic a => a -> a -> SBool
.== SInteger
2
                         forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain forall a b. (a -> b) -> a -> b
$ SBinOp
o forall a. EqSymbolic a => a -> a -> SBool
.== SBinOp
sExpt   SBool -> SBool -> SBool
.=> SInteger
r forall a. EqSymbolic a => a -> a -> SBool
.== SInteger
0
                         forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a v. (Eq a, SymVal a, Mergeable v) => SBV a -> [(a, v)] -> v
sCase SBinOp
o
                                    [ (BinOp
Plus,    SInteger
lforall a. Num a => a -> a -> a
+SInteger
r)
                                    , (BinOp
Minus,   SInteger
lforall a. Num a => a -> a -> a
-SInteger
r)
                                    , (BinOp
Times,   SInteger
lforall a. Num a => a -> a -> a
*SInteger
r)
                                    , (BinOp
Divide,  SInteger
l forall a. SDivisible a => a -> a -> a
`sDiv` SInteger
r)
                                    , (BinOp
Expt,    SInteger
1)   -- exponent is restricted to 0, so the value is 1
                                    ]

        uOp :: SUnOp -> SInteger -> Symbolic SInteger
        uOp :: SUnOp -> SInteger -> Symbolic SInteger
uOp SUnOp
o SInteger
v = do forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain forall a b. (a -> b) -> a -> b
$ SUnOp
o forall a. EqSymbolic a => a -> a -> SBool
.== SUnOp
sSqrt      SBool -> SBool -> SBool
.=> SInteger
v forall a. EqSymbolic a => a -> a -> SBool
.== SInteger
4
                     forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain forall a b. (a -> b) -> a -> b
$ SUnOp
o forall a. EqSymbolic a => a -> a -> SBool
.== SUnOp
sFactorial SBool -> SBool -> SBool
.=> SInteger
v forall a. EqSymbolic a => a -> a -> SBool
.== SInteger
4
                     forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a v. (Eq a, SymVal a, Mergeable v) => SBV a -> [(a, v)] -> v
sCase SUnOp
o
                                [ (UnOp
Negate,    -SInteger
v)
                                , (UnOp
Sqrt,       SInteger
2)  -- argument is restricted to 4, so the value is 2
                                , (UnOp
Factorial, SInteger
24)  -- argument is restricted to 4, so the value is 24
                                ]

-- | In the query mode, find a filling of a given tree shape /t/, such that it evaluates to the
-- requested number /i/. Note that we return back a concrete tree.
generate :: Integer -> T () () -> IO (Maybe (T BinOp UnOp))
generate :: Integer -> T () () -> IO (Maybe (T BinOp UnOp))
generate Integer
i T () ()
t = forall a. Symbolic a -> IO a
runSMT forall a b. (a -> b) -> a -> b
$ do T SBinOp SUnOp
symT <- T () () -> Symbolic (T SBinOp SUnOp)
fill T () ()
t
                           SInteger
val  <- T SBinOp SUnOp -> Symbolic SInteger
eval T SBinOp SUnOp
symT
                           forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain forall a b. (a -> b) -> a -> b
$ SInteger
val forall a. EqSymbolic a => a -> a -> SBool
.== forall a. SymVal a => a -> SBV a
literal Integer
i
                           forall a. Query a -> Symbolic a
query forall a b. (a -> b) -> a -> b
$ do CheckSatResult
cs <- Query CheckSatResult
checkSat
                                      case CheckSatResult
cs of
                                        CheckSatResult
Sat -> forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {u} {b}.
(SymVal u, SymVal b) =>
T (SBV b) (SBV u) -> QueryT IO (T b u)
construct T SBinOp SUnOp
symT
                                        CheckSatResult
_   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
    where -- Walk through the tree, ask the solver for
          -- the assignment to symbolic operators and fill back.
          construct :: T (SBV b) (SBV u) -> QueryT IO (T b u)
construct T (SBV b) (SBV u)
F           = forall (m :: * -> *) a. Monad m => a -> m a
return forall b u. T b u
F
          construct (U SBV u
o T (SBV b) (SBV u)
s')    = do u
uo <- forall a. SymVal a => SBV a -> Query a
getValue SBV u
o
                                     forall b u. u -> T b u -> T b u
U u
uo forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> T (SBV b) (SBV u) -> QueryT IO (T b u)
construct T (SBV b) (SBV u)
s'
          construct (B SBV b
b T (SBV b) (SBV u)
l' T (SBV b) (SBV u)
r') = do b
bo <- forall a. SymVal a => SBV a -> Query a
getValue SBV b
b
                                     forall b u. b -> T b u -> T b u -> T b u
B b
bo forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> T (SBV b) (SBV u) -> QueryT IO (T b u)
construct T (SBV b) (SBV u)
l' forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> T (SBV b) (SBV u) -> QueryT IO (T b u)
construct T (SBV b) (SBV u)
r'

-- | Given an integer, walk through all possible tree shapes (at most 640 of them), and find a
-- filling that solves the puzzle.
find :: Integer -> IO ()
find :: Integer -> IO ()
find Integer
target = [T () ()] -> IO ()
go [T () ()]
allPossibleTrees
  where go :: [T () ()] -> IO ()
go []     = String -> IO ()
putStrLn forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show Integer
target forall a. [a] -> [a] -> [a]
++ String
": No solution found."
        go (T () ()
t:[T () ()]
ts) = do Maybe (T BinOp UnOp)
chk <- Integer -> T () () -> IO (Maybe (T BinOp UnOp))
generate Integer
target T () ()
t
                       case Maybe (T BinOp UnOp)
chk of
                         Maybe (T BinOp UnOp)
Nothing -> [T () ()] -> IO ()
go [T () ()]
ts
                         Just T BinOp UnOp
r  -> do let ok :: Bool
ok  = T BinOp UnOp -> Integer
concEval T BinOp UnOp
r forall a. Eq a => a -> a -> Bool
== Integer
target
                                           tag :: String
tag = if Bool
ok then String
" [OK]: " else String
" [BAD]: "
                                           sh :: a -> String
sh a
i | a
i forall a. Ord a => a -> a -> Bool
< a
10 = Char
' ' forall a. a -> [a] -> [a]
: forall a. Show a => a -> String
show a
i
                                                | Bool
True   =       forall a. Show a => a -> String
show a
i

                                       String -> IO ()
putStrLn forall a b. (a -> b) -> a -> b
$ forall {a}. (Ord a, Num a, Show a) => a -> String
sh Integer
target forall a. [a] -> [a] -> [a]
++ String
tag forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show T BinOp UnOp
r

        -- Make sure the result is correct!
        concEval :: T BinOp UnOp -> Integer
        concEval :: T BinOp UnOp -> Integer
concEval T BinOp UnOp
F         = Integer
4
        concEval (U UnOp
u T BinOp UnOp
t)   = UnOp -> Integer -> Integer
uEval UnOp
u (T BinOp UnOp -> Integer
concEval T BinOp UnOp
t)
        concEval (B BinOp
b T BinOp UnOp
l T BinOp UnOp
r) = BinOp -> Integer -> Integer -> Integer
bEval BinOp
b (T BinOp UnOp -> Integer
concEval T BinOp UnOp
l) (T BinOp UnOp -> Integer
concEval T BinOp UnOp
r)

        uEval :: UnOp -> Integer -> Integer
        uEval :: UnOp -> Integer -> Integer
uEval UnOp
Negate    Integer
i = -Integer
i
        uEval UnOp
Sqrt      Integer
i = if Integer
i forall a. Eq a => a -> a -> Bool
== Integer
4 then  Integer
2 else forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"uEval: Found sqrt applied to value: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Integer
i
        uEval UnOp
Factorial Integer
i = if Integer
i forall a. Eq a => a -> a -> Bool
== Integer
4 then Integer
24 else forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"uEval: Found factorial applied to value: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Integer
i

        bEval :: BinOp -> Integer -> Integer -> Integer
        bEval :: BinOp -> Integer -> Integer -> Integer
bEval BinOp
Plus   Integer
i Integer
j = Integer
i forall a. Num a => a -> a -> a
+ Integer
j
        bEval BinOp
Minus  Integer
i Integer
j = Integer
i forall a. Num a => a -> a -> a
- Integer
j
        bEval BinOp
Times  Integer
i Integer
j = Integer
i forall a. Num a => a -> a -> a
* Integer
j
        bEval BinOp
Divide Integer
i Integer
j = Integer
i forall a. Integral a => a -> a -> a
`div` Integer
j
        bEval BinOp
Expt   Integer
i Integer
j = Integer
i forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
j

-- | Solution to the puzzle. When you run this puzzle, the solver can produce different results
-- than what's shown here, but the expressions should still be all valid!
--
-- @
-- ghci> puzzle
--  0 [OK]: (4 - (4 + (4 - 4)))
--  1 [OK]: (4 / (4 + (4 - 4)))
--  2 [OK]: sqrt((4 + (4 * (4 - 4))))
--  3 [OK]: (4 - (4 ^ (4 - 4)))
--  4 [OK]: (4 + (4 * (4 - 4)))
--  5 [OK]: (4 + (4 ^ (4 - 4)))
--  6 [OK]: (4 + sqrt((4 * (4 / 4))))
--  7 [OK]: (4 + (4 - (4 / 4)))
--  8 [OK]: (4 - (4 - (4 + 4)))
--  9 [OK]: (4 + (4 + (4 / 4)))
-- 10 [OK]: (4 + (4 + (4 - sqrt(4))))
-- 11 [OK]: (4 + ((4 + 4!) / 4))
-- 12 [OK]: (4 * (4 - (4 / 4)))
-- 13 [OK]: (4! + ((sqrt(4) - 4!) / sqrt(4)))
-- 14 [OK]: (4 + (4 + (4 + sqrt(4))))
-- 15 [OK]: (4 + ((4! - sqrt(4)) / sqrt(4)))
-- 16 [OK]: (4 * (4 * (4 / 4)))
-- 17 [OK]: (4 + ((sqrt(4) + 4!) / sqrt(4)))
-- 18 [OK]: -(4 + (4 - (sqrt(4) + 4!)))
-- 19 [OK]: -(4 - (4! - (4 / 4)))
-- 20 [OK]: (4 * (4 + (4 / 4)))
-- @
puzzle :: IO ()
puzzle :: IO ()
puzzle = forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Integer -> IO ()
find [Integer
0 .. Integer
20]