Haskell space leak in implementation of BFS - haskell

Question

I have been banging my head against a Haskell space leak (of the stack overflow kind, naturally) for a few straight days. It's frustrating because I'm attempting to mimic the BFS algorithm straight from CLR, which is not naturally recursive. NB: I have enabled BangPatterns and I have put a bang in front of every possible place where one can go, in an attempt to branch-and-bound this problem, with no effect. I have battled through space leaks before, and I am loth to give up and cry for help on this one, but at this point I'm stuck. I love coding in Haskell, and I understand the Zen of functional programming pretty well, but debugging space leaks is about as much fun as rolling around on a floor full of thumbtacks.
That said, my trouble appears to be a space leak of the typical "accumulator" kind. The stack evidently builds up around calls to bfs' in the code below. Any space-leak protips much appreciated.
import qualified Data.Map as M
import qualified Data.IntSet as IS
import qualified Data.Sequence as S
import qualified Data.List as DL
data BfsColor = White | Gray | Black deriving Show
data Node =
Node {
neighbors :: !IS.IntSet,
color :: !BfsColor,
depth :: !Int
}
type NodeID = Int
type NodeQueue = S.Seq NodeID
type Graph = M.Map NodeID Node
bfs :: Graph -> NodeID -> Graph
bfs graph start_node =
bfs' (S.singleton start_node) graph
bfs' :: NodeQueue -> Graph -> Graph
bfs' !queue !graph
| S.null queue = graph
| otherwise =
let (u,q1) = pop_left queue
Node children _ n = graph M.! u
(g2,q2) = IS.fold (enqueue_child_at_depth $ n+1) (graph,q1) children
g3 = set_color u Black g2
in bfs' q2 g3
enqueue_child_at_depth :: Int -> NodeID -> (Graph, NodeQueue)
-> (Graph, NodeQueue)
enqueue_child_at_depth depth child (graph,!queue) =
case get_color child graph of
White -> (set_color child Gray $ set_depth child depth graph,
queue S.|> child)
otherwise -> (graph,queue)
pop_left :: NodeQueue -> (NodeID, NodeQueue)
pop_left queue =
let (a,b) = S.splitAt 1 queue
in (a `S.index` 0, b)
set_color :: NodeID -> BfsColor -> Graph -> Graph
set_color node_id c graph =
M.adjust (\node -> node{color=c}) node_id graph
get_color :: NodeID -> Graph -> BfsColor
get_color node_id graph = color $ graph M.! node_id
set_depth :: NodeID -> Int -> Graph -> Graph
set_depth node_id d graph =
M.adjust (\node -> node{depth=d}) node_id graph

Answer 1

That looks much easier to understand. (You can still shrink the code by 1/2, though.)
Now, the nature of the space leak becomes apparent. Namely, the one thing that is never evaluated is the depth. It will pile up to a big expression 1+1+.... You can remove all the bang patterns and add a single one at
enqueue_child_at_depth !depth child (graph,queue)
to get rid of the space leak.
(Further code tips: You can replace the IS.IntSet by a simple list. The queue is best deconstructed and reconstructed along the lines of
go depth qs graph = case viewl qs of
EmptyL -> graph
q :< qs ->
let
qs' = (qs ><) . Seq.fromList
. filter (\q -> isWhite q graph)
. neighbors q $ graph
in ...
)

Answer 2

First of all, if would be very helpful if you could provide some simple test case (in the form of code) which demonstrates how this thing stack overflows.
Without it I, personally, can only speculate on the subject of reason for that.
As a speculation: is IS.fold strict enough? Well, for example the following simplest code stack overflows as well (GHC with -O2):
{-# LANGUAGE BangPatterns #-}
import qualified Data.IntSet as IS
test s = IS.fold it 1 s
where it !e !s = s+e
main = print $ test (IS.fromList [1..1000000])
The overflow problem with this code can be hackafixed (is there a better way?) like that:
test s = foldl' it 1 (IS.toList s)
where it !e !s = s+e
Maybe you want to look at IS.fold in your code as well.