Advent of Code 2020 day 15

Day 15 started as a fiddle with off-by-one errors, then dived into the detail of the ST monad, mutable values, and Control.Monad.Loops for control structures.

The problem was the same in both parts, but different sizes. The instance in part 1 was small enough for a pure functional version to be quick enough, but part 2 required mutable storage to be quick.

The functional solution

Progress in the game was built around the Game record. This records the state of the game at the end of a particular round. round is the number of the round just finished; word is the word/number spoken this round; history is the history of what was spoken in previous rounds.

data Game = Game { round :: Int
                 , word :: Int
                 , history :: M.IntMap Int
                 } deriving (Show, Eq)

I create a game from the seed string, taking care not to include the current round in the history.

seedGame seed = Game {..}
  where round = length seed 
        word = last seed
        history = M.fromList $ zip (init seed) [1..]

Given a game at the end of round n, I can create the game at the end of round n + 1 by looking up the word spoken and updating the history. (If a word has not been spoken before, say zero, otherwise say the gap.)

gameStep Game{..} = 
  Game { round = round + 1
       , word = word'
       , history = history'
       }
  where 
    word' = speakWord (M.lookup word history) round
    history' = M.insert word round history

speakWord Nothing _ = 0
speakWord (Just prev) now = now - prev

To find the desired state, I iterate to create an infinite list of states, and drop them while I'm notYet finished.

infiniteGame g = iterate gameStep g

gameRound r game0 = head $ dropWhile notYet $ infiniteGame game0
  where notYet game = round game < r

And then call that from main

main :: IO ()
main = 
  do  let seed = [20, 0, 1, 11, 6, 3]
      print $ part1 seed

part1 = word . (gameRound 2020) . seedGame

This version is available as advent15slow.hs.

The imperative solution

This solution works for part 2, but it takes about 6 minutes to complete the 30,000,000 rounds. I may have been satisfied with that, but I've been keeping an eye on Michael Wong's blog, and his solution to day 15 was quick, by using mutable storage in the ST monad. He's a beginner, and I'd never done anything with ST before, so… challenge accepted!

This version revolves around three mutable variables, round, word, and history. The first two are Ints, the last is Vector of Ints.

seedGame :: [Int] -> Int -> ST s (STRef s Int, STRef s Int,  V.MVector s Int)
seedGame seed historySize = 
  do round <- newSTRef $ length seed
     word <- newSTRef $ last seed
     history <- V.replicate historySize 0
     forM_ (zip (init seed) [1..]) $ \(t, s) -> V.write history t s
     return (round, word, history)

The s state in the type signatures is just to enforce the encapsulation of the mutation to the ST monad's scope. It's not used anywhere and you don't define it.

These variables are updated every round of the game. At the end of each call of gameSteps, these variables reflect the state of the game at the end of the current round. Calculating each game step is just shuffling values around these variables. Note that I also check for finishing the game in the same loop.

gameSteps :: Int -> STRef s Int -> STRef s Int -> V.MVector s Int -> ST s ()
gameSteps targetRound round word history =
  do roundVal <- readSTRef round
     if roundVal == targetRound
     then return ()
     else do 
           wordVal <- readSTRef word
           wordH <- V.read history wordVal
           let word' = speakWord wordH roundVal
           V.write history wordVal roundVal
           modifySTRef round (+1)
           writeSTRef word word'
           gameSteps targetRound round word history

This is all called with runGame, which sets up the ST monad to allow the mutability, seeds the game, runs it, then looks up the final word from the mutable variable.

runGame :: [Int] -> Int -> Int
runGame seed roundsNeeded =
  runST $ 
    do (round, word, history) <- seedGame seed roundsNeeded
       gameSteps roundsNeeded round word history
       readSTRef word

How depressingly imperative. It does, however, complete in less than two seconds.

This version is available as advent15.hs.

Separating the loop

That mixing of update and termination checking in gameSteps slightly bothered me, so I took a look at the Control.Monad.Loops library. Unfortunately, the documentation there is terse, to say the least. But I did find a nice post looking at the library, which was a great help.

I can simplify gameSteps into gameStep, that just does one game update.

gameStep :: STRef s Int -> STRef s Int -> V.MVector s Int -> ST s ()
gameStep round word history =
  do roundVal <- readSTRef round
     wordVal <- readSTRef word
     wordH <- V.read history wordVal
     let word' = speakWord wordH roundVal
     V.write history wordVal roundVal
     modifySTRef round (+1)
     writeSTRef word word'
     return ()

Control of this is handled in gameLoop. whileM_ takes two monadic expressions. The first returns an ST s Bool and determines whether the loop should continue. The second is the body of the loop.

gameLoop targetRound round word history =
    do whileM_ ((/= targetRound) <$> readSTRef round)
               $ gameStep round word history
       return ()

There's an alternative version using untilM_, but with the same two monadic expressions.

gameLoop targetRound round word history =
    do gameStep round word history 
         `untilM_` ((== targetRound) <$> readSTRef round)
       return ()

It's a matter of taste as to which one is clearer.

This version is available as advent15loop.hs, including a number of ways of writing gameLoop as comments.

Code

You can find the code here or on Gitlab.