{-# Language LambdaCase, BlockArguments, ImportQualifiedPost #-}
{-|
Module      : Main
Description : Day 24 solution
Copyright   : (c) Eric Mertens, 2021
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2021/day/24>

Solve an arbitrary program using an SMT solver.

-}
module Main (main) where

import Advent.Input (getInputLines)
import Advent.ReadS ( P(..), runP )
import Control.Applicative ( Alternative((<|>), empty) )
import Control.Monad ( foldM )
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe ( fromJust, isJust )
import Data.SBV

-- | >>> :main
-- 49917929934999
-- 11911316711816
main :: IO ()
IO ()
main =
 do [Op]
pgm <- (String -> Op) -> [String] -> [Op]
forall a b. (a -> b) -> [a] -> [b]
map (P Op -> String -> Op
forall a. P a -> String -> a
runP P Op
pOp) ([String] -> [Op]) -> IO [String] -> IO [Op]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> IO [String]
getInputLines Int
2021 Int
24
    String -> IO ()
putStrLn (String -> IO ()) -> IO String -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Opt -> [Op] -> IO String
findAnswer Opt
forall a. Metric a => String -> SBV a -> Symbolic ()
maximize [Op]
pgm
    String -> IO ()
putStrLn (String -> IO ()) -> IO String -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Opt -> [Op] -> IO String
findAnswer Opt
forall a. Metric a => String -> SBV a -> Symbolic ()
minimize [Op]
pgm

findAnswer :: Opt -> [Op] -> IO String
findAnswer :: Opt -> [Op] -> IO String
findAnswer Opt
opt [Op]
pgm =
 do LexicographicResult SMTResult
r <- OptimizeStyle -> Symbolic () -> IO OptimizeResult
forall a. Provable a => OptimizeStyle -> a -> IO OptimizeResult
optimize OptimizeStyle
Lexicographic (Opt -> [Op] -> Symbolic ()
runProgram Opt
opt [Op]
pgm)
    String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> IO String) -> String -> IO String
forall a b. (a -> b) -> a -> b
$
      (Maybe Int64 -> String) -> [Maybe Int64] -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Int64 -> String
forall a. Show a => a -> String
show (Int64 -> String)
-> (Maybe Int64 -> Int64) -> Maybe Int64 -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Int64 -> Int64
forall a. HasCallStack => Maybe a -> a
fromJust) ([Maybe Int64] -> String) -> [Maybe Int64] -> String
forall a b. (a -> b) -> a -> b
$
      (Maybe Int64 -> Bool) -> [Maybe Int64] -> [Maybe Int64]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile Maybe Int64 -> Bool
forall a. Maybe a -> Bool
isJust
      [String -> SMTResult -> Maybe Int64
forall b. SymVal b => String -> SMTResult -> Maybe b
forall a b. (Modelable a, SymVal b) => String -> a -> Maybe b
getModelValue (String
"in"String -> String -> String
forall a. [a] -> [a] -> [a]
++Int -> String
forall a. Show a => a -> String
show Int
i) SMTResult
r :: Maybe Int64 | Int
i <- [Int
1::Int ..]]

-- * Programs

-- | Program opcodes
data Op = Inp Reg | Mul Reg Term | Add Reg Term | Div Reg Term | Mod Reg Term | Eql Reg Term
  deriving Int -> Op -> String -> String
[Op] -> String -> String
Op -> String
(Int -> Op -> String -> String)
-> (Op -> String) -> ([Op] -> String -> String) -> Show Op
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> Op -> String -> String
showsPrec :: Int -> Op -> String -> String
$cshow :: Op -> String
show :: Op -> String
$cshowList :: [Op] -> String -> String
showList :: [Op] -> String -> String
Show

-- | Machine registers
data Reg = W | X | Y | Z
  deriving (Reg -> Reg -> Bool
(Reg -> Reg -> Bool) -> (Reg -> Reg -> Bool) -> Eq Reg
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Reg -> Reg -> Bool
== :: Reg -> Reg -> Bool
$c/= :: Reg -> Reg -> Bool
/= :: Reg -> Reg -> Bool
Eq, Eq Reg
Eq Reg
-> (Reg -> Reg -> Ordering)
-> (Reg -> Reg -> Bool)
-> (Reg -> Reg -> Bool)
-> (Reg -> Reg -> Bool)
-> (Reg -> Reg -> Bool)
-> (Reg -> Reg -> Reg)
-> (Reg -> Reg -> Reg)
-> Ord Reg
Reg -> Reg -> Bool
Reg -> Reg -> Ordering
Reg -> Reg -> Reg
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Reg -> Reg -> Ordering
compare :: Reg -> Reg -> Ordering
$c< :: Reg -> Reg -> Bool
< :: Reg -> Reg -> Bool
$c<= :: Reg -> Reg -> Bool
<= :: Reg -> Reg -> Bool
$c> :: Reg -> Reg -> Bool
> :: Reg -> Reg -> Bool
$c>= :: Reg -> Reg -> Bool
>= :: Reg -> Reg -> Bool
$cmax :: Reg -> Reg -> Reg
max :: Reg -> Reg -> Reg
$cmin :: Reg -> Reg -> Reg
min :: Reg -> Reg -> Reg
Ord, Int -> Reg -> String -> String
[Reg] -> String -> String
Reg -> String
(Int -> Reg -> String -> String)
-> (Reg -> String) -> ([Reg] -> String -> String) -> Show Reg
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> Reg -> String -> String
showsPrec :: Int -> Reg -> String -> String
$cshow :: Reg -> String
show :: Reg -> String
$cshowList :: [Reg] -> String -> String
showList :: [Reg] -> String -> String
Show)

