Dijkstra with Monads!

Written By James Bowen

Last time on the blog, we considered this library version of Dijkstra's algorithm that you can find on Hackage.

dijkstra ::
    (Foldable f, Num cost, Ord cost, Ord State)
  => (state -> f state) -- Function to generate list of neighbors
  -> (state -> state -> cost) -- Function for cost generation
  -> (state -> Bool) -- Destination predicate
  -> state -- Initial state
  -> Maybe (cost, [state]) -- Solution cost and path, Nothing if goal is unreachable

However, there are a number of situations where this might be insufficient. In this article we'll consider some reasons why you would want to introduce monads into your solution for Dijkstra's algorithm. Let's explore some of these reasons!

Note! This is a "coding ideas" blog, rather than an "In Depth Tutorial" blog (see this article for a summary of different reading styles). Some of the code sampled are pretty well fleshed out, but some of them are more hypothetical ideas for you to try out on your own!

The Monadic Version

In addition to the "pure" version of Dijkstra's algorithm, the Algorithm.Search library also provides a "monadic" version. This version allows each of the input functions to act within a monad m, and of course gives its final result within this monad.

dijkstraM ::
    (Monad m, Foldable f, Num cost, Ord cost, Ord State)
  => (state -> m (f state))
  -> (state -> state ->m cost)
  -> (state -> m Bool)
  -> state
  -> m (Maybe (cost, [state]))

Now, if you've read our Monads Series, you'll know that a monad is a computational context. What are the kinds of contexts we might find ourselves in while performing Dijkstra's algorithm? Here are a few ideas to start with.

  1. We're reading our graph from a global state (mutable or immutable) 2.Our graph functions require reading a file or making a network call
  2. We would like to log the actions taken in our graph.

Let's go through some pseudocode examples to see how each of these could be useful.

Using a Global State

A global mutable state is represented with (of course) the State monad. An immutable global state uses the Reader monad to represent this context. Now, taken in a simple way, the Reader context could allow us to "pass the graph" without actually including it as an argument:

import qualified Data.Array as A
import Control.Monad.Reader
import Algorithm.Search (dijkstraM)

newtype Graph2D = Graph2D (A.Array (Int, Int) Int)

getNeighbors :: A.Array (Int, Int) Int -> (Int, Int) -> [(Int, Int)]

findShortestPath :: Graph2D -> (Int, Int) -> (Int, Int) -> Maybe (Int, [(Int, Int)])
findShortestPath graph start end = runReader
  (dijkstraM neighbors cost (return . (== end)) start)
  graph
  where
    cost :: (Int, Int) -> (Int, Int) -> Reader Graph2D Int
    cost _ b = do
      (Graph2D gr) <- ask
      return $ gr A.! b

    neighbors :: (Int, Int) -> Reader Graph2D [(Int, Int)]
    neighbors source = do
      (Graph2D gr) <- ask
      return $ getNeighbors gr source

If we're already in this monad for whatever reason, then this could make sense. But on its own, it's not necessarily much of an improvement over partial function application.

A mutable state could be useful in certain circumstances as well. We likely wouldn't want to mutate the graph itself during iteration, as this would invalidate the algorithm. However, we could store certain metadata about what is happening during the search. For instance, we might want to track how often certain nodes are returned as a potential neighbor.

import qualified Data.HashMap.Strict as HM
import Control.Monad.State
import Data.Maybe (fromMaybe)
import Data.Foldable (find)

newtype Graph = Graph
   { edges :: HM.HashMap String [(String, Int)] }

type Metadata = HM.HashMap String Int
incrementKey :: String -> Metadata -> Metadata
incrementKey k metadata = HM.insert k (count + 1) metadata
  where
    count = fromMaybe 0 (HM.lookup k metadata)

