It's finally time to get back to the Advent of Code solutions and optimise the very slow ones (as described in my Overview of Advent of Code 2023). In this post, I tackle the slowest-running day 14, with a run-time of a whopping 97 seconds.

The first thing is to profile the code to see what's actually causing the problem. That clearly showed that 94% of the time was in the rollStep function, that did the rolling of rocks around the grid.

The main problem here is my choice of representation for the underlying grid. I chose to represent the grid as a list of lists.

data Element = Empty | Cube | Round deriving (Show, Eq, Ord)
type Grid = [[Element]]

Rolling rocks across the grid was done by moving elements around this list of lists. However, we need to do a lot of that, and Haskell has immutable data structures. That means I can't update the list, but must instead create a new one after each change. And the structure of lists means that I have to walk along the prefix of the list each time, to find the section to copy from.

A better approach is to get closer to the metal and use a constant-time-access, mutable data structure. (Haskell allows this by wrapping the mutable-data code in an ST monad, so it appears immutable to the rest of the program.) The data structure I used was a mutable unboxed array.

Data structures

This change required some new data structures: the mutable array used during the processing, an immutable unboxed array for passing around the grid outside the mutation, and a mutable variable holding a queue of vacant positions.

type Position = V2 Int
type Grid = U.UArray Position Bool
type MGrid s = A.STUArray s Position Bool
type Gaps s = STRef s (Q.Seq Position)

This structure means I need two arrays to fully represent the state: one to hold the round rocks (that move) and one to hold the cube rocks (that don't, and block the round rocks).

Rolling

Rolling needs to walk over the grid, examining each position in turn and rolling any round rocks as I find them. For instance, when rolling rocks to the north, I start at the north-west corner (column 0, row 0). I roll the rocks in each column in turn, incrementing the columns as I go (the "major step"). For each column, I roll the rocks row by row in that column, incrementing the row as I go (the "minor step"). To roll west, I start in the north-west corner, increment the rows, and decrement the column within each row.

That means my roll function needs a group of three values to fix the order of operations when rolling all the round rocks in the grid.

rollNorth, rollCycle :: Grid -> Grid -> Grid
rollNorth rGrid cGrid = roll [(V2 0 0, V2 0 1, V2 1 0)] cGrid rGrid

rollCycle rGrid cGrid = roll [ (V2 0 0, V2 0 1, V2 1 0)
                             , (V2 0 0, V2 1 0, V2 0 1)
                             , (V2 r 0, V2 0 1, V2 -1 0)
                             , (V2 0 c, V2 1 0, V2 0 -1)
                             ]
                             cGrid rGrid
  where (_, V2 r c) = U.bounds rGrid

The roll function itself is a pair of nested loops, going across the the major steps then the minor steps. Rolling is controlled by a queue of "holes" that a round rock can roll into.

roll :: [(Position, Position, Position)] -> Grid -> Grid -> Grid
roll moveSpecs cGrid rGrid =
  A.runSTUArray $ 
    do grid <- A.thaw rGrid 
       holes <- newSTRef Q.Empty
       forM_ moveSpecs $ \(start, majorStep, minorStep) -> 
         forM_ (takeWhile (inBounds rGrid) $ iterate (^+^ majorStep) start) $ \maj ->
           do writeSTRef holes Q.Empty
              forM_ (takeWhile (inBounds rGrid) $ iterate (^+^ minorStep) maj) $ \here ->
                rollPosition grid cGrid holes here
       return grid

As the current position moves along in minor steps, I maintain a queue of vacant positions. When I get to a round rock, the queue tells me all the positions it will roll across until it stops.

-- Direction of rolling -->
Rocks:           ..O...#..
Queue:              xxx
Current position:  ^
After rolling:   .....O#..
                   xxx

The update of the queue (and rolling of round rocks) is handled by the rollPosition function. It covers a few cases:

1. If the current position is a cube rock, make the queue empty
2. If the current position is empty, add it to the queue
3. If the current position is a round rock, look at the queue
   1. If the queue is empty, do nothing
   2. Otherwise, move the rock to the head of the queue and add this position to the tail of the queue (it's now empty)

rollPosition :: (MGrid s) -> Grid -> (Gaps s) -> Position -> ST s ()
rollPosition grid cGrid holes here
  | cGrid U.! here = writeSTRef holes Q.Empty                          (1)
  | otherwise = do roundHere <- A.readArray grid here
                   holesVal <- readSTRef holes
                   if roundHere then
                     case holesVal of
                       Q.Empty -> return ()                            (3.1)
                       (h :<| hs) -> do A.writeArray grid h True       (3.2)
                                        A.writeArray grid here False
                                        writeSTRef holes (hs :|> here)
                   else modifySTRef holes (:|> here)                    (2)

A couple of other bits had to change, such as the load calculation, but nothing of any great significance.

Performance

The runtime changed from 65 seconds to 0.55 seconds, a speed up of about 120 times. Not too shabby for a small change in representation!

Code

You can get the code from my locally-hosted Git repo, or from Gitlab.