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

<https://adventofcode.com/2016/day/25>

Assembly that can output binary numbers. We run them
until we can establish that the state is looping while
producing the desired output sequence.

-}
module Main where

import Advent.Input ( getInputLines )
import Advent.ReadS ( P(..), runP )
import AsmProg
import Control.Applicative ( Alternative(empty) )
import Control.Lens (use, (+=), (-=), (.=), (<~), makeLenses, Contains(contains))
import Control.Monad.Trans.State ( evalState )
import Data.List ( find )
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Vector (Vector)
import Data.Vector qualified as Vector

-- | Expected next output
data Progress = NeedOne | NeedZero

-- | State of the interpreter
data Machine = Machine
  { Machine -> Registers
_machRegisters :: !Registers
  , Machine -> Progress
_machProgress  :: !Progress
  , Machine -> Set (Int, Registers)
_machTargets   :: !(Set (Int,Registers))
  }

makeLenses ''Machine

instance HasRegisters Machine where
  reg :: forall (f :: * -> *).
Functor f =>
Register -> LensLike' f Machine Int
reg Register
r = (Registers -> f Registers) -> Machine -> f Machine
Lens' Machine Registers
machRegisters ((Registers -> f Registers) -> Machine -> f Machine)
-> ((Int -> f Int) -> Registers -> f Registers)
-> (Int -> f Int)
-> Machine
-> f Machine
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Register -> (Int -> f Int) -> Registers -> f Registers
forall a (f :: * -> *).
(HasRegisters a, Functor f) =>
Register -> LensLike' f a Int
forall (f :: * -> *).
Functor f =>
Register -> LensLike' f Registers Int
reg Register
r
  {-# INLINE reg #-}

-- | >>> :main
-- Just 192
main :: IO ()
IO ()
main =
 do Vector Inst
program <- [Inst] -> Vector Inst
forall a. [a] -> Vector a
Vector.fromList ([Inst] -> Vector Inst)
-> ([String] -> [Inst]) -> [String] -> Vector Inst
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String -> Inst) -> [String] -> [Inst]
forall a b. (a -> b) -> [a] -> [b]
map (P Inst -> String -> Inst
forall a. P a -> String -> a
runP P Inst
pInst) ([String] -> Vector Inst) -> IO [String] -> IO (Vector Inst)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> IO [String]
getInputLines Int
2016 Int
25
    Maybe Int -> IO ()
forall a. Show a => a -> IO ()
print (Maybe Int -> IO ()) -> Maybe Int -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> Maybe Int
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (Vector Inst -> Int -> Bool
execute Vector Inst
program) [Int
1..]

data Inst
  = Copy Value Register
  | Inc Register
  | Dec Register
  | Jnz Value Value
  | Out Value
 deriving Int -> Inst -> ShowS
[Inst] -> ShowS
Inst -> String
(Int -> Inst -> ShowS)
-> (Inst -> String) -> ([Inst] -> ShowS) -> Show Inst
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Inst -> ShowS
showsPrec :: Int -> Inst -> ShowS
$cshow :: Inst -> String
show :: Inst -> String
$cshowList :: [Inst] -> ShowS
showList :: [Inst] -> ShowS
Show

pInst :: P Inst
pInst :: P Inst
pInst = ReadS String -> P String
forall a. ReadS a -> P a
P ReadS String
lex P String -> (String -> P Inst) -> P Inst
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
"cpy" -> Value -> Register -> Inst
Copy (Value -> Register -> Inst) -> P Value -> P (Register -> Inst)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Value
pValue P (Register -> Inst) -> P Register -> P Inst
forall a b. P (a -> b) -> P a -> P b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> P Register
pReg
  String
"jnz" -> Value -> Value -> Inst
Jnz  (Value -> Value -> Inst) -> P Value -> P (Value -> Inst)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Value
pValue P (Value -> Inst) -> P Value -> P Inst
forall a b. P (a -> b) -> P a -> P b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> P Value
pValue
  String
"inc" -> Register -> Inst
Inc  (Register -> Inst) -> P Register -> P Inst
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Register
pReg
  String
"dec" -> Register -> Inst
Dec  (Register -> Inst) -> P Register -> P Inst
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Register
pReg
  String
"out" -> Value -> Inst
Out  (Value -> Inst) -> P Value -> P Inst
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Value
pValue
  String
_     -> P Inst
forall a. P a
forall (f :: * -> *) a. Alternative f => f a
empty

execute :: Vector Inst -> Int -> Bool
execute :: Vector Inst -> Int -> Bool
execute Vector Inst
program Int
a =
  State Machine Bool -> Machine -> Bool
forall s a. State s a -> s -> a
evalState State Machine Bool
theMain (Registers -> Progress -> Set (Int, Registers) -> Machine
Machine Registers
zeroRegisters Progress
NeedZero Set (Int, Registers)
forall a. Monoid a => a
mempty)
  where
    theMain :: State Machine Bool
theMain = do Register -> LensLike' Identity Machine Int
forall a (f :: * -> *).
(HasRegisters a, Functor f) =>
Register -> LensLike' f a Int
forall (f :: * -> *).
Functor f =>
Register -> LensLike' f Machine Int
reg Register
A LensLike' Identity Machine Int -> Int -> StateT Machine Identity ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= Int
a
                 Int -> State Machine Bool
goto Int
0

    step :: Int -> Inst -> State Machine Bool
step Int
pc = \case
      Out Value
o ->
         do Int
v <- Value -> StateT Machine Identity Int
forall r (m :: * -> *).
(MonadState r m, HasRegisters r) =>
Value -> m Int
rval Value
o
            Progress
