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