findShortestPath :: Graph -> String -> String -> Maybe (Int, [String])
findShortestPath graph start end = evalState
  (dijkstraM neighbors cost (return . (== end)) start)
  HM.empty
  where
    cost :: String -> String -> State Metadata Int
    cost n1 n2 = 
      let assocs = fromMaybe [] (HM.lookup n1 (edges graph))
          costForN2 = find (\(n, _) -> n == n2) assocs
      in  case costForN2 of
            Nothing -> return maxBound
            Just (_, x) -> return x
    neighbors :: String -> State Metadata [String]
    neighbors node = do
      let neighbors = fst <$> fromMaybe [] (HM.lookup node (edges graph))
      metadata <- get
      put $ foldr incrementKey metadata neighbors
      return neighbors

In this implementation, we end up discarding our metadata, but if we wanted to we could include it as an additional output to help us understand what's happening in our search.

Reading from Files

In many cases, our "graph" is actually too big to fit within memory. In various cases, the entire graph could be distributed across many files on our system. Consider this simplified example:

data Location = Location
  { filename :: FilePath
  , tag :: String
  ...
  }

Each file could track a certain "region" of your map, with references to certain locations "on the edge" whose primary data must be found in a different file. This means you'll need to have access to the file system to ensure you can find all the "neighbors" of a particular location: This means you'll need the IO monad in Haskell!

getMapNeighbors :: Location -> IO [Location]
-- Open original locations file
-- Find tag and include neighboring tags together with references to other files

This matches the signature of the "neighbor generator" function in dijkstraM, so we'll be able to pass this function as the first argument.

Using Network Calls

Here's a fun example. Consider wiki-racing - finding the shortest path between the Wikipedia pages of two topics using only the links in the bodies of those pages. You could (theoretically) write a program to do this for you. You might create a type like this:

data WikiPage = WikiPage
  { pageTitle :: Text
  , url :: URL
  , bodyContentHtml :: Text
  }

In order to find the "neighbors" of this page, you would first have to parse the body HTML and find all the wikipedia links within it. This could be done in a pure fashion. But in order to create the WikiPage objects for each of those links, you would then need to send an HTML GET request to get their body HTML. Such a network call would require the IO monad (or some other MonadIO), so you're function will necessarily look like:

getWikiNeighbors :: WikiPage -> IO [Wikipage]

But if you successfully implement that function, it's very easy to apply dijkstraM because the "cost" of each hop is always 1!

findShortestWikiPath :: Text -> Text -> IO (Maybe (Int, [WikiPage]))
findShortestWikiPath start end = do
  firstPage <- findWikiPageFromTitle start
  dijkstraM getWikiNeighbors (\_ _ -> return 1) (return . (== end)) firstPage

findWikiPageFromTitle :: Text -> IO WikiPage
...

Of course, because the cost is always 1 this is actually a case where breadth first search would work more simply than Dijkstra's algorithm, so you could use the function bfsM from the same library!

Logging

Another common context for problem solving is the logging context. While we are solving our problem, we might want to record helpful statements telling us what is happening so that we can debug when things are going wrong. This happens using the MonadLogger typeclass, with a few interesting functions we can use, indicating different "levels" of logging.

class MonadLogger m where
  ...

logDebugN :: (MonadLogger m) => Text -> m ()
logInfoN :: (MonadLogger m) => Text -> m ()
logWarnN :: (MonadLogger m) => Text -> m ()
logErrorN :: (MonadLogger m) => Text -> m ()

Now, unlike the previous two examples, this doesn't require the IO monad. A couple of the most common implementations of this monad class will, in fact, use IO functionality (printing to the screen or logging to a file). But this isn't necessary. You can still do logging in a "pure" way by storing the log messages in a sequence or other structure so you can examine them at the end of your program.

When would we want this for Dijkstra's algorithm? Well, sometimes the process of determining neighbors and costs can be complicated! I'll motivate this by introducing a more complicated example of a Dijkstra's algorithm problem.

A Complicated Example

Here's an example from last year's Advent of Code challenge. You can read the full description on that page. This problem demonstrates a less intuitive use of Dijkstra's algorithm.

The problem input is a "map" of sorts, showing a diagram of 4 rooms leading into one shared hallway.

