advent05/src/advent05.hs

   1 import Debug.Trace
   2
   3 import Data.Text (Text)
   4 import qualified Data.Text.IO as TIO
   5
   6 import Data.Void (Void)
   7
   8 import Text.Megaparsec hiding (State)
   9 import Text.Megaparsec.Char
  10 import qualified Text.Megaparsec.Char.Lexer as L
  11 import qualified Control.Applicative as CA
  12
  13 import Control.Monad (unless)
  14 import Control.Monad.State.Strict
  15 import Control.Monad.Reader
  16 import Control.Monad.Writer
  17
  18 import qualified Data.IntMap.Strict as M
  19 import Data.IntMap.Strict ((!))
  20 import Data.List
  21
  22 type Memory = M.IntMap Int
  23
  24 data Machine = Machine { _memory :: Memory
  25                        , _ip :: Int
  26                        , _inputIndex :: Int
  27                        }
  28                deriving (Show, Eq)
  29
  30 type ProgrammedMachine = WriterT [Int] (ReaderT ([Int]) (State Machine)) ()
  31
  32 data ParameterMode = Position | Immediate deriving (Ord, Eq, Show)
  33
  34
  35 main :: IO ()
  36 main = do
  37         text <- TIO.readFile "data/advent05.txt"
  38         let mem = successfulParse text
  39         print $ findMachineOutput [1] mem
  40         print $ findMachineOutput [5] mem
  41
  42
  43 findMachineOutput inputs program = last output
  44     where   finalStack =
  45                 runState (
  46                     runReaderT (
  47                         runWriterT runAll
  48                                )
  49                         inputs
  50                          )
  51                          (makeMachine program)
  52             ((_retval, output), _machine) = finalStack
  53
  54 makeMachine :: [Int] -> Machine
  55 makeMachine memory = Machine    {_ip = 0, _inputIndex = 0
  56                                 , _memory = M.fromList $ zip [0..] memory
  57                                 }
  58
  59
  60 runAll :: ProgrammedMachine
  61 runAll = do mem <- gets _memory
  62             ip <- gets _ip
  63             unless (mem!ip == 99)
  64                 do runStep
  65                    runAll
  66
  67 runStep :: ProgrammedMachine
  68 runStep =
  69     do mem <- gets _memory
  70        ip <- gets _ip
  71        let opcode = (mem!ip) `mod` 100
  72        let modes = parameterModes ((mem!ip) `div` 100)
  73        fetchInput opcode
  74        putOutput opcode modes
  75        mem' <- gets _memory
  76        let (mem'', ip') = perform opcode ip modes mem'
  77        modify (\m -> m {_ip = ip', _memory = mem''})
  78
  79
  80 -- fetchInput opcode | trace ("Input with opcode " ++ show opcode) False = undefined
  81 fetchInput 3 =
  82     do mem <- gets _memory
  83        ip <- gets _ip
  84        inputIndex <- gets _inputIndex
  85        inputs <- ask
  86        let mem' = iInsert (ip + 1) (inputs!!inputIndex) mem
  87        modify (\m -> m {_inputIndex = inputIndex + 1, _memory = mem'})
  88 fetchInput _ = return ()
  89
  90
  91 -- putOutput opcode _modes | trace ("Output with opcode " ++ show opcode) False = undefined
  92 putOutput 4 modes =
  93     do mem <- gets _memory
  94        ip <- gets _ip
  95        let v = getMemoryValue (ip + 1) (modes!!0) mem
  96        tell [v]
  97 putOutput _ _ = return ()
  98
  99
 100 perform :: Int -> Int -> [ParameterMode] -> Memory -> (Memory, Int)
 101 -- perform instr ip modes mem | trace ("Perform ip " ++ show ip ++ " opcode " ++ show instr ++ " modes " ++ (show (take 3 modes)) ++ " args " ++ (intercalate ", " (map show [(mem!(ip+1)), (mem!(ip+2)), (mem!(ip+3))]))) False = undefined
 102 perform 1 ip modes mem = (iInsert (ip + 3) (a + b) mem, ip + 4)
 103     where a = getMemoryValue (ip + 1) (modes!!0) mem
 104           b = getMemoryValue (ip + 2) (modes!!1) mem
 105 perform 2 ip modes mem = (iInsert (ip + 3) (a * b) mem, ip + 4)
 106     where a = getMemoryValue (ip + 1) (modes!!0) mem
 107           b = getMemoryValue (ip + 2) (modes!!1) mem
 108 perform 3 ip _ mem = (mem, ip + 2)
 109 perform 4 ip _ mem = (mem, ip + 2)
 110 perform 5 ip modes mem = (mem, ip')
 111     where a = getMemoryValue (ip + 1) (modes!!0) mem
 112           b = getMemoryValue (ip + 2) (modes!!1) mem
 113           ip' = if a /= 0 then b else ip + 3
 114 perform 6 ip modes mem = (mem, ip')
 115     where a = getMemoryValue (ip + 1) (modes!!0) mem
 116           b = getMemoryValue (ip + 2) (modes!!1) mem
 117           ip' = if a == 0 then b else ip + 3
 118 perform 7 ip modes mem = (iInsert (ip + 3) res mem, ip + 4)
 119     where a = getMemoryValue (ip + 1) (modes!!0) mem
 120           b = getMemoryValue (ip + 2) (modes!!1) mem
 121           res = if a < b then 1 else 0
 122 perform 8 ip modes mem = (iInsert (ip + 3) res mem, ip + 4)
 123     where a = getMemoryValue (ip + 1) (modes!!0) mem
 124           b = getMemoryValue (ip + 2) (modes!!1) mem
 125           res = if a == b then 1 else 0
 126 perform _ ip _ mem = (mem, ip)
 127
 128
 129 getMemoryValue loc Position mem = mem!>loc
 130 getMemoryValue loc Immediate mem = mem!loc
 131
 132
 133 parameterModes :: Int -> [ParameterMode]
 134 parameterModes modeCode = unfoldr generateMode modeCode
 135
 136 generateMode :: Int -> Maybe (ParameterMode, Int)
 137 generateMode modeCode = Just (mode, modeCode `div` 10)
 138     where mode = case (modeCode `mod` 10) of
 139                     0 -> Position
 140                     1 -> Immediate
 141
 142
 143 -- Some IntMap utility functions, for syntactic sugar
 144
 145 -- prefix version of (!)
 146 lkup k m = m!k
 147
 148 -- indirect lookup
 149 (!>) m k = m!(m!k)
 150
 151 -- indirect insert
 152 iInsert k v m = M.insert (m!k) v m
 153
 154
 155
 156 -- Parse the input file
 157 type Parser = Parsec Void Text
 158
 159 sc :: Parser ()
 160 sc = L.space (skipSome spaceChar) CA.empty CA.empty
 161 -- sc = L.space (skipSome (char ' ')) CA.empty CA.empty
 162
 163 lexeme  = L.lexeme sc
 164 integer = lexeme L.decimal
 165 signedInteger = L.signed sc integer
 166 symb = L.symbol sc
 167 comma = symb ","
 168
 169 memoryP = signedInteger `sepBy` comma
 170
 171 successfulParse :: Text -> [Int]
 172 successfulParse input =
 173         case parse memoryP "input" input of
 174                 Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 175                 Right memory -> memory