AOC008.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}

module AOC008 () where

import Console
  (Instruction(..), Console, Error(..), ConsoleHook, toConsole, acc, runWithHook)
import Control.Lens
  ((.=), view, use, _1)
import Control.Monad
  (when)
import Control.Monad.Except
  (throwError, catchError, runExceptT)
import Control.Monad.State
  (get, put, execState)
import Data.Has
  (Has(..))
import Data.HashSet
  (HashSet, member)
import Problem
  (Solved(..), Part(..))

repeatHook :: Has (HashSet Integer) s => ConsoleHook s m
repeatHook cpc i continue = do
  pcs <- use hasLens
  when (cpc `member` pcs) $ throwError Reexecuted
  hasLens .= [cpc] <> pcs
  continue i

flipHook :: (Has (HashSet (Integer, Bool)) s, Has Bool s) => ConsoleHook s m
flipHook cpc inst continue = do
  pcs <- use hasLens
  modified <- use hasLens
  when ((cpc, modified) `member` pcs) $ throwError Reexecuted
  hasLens .= [(cpc, modified)] <> pcs

  if not modified && modifiable inst then do
    og <- get
    hasLens .= [(cpc, True)] <> pcs
    hasLens .= True
    -- try to run the new program
    -- if it fails, restore the state
    continue (modify inst) `catchError` \_ -> do
      npcs <- use $ hasLens @(HashSet (Integer, Bool))
      put og
      hasLens .= npcs
      continue inst
  else
    continue inst
  where modify (Jmp i) = Nop i
        modify (Nop i) = Jmp i
        modify i = i
        modifiable (Jmp _) = True
        modifiable (Nop _) = True
        modifiable _ = False

checkLoop :: Console -> (Console, HashSet Integer)
checkLoop = execState (runExceptT $ runWithHook repeatHook) . (,mempty)

checkFlip :: Console -> (Console, HashSet (Integer, Bool), Bool)
checkFlip = execState (runExceptT $ runWithHook flipHook) . (,mempty,False)

instance Solved 8 'PartA where
  solve = show . view (_1 . acc) . checkLoop . toConsole . fmap read . lines
instance Solved 8 'PartB where
  solve = show . view (_1 . acc) . checkFlip . toConsole . fmap read . lines

Console.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE RankNTypes #-}

module Console
  ( Instruction(..)
  , Console
  , acc
  , pc
  , prog
  , halted
  , Program
  , Error(..)
  , ConsoleHook
  , toConsole
  , toProgram
  , runWithHook
  , run
  ) where

import Control.Lens
  ((.=), (+=), use, preuse, ix)
import Control.Lens.TH
  (makeLenses)
import Control.Monad
  (unless)
import Data.Hashable
  (Hashable)
import Control.Monad.Except
  (MonadError, throwError)
import Control.Monad.State
  (MonadState)
import Data.Char
  (isDigit)
import Data.Has
  (Has(..))
import Data.HashMap.Strict
  (HashMap, fromList)
import GHC.Generics
  (Generic)
import Text.ParserCombinators.ReadP
  (choice, string, char, munch1)
import Text.ParserCombinators.ReadPrec
  (lift)
import Text.Read
  (Read(..))

data Instruction = Acc Integer
                 | Jmp Integer
                 | Nop Integer
                 | Hlt
  deriving (Show, Eq, Generic)
instance Hashable Instruction

type Program = HashMap Integer Instruction

data Console =
  Console
  { _acc    :: Integer
  , _pc     :: Integer
  , _prog   :: Program
  , _halted :: Bool
  }
makeLenses ''Console

data Error = NoInst
           | Reexecuted
  deriving (Show, Eq, Enum, Bounded)

type ConsoleHook s m = (Has Console s, MonadState s m, MonadError Error m)
                     => (Integer -> Instruction -> (Instruction -> m ()) -> m ())

instance Read Instruction where
  readPrec = lift $ do
    ins <- choice [string "acc" >> pure Acc, string "jmp" >> pure Jmp, string "nop" >> pure Nop]
    _ <- char ' '
    sgn <- choice [char '+' >> pure id, char '-' >> pure negate]
    arg <- read <$> munch1 isDigit
    return $ ins (sgn arg)

toConsole :: [Instruction] -> Console
toConsole is = Console 0 0 (toProgram is) False

toProgram :: [Instruction] -> Program
toProgram = fromList . zip [0..] . (<>[Hlt])

runWithHook :: (Has Console s, MonadState s m, MonadError Error m)
            => ConsoleHook s m -> m ()
runWithHook hook = do
  cpc <- use $ hasLens . pc
  mins <- preuse $ hasLens . prog . ix cpc
  case mins of
    Just ins -> hook cpc ins (step hook)
    Nothing -> throwError NoInst

step :: (Has Console s, MonadState s m, MonadError Error m)
     => ConsoleHook s m -> Instruction -> m ()
step hook ins = do
  case ins of
    Acc i -> hasLens . acc += i >> hasLens . pc += 1
    Jmp o -> hasLens . pc += o
    Nop _ -> hasLens . pc += 1
    Hlt   -> hasLens . halted .= True
  stop <- use $ hasLens . halted
  unless stop (runWithHook hook)

run :: (Has Console s, MonadState s m, MonadError Error m) => m ()
run = runWithHook $ \_ i s -> s i