#############
#...........#
###B#C#B#D###
  #A#D#C#A#
  #########

Each of the four rooms is filled with "tokens", which come in 4 different varieties, A, B, C, D. (The Advent of Code description refers to them as "Amphipods", but that takes a while to write out, so I'm simplifying to "tokens").

We want to move the tokens around so that the A tokens end in the far left room, the B tokens in the room next to them, and so on.

#############
#...........#
###A#B#C#D###
  #A#B#C#D#
  #########

But there are rules on how these tokens move. You can only move each token twice. Once to get it into an empty space in the hallway, and once to get it from the hallway to its final room. And tokens can't move "past" each other within the hallway.

Now each token has a specific cost for each space it moves.

A = 1 energy per move
B = 10 energy per move
C = 100 energy per move
D = 1000 energy per move

So you want to move the token's into the final state with the lowest total cost.

Using Dijkstra's Algorithm

It turns out the most efficient solution (especially at a larger scale) is to treat this like a graph problem and use Dijkstra's algorithm! Each "state" of the problem is like a node in our graph, and we can move to certain "neighboring" nodes by moving tokens at a certain cost.

But the implementation turns out to be quite tricky! To give you an idea of this, here are some of the data type names and functions I came up with. 

data Token = ...

data HallSpace = ...

data TokenGraphState = ...

tokenEdges :: TokenGraphState -> [(TokenGraphState, Int)]

updateStateWithMoveFromRoom :: Token -> HallSpace -> Int -> TokenGraphState -> (TokenGraphState, Int)

updateStateWithMoveFromHall :: Token -> HallSpace -> Int -> TokenGraphState -> (TokenGraphState, Int)

validMovesToHall :: Token -> TokenGraphState -> [(HallSpace, Int)]

validMoveToRoom :: TokenGraphState -> (HallSpace, TokenGraphState -> Maybe Token) -> Maybe (Int, Token, HallSpace)

And these are just the functions with complex logic! There are even a few more simple helpers beyond this!

But when I ran this implementation, I didn't get the right answer! So how could I learn more about my solution and figure out what's going wrong? Unit testing and applying a formal debugger would be nice, but simply being able to print out what is going on in the problem is a quicker way to get started.

Haskell doesn't let you (safely) print from pure functions like I've written above, nor can you add values to a global logging state. So we can fix this by modifying the type signatures to instead use a MonadLogger constraint.

tokenEdges :: (MonadLogger m) => TokenGraphState -> m [(TokenGraphState, Int)]

updateStateWithMoveFromRoom :: (MonadLogger m) => Token -> HallSpace -> Int -> TokenGraphState -> (TokenGraphState, Int)

updateStateWithMoveFromHall :: (MonadLogger m) => Token -> HallSpace -> Int -> TokenGraphState -> m (TokenGraphState, Int)

validMovesToHall :: (MonadLogger m) => Token -> TokenGraphState -> m [(HallSpace, Int)]

validMoveToRoom :: (MonadLogger m) => TokenGraphState -> (HallSpace, TokenGraphState -> Maybe Token) -> m (Maybe (Int, Token, HallSpace))

Now it's simple enough to modify a function to give us some important information about what's happening. Hopefully this is enough to help us solve the problem.

We would like to limit the number of functions that "need" the monadic action. But in practice, it is frustrating to find you need a monad in a deeper function of your algorithm because you'll need to modify everything on its call stack. So it might be a good idea to add at least a basic monad constraint from the beginner!

(Update: I did a full implementation of this particular problem that incorporates the logger monad!)

Conclusion

Next time on the blog, we'll start talking more generally about this idea of using monads to debug, especially MonadLogger. We'll consider the implementation pattern of "monads first" and different ways to approach this.

Make sure you're staying up to date with the latest news from Monday Morning Haskell by subscribing to our mailing list! This will also give you access to our subscriber resources!

James Bowen
Previous

My New Favorite Monad?
Next

Dijkstra Comparison: Looking at the Library Function