-- | Opcode terms that can evaluate to an integer
data Term = Reg Reg | Int Int64
  deriving Int -> Term -> String -> String
[Term] -> String -> String
Term -> String
(Int -> Term -> String -> String)
-> (Term -> String) -> ([Term] -> String -> String) -> Show Term
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> Term -> String -> String
showsPrec :: Int -> Term -> String -> String
$cshow :: Term -> String
show :: Term -> String
$cshowList :: [Term] -> String -> String
showList :: [Term] -> String -> String
Show

-- * Interpretation

-- | Optimization function ('minimize' or 'maximize')
type Opt = String -> SBV Int64 -> Symbolic ()

-- | Map of values associated with each register
type Regs = Map Reg SInt64

-- | Evaluate a 'Term'
val :: Regs -> Term -> SInt64
val :: Regs -> Term -> SBV Int64
val Regs
regs (Reg Reg
r) = Regs
regs Regs -> Reg -> SBV Int64
forall k a. Ord k => Map k a -> k -> a
Map.! Reg
r
val Regs
_    (Int Int64
i) = Int64 -> SBV Int64
forall a. SymVal a => a -> SBV a
literal Int64
i

-- | Optimize the given program ensuring it ends with @z == 0@
runProgram :: Opt -> [Op] -> Symbolic ()
runProgram :: Opt -> [Op] -> Symbolic ()
runProgram Opt
opt [Op]
pgm =
 do let regs :: Regs
regs = [(Reg, SBV Int64)] -> Regs
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Reg
W,SBV Int64
0),(Reg
X,SBV Int64
0),(Reg
Y,SBV Int64
0),(Reg
Z,SBV Int64
0)]
    (Int
_, Regs
regs') <- ((Int, Regs) -> Op -> SymbolicT IO (Int, Regs))
-> (Int, Regs) -> [Op] -> SymbolicT IO (Int, Regs)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Opt -> (Int, Regs) -> Op -> SymbolicT IO (Int, Regs)
step Opt
opt) (Int
1, Regs
regs) [Op]
pgm
    String -> SBool -> Symbolic ()
forall (m :: * -> *). SolverContext m => String -> SBool -> m ()
namedConstraint String
"z==0" (Regs
regs' Regs -> Reg -> SBV Int64
forall k a. Ord k => Map k a -> k -> a
Map.! Reg
Z SBV Int64 -> SBV Int64 -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== SBV Int64
0)

step :: Opt -> (Int, Regs) -> Op -> Symbolic (Int, Regs)
step :: Opt -> (Int, Regs) -> Op -> SymbolicT IO (Int, Regs)
step Opt
opt (Int
n, Regs
regs) (Inp Reg
r) =
 do SBV Int64
i <- String -> Symbolic (SBV Int64)
forall a. SymVal a => String -> Symbolic (SBV a)
free (String
"in" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n)
    Opt
opt String
"opt" SBV Int64
i
    SBool -> Symbolic ()
forall (m :: * -> *). SolverContext m => SBool -> m ()
constrain (SBV Int64 -> (SBV Int64, SBV Int64) -> SBool
forall a. OrdSymbolic a => a -> (a, a) -> SBool
inRange SBV Int64
i (SBV Int64
1,SBV Int64
9))
    (Int, Regs) -> SymbolicT IO (Int, Regs)
