Taking advantage of a neat trick for using $ rather than a lambda
[advent-of-code-17.git] / src / advent05 / advent05vector.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