src/advent05/advent05.hs

   1 {-# LANGUAGE NegativeLiterals #-}
   2 {-# LANGUAGE FlexibleContexts #-}
   3
   4 import qualified Data.IntMap.Strict as M
   5 import Data.IntMap.Strict ((!))
   6 import Control.Monad.State.Lazy
   7
   8 data Machine = Machine { location :: Int
   9                        , steps :: Int
  10                        , memory :: M.IntMap Int
  11                        } deriving (Show, Eq)
  12
  13 main :: IO ()
  14 main = do
  15         text <- readFile "data/advent05.txt"
  16         let locations = map (readJump) $ lines text
  17         let m0 = makeMachine locations
  18         print $ evalState stepAll m0
  19         print $ evalState stepAllB m0
  20
  21
  22 readJump :: String -> Int
  23 readJump = read
  24
  25 makeMachine :: [Int] -> Machine
  26 makeMachine locations = Machine {location = 0, steps = 0,
  27     memory = M.fromList $ zip [0..] locations}
  28
  29 stepAll :: State Machine Int
  30 stepAll = do
  31             m0 <- get
  32             if M.member (location m0) (memory m0)
  33             then do stepOnce
  34                     stepAll
  35             else return (steps m0)
  36
  37 stepAllB :: State Machine Int
  38 stepAllB = do
  39             m0 <- get
  40             if M.member (location m0) (memory m0)
  41             then do stepOnceB
  42                     stepAllB
  43             else return (steps m0)
  44
  45 stepOnce :: State Machine ()
  46 stepOnce =
  47     do m0 <- get
  48        let mem = memory m0
  49        let loc = location m0
  50        let loc' = mem!loc + loc
  51        let steps' = steps m0 + 1
  52        let mem' = M.insert loc (mem!loc + 1) mem
  53        put m0 {location = loc', steps = steps', memory = mem'}
  54
  55 stepOnceB :: State Machine ()
  56 stepOnceB =
  57     do m0 <- get
  58        let mem = memory m0
  59        let loc = location m0
  60        let loc' = mem!loc + loc
  61        let steps' = steps m0 + 1
  62        let newVal = if mem!loc >= 3 then (mem!loc - 1) else (mem!loc + 1)
  63        let mem' = M.insert loc newVal mem
  64        put m0 {location = loc', steps = steps', memory = mem'}
  65