forall a. a -> SymbolicT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1, Reg -> SBV Int64 -> Regs -> Regs
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Reg
r SBV Int64
i Regs
regs)
step Opt
_ (Int
n, Regs
regs) (Add Reg
x Term
y) = (Int, Regs) -> SymbolicT IO (Int, Regs)
forall a. a -> SymbolicT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
n, (SBV Int64 -> SBV Int64) -> Reg -> Regs -> Regs
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
Map.adjust (Regs -> Term -> SBV Int64
val Regs
regs Term
y SBV Int64 -> SBV Int64 -> SBV Int64
forall a. Num a => a -> a -> a
+)                     Reg
x Regs
regs)
step Opt
_ (Int
n, Regs
regs) (Mul Reg
x Term
y) = (Int, Regs) -> SymbolicT IO (Int, Regs)
forall a. a -> SymbolicT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
n, (SBV Int64 -> SBV Int64) -> Reg -> Regs -> Regs
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
Map.adjust (Regs -> Term -> SBV Int64
val Regs
regs Term
y SBV Int64 -> SBV Int64 -> SBV Int64
forall a. Num a => a -> a -> a
*)                     Reg
x Regs
regs)
step Opt
_ (Int
n, Regs
regs) (Div Reg
x Term
y) = (Int, Regs) -> SymbolicT IO (Int, Regs)
forall a. a -> SymbolicT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
n, (SBV Int64 -> SBV Int64) -> Reg -> Regs -> Regs
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
Map.adjust (SBV Int64 -> SBV Int64 -> SBV Int64
forall a. SDivisible a => a -> a -> a
`sDiv` Regs -> Term -> SBV Int64
val Regs
regs Term
y)                Reg
x Regs
regs)
step Opt
_ (Int
n, Regs
regs) (Mod Reg
x Term
y) = (Int, Regs) -> SymbolicT IO (Int, Regs)
forall a. a -> SymbolicT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
n, (SBV Int64 -> SBV Int64) -> Reg -> Regs -> Regs
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
Map.adjust (SBV Int64 -> SBV Int64 -> SBV Int64
forall a. SDivisible a => a -> a -> a
`sMod` Regs -> Term -> SBV Int64
val Regs
regs Term
y)                Reg
x Regs
regs)
step Opt
_ (Int
n, Regs
regs) (Eql Reg
x Term
y) = (Int, Regs) -> SymbolicT IO (Int, Regs)
forall a. a -> SymbolicT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
n, (SBV Int64 -> SBV Int64) -> Reg -> Regs -> Regs
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
Map.adjust (\SBV Int64
a -> SBool -> SBV Int64 -> SBV Int64 -> SBV Int64
forall a. Mergeable a => SBool -> a -> a -> a
ite (SBV Int64
a SBV Int64 -> SBV Int64 -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== Regs -> Term -> SBV Int64
val Regs
regs Term
y) SBV Int64
1 SBV Int64
0) Reg
x Regs
regs)

-- * Parsing

pOp :: P Op
pOp :: P Op
pOp = ReadS String -> P String
forall a. ReadS a -> P a
P ReadS String
lex P String -> (String -> P Op) -> P Op
forall a b. P a -> (a -> P b) -> P b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  String
"inp" -> Reg -> Op
Inp (Reg -> Op) -> P Reg -> P Op
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Reg
pReg
  String
"add" -> Reg -> Term -> Op
Add (Reg -> Term -> Op) -> P Reg -> P (Term -> Op)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Reg
pReg P (Term -> Op) -> P Term -> P Op
forall a b. P (a -> b) -> P a -> P b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> P Term
pTerm
  String
"mul" -> Reg -> Term -> Op
Mul (Reg -> Term -> Op) -> P Reg -> P (Term -> Op)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Reg
pReg P (Term -> Op) -> P Term -> P Op
forall a b. P (a -> b) -> P a -> P b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> P Term
pTerm
  String
"div" -> Reg -> Term -> Op
Div (Reg -> Term -> Op) -> P Reg -> P (Term -> Op)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Reg
pReg P (Term -> Op) -> P Term -> P Op
forall a b. P (a -> b) -> P a -> P b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> P Term
pTerm
  String
"mod" -> Reg -> Term -> Op
Mod (Reg -> Term -> Op) -> P Reg -> P (Term -> Op)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Reg
pReg P (Term -> Op) -> P Term -> P Op
forall a b. P (a -> b) -> P a -> P b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> P Term
pTerm
  String
"eql" -> Reg -> Term -> Op
Eql (Reg -> Term -> Op) -> P Reg -> P (Term -> Op)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Reg
pReg P (Term -> Op) -> P Term -> P Op
forall a b. P (a -> b) -> P a -> P b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> P Term
pTerm
  String
_     -> P Op
forall a. P a
forall (f :: * -> *) a. Alternative f => f a
empty

pReg :: P Reg
pReg :: P Reg
pReg = ReadS String -> P String
forall a. ReadS a -> P a
P ReadS String
lex P String -> (String -> P Reg) -> P Reg
forall a b. P a -> (a -> P b) -> P b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  String
"w" -> Reg -> P Reg
forall a. a -> P a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Reg
W
  String
"x" -> Reg -> P Reg
forall a. a -> P a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Reg
X
  String
"y" -> Reg -> P Reg
forall a. a -> P a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Reg
Y
  String
"z" -> Reg -> P Reg
forall a. a -> P a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Reg
Z
  String
_   -> P Reg
forall a. P a
forall (f :: * -> *) a. Alternative f => f a
empty

pTerm :: P Term
pTerm :: P Term
pTerm = Int64 -> Term
Int (Int64 -> Term) -> P Int64 -> P Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ReadS Int64 -> P Int64
forall a. ReadS a -> P a
P ReadS Int64
forall a. Read a => ReadS a
reads P Term -> P Term -> P Term
forall a. P a -> P a -> P a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Reg -> Term
Reg (Reg -> Term) -> P Reg -> P Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Reg
pReg