progress <- Getting Progress Machine Progress
-> StateT Machine Identity Progress
forall s (m :: * -> *) a. MonadState s m => Getting a s a -> m a
use Getting Progress Machine Progress
Lens' Machine Progress
machProgress
            case (Progress
progress, Int
v) of
              (Progress
NeedOne,  Int
1) ->
                  do (Progress -> Identity Progress) -> Machine -> Identity Machine
Lens' Machine Progress
machProgress ((Progress -> Identity Progress) -> Machine -> Identity Machine)
-> Progress -> StateT Machine Identity ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= Progress
NeedZero
                     Int -> State Machine Bool
goto (Int
pcInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

              (Progress
NeedZero, Int
0) ->
                  do Registers
registers <- Getting Registers Machine Registers
-> StateT Machine Identity Registers
forall s (m :: * -> *) a. MonadState s m => Getting a s a -> m a
use Getting Registers Machine Registers
Lens' Machine Registers
machRegisters
                     Set (Int, Registers)
targets   <- Getting (Set (Int, Registers)) Machine (Set (Int, Registers))
-> StateT Machine Identity (Set (Int, Registers))
forall s (m :: * -> *) a. MonadState s m => Getting a s a -> m a
use Getting (Set (Int, Registers)) Machine (Set (Int, Registers))
Lens' Machine (Set (Int, Registers))
machTargets
                     let now :: (Int, Registers)
now = (Int
pc,Registers
registers)
                     if (Int, Registers) -> Set (Int, Registers) -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member (Int, Registers)
now Set (Int, Registers)
targets then
                       Bool -> State Machine Bool
forall a. a -> StateT Machine Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
                     else
                       do (Set (Int, Registers) -> Identity (Set (Int, Registers)))
-> Machine -> Identity Machine
Lens' Machine (Set (Int, Registers))
machTargets ((Set (Int, Registers) -> Identity (Set (Int, Registers)))
 -> Machine -> Identity Machine)
-> ((Bool -> Identity Bool)
    -> Set (Int, Registers) -> Identity (Set (Int, Registers)))
-> (Bool -> Identity Bool)
-> Machine
-> Identity Machine
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Index (Set (Int, Registers)) -> Lens' (Set (Int, Registers)) Bool
forall m. Contains m => Index m -> Lens' m Bool
contains (Int, Registers)
Index (Set (Int, Registers))
now ((Bool -> Identity Bool) -> Machine -> Identity Machine)
-> Bool -> StateT Machine Identity ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= Bool
True
                          (Progress -> Identity Progress) -> Machine -> Identity Machine
Lens' Machine Progress
machProgress               ((Progress -> Identity Progress) -> Machine -> Identity Machine)
-> Progress -> StateT Machine Identity ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= Progress
NeedOne
                          Int -> State Machine Bool
goto (Int
pcInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

              (Progress, Int)
_ -> Bool -> State Machine Bool
forall a. a -> StateT Machine Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

      Copy Value
i Register
o -> do Register -> LensLike' Identity Machine Int
forall a (f :: * -> *).
(HasRegisters a, Functor f) =>
Register -> LensLike' f a Int
forall (f :: * -> *).
Functor f =>
Register -> LensLike' f Machine Int
reg Register
o LensLike' Identity Machine Int
-> StateT Machine Identity Int -> StateT Machine Identity ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> m b -> m ()
<~ Value -> StateT Machine Identity Int
forall r (m :: * -> *).
(MonadState r m, HasRegisters r) =>
Value -> m Int
rval Value
i
                     Int -> State Machine Bool
goto (Int
pcInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

      Inc Register
r    -> do Register -> LensLike' Identity Machine Int
forall a (f :: * -> *).
(HasRegisters a, Functor f) =>
Register -> LensLike' f a Int
forall (f :: * -> *).
Functor f =>
Register -> LensLike' f Machine Int
reg Register
r LensLike' Identity Machine Int -> Int -> StateT Machine Identity ()
forall s (m :: * -> *) a.
(MonadState s m, Num a) =>
ASetter' s a -> a -> m ()
+= Int
1
                     Int -> State Machine Bool
goto (Int
pcInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

      Dec Register
r    -> do Register -> LensLike' Identity Machine Int
forall a (f :: * -> *).
(HasRegisters a, Functor f) =>
Register -> LensLike' f a Int
forall (f :: * -> *).
Functor f =>
Register -> LensLike' f Machine Int
reg Register
r LensLike' Identity Machine Int -> Int -> StateT Machine Identity ()
forall s (m :: * -> *) a.
(MonadState s m, Num a) =>
ASetter' s a -> a -> m ()
-= Int
1
                     Int -> State Machine Bool
goto (Int
pcInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

      Jnz Value
i Value
o  -> do Int
v  <- Value -> StateT Machine Identity Int
forall r (m :: * -> *).
(MonadState r m, HasRegisters r) =>
Value -> m Int
rval Value
i
                     Int
o' <- Value -> StateT Machine Identity Int
forall r (m :: * -> *).
(MonadState r m, HasRegisters r) =>
Value -> m Int
rval Value
o
                     let pcOff :: Int
pcOff = if Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Int
1 else Int
o'
                     Int -> State Machine Bool
goto (Int
pcInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
pcOff)

    goto :: Int -> State Machine Bool
goto Int
pc =
      case Vector Inst
program Vector Inst -> Int -> Maybe Inst
forall a. Vector a -> Int -> Maybe a
Vector.!? Int
pc of
        Maybe Inst
Nothing -> Bool -> State Machine Bool
forall a. a -> StateT Machine Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
        Just Inst
o  -> Int -> Inst -> State Machine Bool
step Int
pc Inst
o