Search code examples
haskellspace-leak

Avoiding space leaks with `mapM` and `foldM` over `State` monad


How do I avoid space leaks while using foldM and mapM over a State monad?

Last year's Advent of Code day 20 has a puzzle of generating a map of a maze from instructions on how to walk across it. For instance, the instructions NN gives the maze

   |
   |
   *

(a straight corridor two steps northwards), and the instructions NNN(EE|WW)S gives the maze

 +-+-+
 | | |
   |
   *

(go north a bit, then either go east then south or west then south).

The way I'm trying to solve this involves having a State monad, where the state is the Set of all the corridor sections (termed Doors below), and the value is the list of positions you could be working from.

If you're just following a corridor Path, I use foldM to walk along it, updating the current position. If you're at a junction, follow each branch of the junction and collect all the positions you end up.

This code produces the correct results on small test inputs, but there's a huge space leak when working on the full example.

Profiling indicates it's spending most of its time in includeDoor.

So, questions.

  1. Is there a space leak? If so, where, and how can you tell.
  2. How do I fix it?

(I think what's happening is that Haskell isn't strictly adding fully-evaluated Doors to the Set as soon as it can. In this case, I don't want any laziness anywhere.)

(I parse the input into a bunch of two-element vectors that indicate the step to take for each instruction. That code works fine, and quickly.)

import qualified Data.Set as S

import Linear (V2(..))

import Control.Monad.State.Strict
import Control.Monad.Extra (concatMapM)

type Coord = V2 Integer -- x, y, with north and east incresing values (origin a bottom left)
data Door = Door Coord Coord deriving (Show, Eq, Ord)
type Doors = S.Set Door

data MazeSection = Path [Coord] | Junction [Maze] deriving (Show, Eq)
type Maze = [MazeSection]

type Mapper = State Doors [Coord]

makeDoor :: Coord -> Coord -> Door
makeDoor !a !b 
    | a < b = Door a b
    | otherwise = Door b a

emptyMap = S.empty

part1 maze = 
    do 
        let start = V2 0 0
        let doors = execState (mapMaze [start] maze) emptyMap
        print $ length doors

mapMaze :: [Coord] -> Maze -> Mapper
mapMaze !starts !sections =
    foldM (\heres section -> mapMazeSection heres section) starts sections

mapMazeSection :: [Coord] -> MazeSection -> Mapper
mapMazeSection !starts (Junction mazes) = 
    concatMapM (\maze -> mapMaze starts maze) mazes
mapMazeSection !starts (Path steps) = 
    mapM mapPath starts
    where mapPath start = foldM (\here step -> includeDoor here step) start steps

includeDoor :: Coord -> Coord -> State Doors Coord
includeDoor !here !step = 
    do let there = (here + step)
       let door = there `seq` makeDoor here there
       modify' (door `seq` S.insert door)
       return there

Solution

  • Turns out, it wasn't a space leak! It was me failing to deal with some pathological input. Once I sorted out how to handle that, it worked, and very quickly.