Breadth-First Search in Haskell

In 2016 (strictly speaking, in August), I, along with two others, tried my hand at the year’s ICFP competition. This was pretty fun – although a nice reminder of how much computational geometry I’d forgotten – and although we didn’t achieve all that much in the way of concrete results, we certainly learned a few things (see the repo for details). One is that I’d quite like to have a few more people, so that we could try multiple approaches.

What I really missed when doing this was a general technique for search in Haskell (sure, we could have not used Haskell, but there was a certain appeal to it). I wasn’t going to write one of those while trying to solve a complex problem for no stakes under a time limit, but I figured I’d come back to it.

Enter Reality

I didn’t, and eventually December rolled around, and with that came the second round of Advent of Code. I’d started playing with the site in 2015, but this time around there were a few coworkers and I was very much in the mood for something that wasn’t like my day job.

In one of those entertaining moments, I really missed having a general technique for search in Haskell. To give an idea of why Haskell comes up a lot, I write in Ruby most of the time, so something like Haskell makes for a fun change of pace. Also, I find the type system interesting: it offers a mechanism I can use to guide software development, and lets me avoid depending on more than I think I am.

I didn’t have a great deal of spare time in December (bar the Christmas period), so again I didn’t really get into solving the general search problem until fairly late. Given some of the complexities the Advent of Code problems established, this was probably a good thing.

What I’m (Not) Trying to Do

This is not a post about building a full-speed, highly-optimised heuristic search engine. That’s an interesting problem, but it’s been a while since I was really thinking about search and practical optimisation is something that benefits both from significant experience with a language in production and several other people to talk about possibilities with.

Instead, I want to construct a useful, comprehensible implementation of a search tool that can illustrate problems. Maximal speed is not critical (but we should avoid doing anything really terrible). Because I’m not completely insane, we’ll require that our branching factor always be finite, and start with a breadth-first implmentation and iterate our way towards a general solution.

Possible Preconditions

In general, we’d like to search a state space until we encounter some state that meets a termination condition. The implication here is that we have some defined type a, corresponding to states, and a termination condition f : a -> Boolean. Effective search in a recurrent space requires us to prune states we have already seen – this implies the existence of a computable¹ equivalence relation. If we suppose that we are generating large numbers of states, we would equally want our pruning to not require traversal of large lists; this suggests we back our collection of seen states by some more intelligent set. I’ve elected to use a hash here; a tree-based set could work fine as well (substitute Ord a below).

These constraints imply Hashable a and Eq a; to simplify matters we will elide the termination condition throughout the following treatment (rather, the consumer will traverse the yielded states and apply termination themselves – non-strict semantics to the rescue)

Secondly, we need a way to identify what possible states we can reach from any particular state. That is, we need the adjacency function adj : a -> [a]. The list here is one of those facets that is not strictly necessary – I suspect that any general Traversable a would be adequate – but we’ve already got enough on our plates.

This suggests a type signature such as the following:

bfs :: (Hashable a, Eq a) => (a -> [a]) -> a -> [a]

which is certainly something we can work with.

Breadth-First Traversal

Now that we have our types, we might as well get into defining the actual breadth-first search implementation. In order to do so, we should first talk about the effective requirements:

• every node within n steps of the origin should be visited before a node that is n+1 steps away
• any node should be visited at most one time (with respect to its atomic existence, not the (minimal) path by which it has been reached)

That ordering requirement is important: it suggests a queue, and queues are handled very poorly by singly-linked lists. Enter Data.Sequence.

In deference to tradition, bfs can just handle setup for an internal variant (consumers probably don’t really care about how we manage our states, after all…)

import           Data.Hashable
import qualified Data.HashSet  as Set
import qualified Data.Sequence as Seq

bfs adj start = bfs' adj seen queue
  where
    seen  = Set.singleton start
    queue = Seq.singleton start

We now have a relatively straightforward process: yield the first element in the queue, construct a new set of seen nodes and queue, and iterate this process. We want to do so without literal recursion (no, really), so we might as well apply unfoldr². For clarity, we will extract the individual steps:

{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE PatternSynonyms #-}
    
import Data.List (nubBy)
    
bfs_step neighbours (seen, queue)
  | Seq.null queue = Nothing
  | otherwise      = Just (current, next)
    where
      (Seq.viewl -> (current Seq.:< remaining)) = queue
      next = (seen', queue')
      seen' = Set.union seen (Set.fromList descendents)
      queue' = remaining Seq.>< (Seq.fromList descendents)
      descendents =
	nubBy (==)
	[
	  child
	| child <- adj current
	, not $ (flip Set.member seen) child
	]

We incur some clarity overhead here – Seq.:< is a bit of a painful invocation, for instance – but the general approach of filtering neighbours should be fairly clear. The need for nubBy is essentially security against implementations of adj that at some points may generate the same entry more than once for a single node.

Now that we have a mechanism for generating each step, implementing the actual (internal) BFS function is essentially immediate:

import Data.List (unfoldr)

bfs' :: (Hashable a, Eq a)
  => (a -> [a])
  -> Set.HashSet a
  -> Seq.Seq a
  -> [a]
bfs' adj seen queue = unfoldr (bfs_step adj) (seen, queue)

Controlled Equivalence

It sometimes comes to pass that we can know that two states will expand in an equivalent manner. Imagine, for instance, the colour of the first vertex in an n-colouring: it is fundamentally irrelevant what choice is made here, due to the symmetries involved. Alternatively, we might know that two components of a search are independent – we may as well preference expanding the one we have already started on.

Rather than try to prune here based on some provided relation, let’s implement this in terms of projection. This is approximately the canonical projection, in the classical set theoretic sense. There is no guarantee that the type of the resultant object is the same as the individual states, so we end up with some function proj : a -> b, and our constraints naturally pass through to b (as these are the objects being used to constrain search expansion…)

For the sake of clarity, we’ll call these searches “masked”.

bfsm mask adj start = bfsm' mask adj seen queue
  where
    seen  = Set.singleton (mask start)
    queue = Seq.singleton start

bfsm' :: (Hashable b, Eq b)
  => (a -> b)
  -> (a -> [a])
  -> Set.HashSet b
  -> Seq.Seq a
  -> [a]
bfsm' mask adj seen queue = unfoldr (bfsm_step mask adj) (seen queue)

This leaves only the core step implementation, which should be awfully familiar:

bfsm_step mask adj (seen, queue)
  | Seq.null queue = Nothing
  | otherwise      = Just (current, next)
    where
      (Seq.viewl -> (current Seq.:< remaining)) = queue
      next = (seen', queue')
      seen' = Set.union seen (Set.fromList masked_descendents)
      queue' = remaining Seq.>< (Seq.fromList descendents)
      masked_descendents = fmap mask descendents
      descendents =
	nubBy (\x y -> mask x == mask y)
	[
	  child
	| child <- adj current
	, not $ (flip Set.member seen) (mask child)
	]

Conveniently, this lets us redefine bfs = bfsm id.

Up Next

Now that we have masked (and unmasked…) search, it’s time to get into heuristics.