Feedback on simple DFS

Question

I am implementing a simple recursive Depth-First Search algorithm on graphs, using StateT and ReaderT for encapsulating all the required information.

Pre-requisite definitions

type Vertex = Int
type Cost = Int
data Edge = Edge { vertex:: Vertex, cost:: Cost } deriving (Eq, Show, Ord)

newtype Graph = Graph { adjacencyList :: Map.Map Vertex [Edge] } deriving (Eq, Show)

type DFSState = ([Vertex], [Vertex])
type Env = Graph
type DFSNextState = ReaderT Env (StateT DFSState Identity) Vertex

and then the main function:

{-# LANGUAGE FlexibleContexts #-}

dfsGetNext :: DFSNextState
dfsGetNext = do
    graph <- ask
    s <- get
    let candidates = fst s
    let visited = snd s
    let ret = case candidates of
                []     -> return 0 :: DFSNextState
                (x:xs) -> if elem x visited 
                            then put (xs, visited) >> (return 1 :: DFSNextState)
                          else put ((getNeighbors x graph) ++ xs, visited ++ [x]) >> (return 1 :: DFSNextState)
    ret

I can then simply run the algorithm for n steps and get the nodes visited in a depth-first fashion using

dfsStartState :: DFSState
dfsStartState = ([0], [])

dfsSteps = replicateM 13 dfsGetNext

runStateT (runReaderT dfsSteps graph) dfsStartState

Questions:
1) Is my style in the main function good? Can it be simplified?
2) Is there a way to remove the magic number (13 in this case) for replicateM? Is there a replicate until which I could use for this?

Update I added a replicateMUntil function myself:

{-# LANGUAGE RankNTypes, KindSignatures #-}
replicateMUntil :: forall (m :: * -> *) a. Monad m => (m a -> m Bool) -> m a -> m [a]
replicateMUntil f ma = do
    res <- f ma
    if res then liftM2 (:) ma (replicateMUntil f ma) else return []

The driver now becomes:

isNextAvailable :: DFSNextState -> ReaderT Env (StateT DFSState Identity) Bool
isNextAvailable nextState = do
    s <- get
    return (fst s /= [])

dfsStartState :: DFSState
dfsStartState = ([0], [])

dfsSteps = replicateMUntil isNextAvailable dfsGetNext

runStateT (runReaderT dfsSteps graph) dfsStartState

Answer 1

The small things

    s <- get
    let candidates = fst s
    let visited = snd s

could be simplified to

    (candidates, visited) <- get

(It is technically a little bit stricter but it's not a problem because you don't rely on laziness here.)

    let ret = blah in
    ret

is equivalent to

    blah

At the end of dfsGetNext you return either 0 or 1, which is strange because the return type is Vertex, and although it is a synonym for Int (so this type checks), 0 and 1 don't look like vertices.

My guess is that you intended either:

• to return a boolean to indicate whether the DFS ended, so Bool would be a more appropriate type;
• to return the visited vertex, so, rather than returning dummy sentinel values, Maybe Vertex would be a more appropriate type.

Assuming that 0 and 1 are meant to be booleans, and with the other changes above, we get:

-- Also isolating the monad, and applying it to the result type in the function instead
type DFS = ReaderT Env (StateT DFSState Identity)

dfsGetNext :: DFS Bool
dfsGetNext = do
    graph <- ask
    (candidates, visited) <- get
    case candidates of
       []     -> return False
       (x:xs) -> if elem x visited 
                 then put (xs, visited) >> return True
                 else put ((getNeighbors x graph) ++ xs, visited ++ [x]) >> return True

Getting fancy

There's a more direct implementation of DFS. To do a DFS starting from a vertex v:

1. If v has already been visited, stop; otherwise...
2. Visit v
3. Recursively run a DFS from each neighbor of v

In particular, this description doesn't require you to explicitly keep track of "pending candidate vertices", nor to iterate a "single step" function like dfsGetNext; the DFS is directly defined.

-- The state now only contains the visited vertices, not the candidates (which is implicit in the recursive structure of dfsFrom)
type DFS = ReaderT Graph (State Visited)

dfsFrom :: Vertex -> DFS ()
dfsFrom v = do
  visited <- get
  if v `elem` visited then                     -- 1
    pure ()
  else do
    put (visited ++ [v])                       -- 2
    graph <- ask
    let neighbors = getNeighbors v graph
    for_ neighbors dfsFrom                     -- 3

Now you can "run" that DFS, interpreting it as a function as follows:

-- dfs from 0
dfs :: Graph -> [Vertex]
dfs graph = dfsFrom 0 `runReaderT` graph `execState` []

Your "single step" version does have one distinguishing feature, which is that you can easily interrupt the DFS at any time to inspect the state for example. But you can also do that with the recursive version, by passing a function (visit below) that will be called when visiting a node. Clients of the DFS are free to define it however they wish.

It's also useful to generalize the DFS monad to support additional effects relevant to users rather than to the DFS itself. This is straightforward since you're already using monad transformers, by changing the base monad from Identity to something else. Here I chose IO for simplicity and concreteness, but it might as well be a type parameter m.

type DFSIO = ReaderT Graph (StateT Visited IO)

dfsFromIO :: (Vertex -> DFSIO ()) -> Vertex -> DFSIO ()
dfsFromIO visit v = do
  visited <- get
  if v `elem` visited then
    pure ()
  else do
    visit v
    put (visited ++ [v])
    graph <- ask
    let neighbors = getNeighbors v graph
    for_ neighbors (dfsFromIO visit)

Another possibility is to call visit at the very start of the function, before checking whether the vertex was visited, or at the very end, on the neighboring vertices, where the function could also have access to the cost of the edges being crossed.

Running this new version dfsFromIO is as straightforward as running dfsFrom (only replacing execState with execStateT since the base monad is no longer Identity):

-- dfs from 0
dfsIO :: (Vertex -> DFSIO ()) -> Graph -> IO [Vertex]
dfsIO visit graph = dfsFromIO visit 0 `runReaderT` graph `execStateT` []

For example you can print every visited vertex as follows:

main :: IO ()
main = void (dfsIO (\v -> liftIO (print v)))

---

Full gist: https://gist.github.com/Lysxia/4733f8d670c84e5b886f74d2d106d1f5