Dijkstra

Dijkstra.hs

module Dijkstra where

import qualified Data.Graph.Inductive.Graph as G
import Data.Graph.Inductive.Graph (mkGraph, labNodes, LEdge)
import Data.Graph.Inductive.PatriciaTree (Gr)
import qualified Data.Array as Ar
import qualified Data.Map as M
import qualified Data.Set as S
import qualified Control.Monad.State as St
import Control.Monad.State (State, runState)
import Control.Monad (foldM)
import Data.Maybe (listToMaybe)
import qualified Data.List as L

-- ====== Main program demonstrating the example graph from wikipedia ======

graphResults :: IO ()
graphResults = do
  putStrLn "The results of finding the shortest distance on this graph:"
  newLine
  putStrLn $ show eg1Graph
  newLine
  putStrLn "...are:"
  putStrLn $ show eg1

newLine :: IO ()
newLine = putStrLn ""

-- ====== Types ======

type Vertex = Int
type EdgeLength = Int

data Edge = Edge
  { es_from_vertex :: Vertex
  , es_to_vertex :: Vertex
  , es_length :: EdgeLength
  }

newtype NodeLabel = NodeLabel Int
  deriving (Show)
type Graph = Gr NodeLabel EdgeLength

-- a distance can either be inifitely far away, or a countable distance:
data Distance = CountableD Int | InfiniteD
  deriving (Ord, Eq, Show)

-- the data type that represents a single step within state manipulations
-- to find the shortest path
data StepState =
  StepState
    { ss_graph :: Graph
    , ss_unvisited_set :: S.Set Vertex
    , ss_distances :: M.Map Vertex Distance
    , ss_parent_nodes :: M.Map Vertex Vertex
    , ss_current_node :: Vertex
    , ss_target_node :: Vertex
    } deriving (Show)

-- ====== Setup functions for Graph ======

genGraph :: [Edge] -> Graph
genGraph edges = mkGraph nodes labeledGraphEdges
  where
    numNodes = nodeCountFromEdges edges
    nodes = (\i -> (i, NodeLabel i)) <$> [1..numNodes]
    labeledGraphEdges = labeledGraphEdgeFromEdge <$> edges
    labeledGraphEdgeFromEdge :: Edge -> (Vertex, Vertex, EdgeLength)
    labeledGraphEdgeFromEdge edge =
      (es_from_vertex edge, es_to_vertex edge, es_length edge)

nodeCountFromEdges :: [Edge] -> Int
nodeCountFromEdges = length . nodesFromEdges
  where
    nodesFromEdges :: [Edge] -> [Vertex]
    nodesFromEdges = L.nub . edgesToVertices
    edgesToVertices :: [Edge] -> [Vertex]
    edgesToVertices edges = [es_from_vertex, es_to_vertex] <*> edges

-- ====== Example Data ======

-- example graph
eg1Graph :: Graph
eg1Graph = genGraph
  [ Edge 1 2 7
  , Edge 1 3 9
  , Edge 1 6 14
  , Edge 2 3 10
  , Edge 2 4 15
  , Edge 3 4 11
  , Edge 3 6 2
  , Edge 4 5 6
  , Edge 5 6 9
  ]

eg1InitStepState :: StepState
eg1InitStepState = initStepState eg1Graph 1 5

eg1 :: [Vertex]
eg1 = St.evalState shortestPath eg1InitStepState

-- ====== Functions to set up initial state ======

-- build an initial StepState from a graph and initial and target
-- vertices, which will then be used to find the shortest path
initStepState :: Graph -> Vertex -> Vertex -> StepState
initStepState graph initialVertex targetVertex =
    StepState
      { ss_graph = graph
      , ss_unvisited_set = set
      , ss_distances = distances
      , ss_parent_nodes = M.empty
      , ss_current_node = initialVertex
      , ss_target_node = targetVertex
      }
  where
    vertices = verticesFromGraph graph
    set = S.fromList vertices
    distances =
      let dists = M.fromList $ zip vertices $ repeat InfiniteD
      in M.insert initialVertex (CountableD 0) dists

-- ====== The main State action; uses various other State actions ======

-- State action that finds the shortest path (as a list of vertices)
-- given an initial StepState
shortestPath :: State StepState [Vertex]
shortestPath = do
  neighboursWithTentativeDistances <- getNeighboursWithTentativeDistances
  updateDistancesAndParentsIfShorter neighboursWithTentativeDistances
  removeCurrentNodeFromUnvistedSet
  setNewCurrentNodeAsMinimumDistanceNeighbourFrom neighboursWithTentativeDistances
  maybePath <- getPathIfDone
  case maybePath of
    Nothing -> shortestPath
    Just path -> pure path

-- ====== Some helper functions ======

addDistances :: Distance -> Distance -> Distance
addDistances (CountableD x) (CountableD y) = (CountableD $ x + y)
addDistances _ _ = InfiniteD

(&+&) :: Distance -> Distance -> Distance
(&+&) = addDistances

currentNode :: StepState -> Vertex
currentNode = ss_current_node

verticesFromGraph :: Graph -> [Vertex]
verticesFromGraph graph =
  fst <$> labNodes graph

-- ====== State actions used to compose the main State action  ======

getPathIfDone :: State StepState (Maybe [Vertex])
getPathIfDone = do
  state <- St.get
  targetNodeVisited <- getTargetNodeVisited
  smallestDistanceOfNodesIsInfinity <- getSmallestDistanceOfNodesIsInfinity
  pure $
    if targetNodeVisited || smallestDistanceOfNodesIsInfinity
    then
      pure $ parentPathFromParentMapAndTarget (ss_parent_nodes state) (ss_target_node state)
    else
      Nothing

getTargetNodeVisited :: State StepState Bool
getTargetNodeVisited = do
  state <- St.get
  pure $ S.notMember (ss_target_node state) (ss_unvisited_set state)

getSmallestDistanceOfNodesIsInfinity :: State StepState Bool
getSmallestDistanceOfNodesIsInfinity = do
  tentativeDistances <- M.elems . ss_distances <$> St.get
  pure $ minimum tentativeDistances == InfiniteD

parentPathFromParentMapAndTarget ::  M.Map Vertex Vertex -> Vertex -> [Vertex]
parentPathFromParentMapAndTarget parentMap targetNode =
    L.reverse $ targetNode : L.unfoldr go targetNode
  where
    go previousNode =
      let maybeFoundParent = M.lookup previousNode parentMap
      in (\foundParent -> (foundParent, foundParent)) <$> maybeFoundParent

updateDistancesAndParentsIfShorter :: [(Vertex, Distance)] -> State StepState ()
updateDistancesAndParentsIfShorter neighboursWithTentativeDistances =
    St.modify $ \state ->
      let currentNode = ss_current_node state
      in foldr (go currentNode) state neighboursWithTentativeDistances
  where
    go currentNode (vertex, tentativeDist) prevState =
      let distances = ss_distances prevState
          needToUpdate =
            maybe True (\distance -> distance > tentativeDist) (M.lookup vertex distances)
      in
        if needToUpdate
          then
            prevState
              { ss_distances = M.insert vertex tentativeDist distances
              , ss_parent_nodes = M.insert vertex currentNode (ss_parent_nodes prevState)
              }
          else
            prevState

setNewCurrentNodeAsMinimumDistanceNeighbourFrom :: [(Vertex, Distance)] -> State StepState ()
setNewCurrentNodeAsMinimumDistanceNeighbourFrom verticesWithDistances =
  St.modify $ \state ->
    let maybeVertex = fst <$>
         case verticesWithDistances of
           [] -> Nothing
           _ -> Just $ L.minimumBy (\x y -> compare (snd x) (snd y)) verticesWithDistances
    in case maybeVertex of
      Just vertex -> state { ss_current_node = vertex }
      Nothing -> state

removeCurrentNodeFromUnvistedSet :: State StepState ()
removeCurrentNodeFromUnvistedSet =
  St.modify $ \state ->
    let currentNode = ss_current_node state
        unvisitedSet = ss_unvisited_set state
    in state { ss_unvisited_set = S.delete currentNode unvisitedSet }


getNeighboursWithTentativeDistances :: State StepState [(Vertex, Distance)]
getNeighboursWithTentativeDistances = do
  currentNodeDistance <- getCurrentNodeDistance
  state <- St.get
  let
    graph :: Graph
    graph = ss_graph state
    currentNode :: Vertex
    currentNode = ss_current_node state
    neighboursAndLengths :: [(EdgeLength, Vertex)]
    neighboursAndLengths = G.lneighbors graph currentNode
    unvisitedSet :: S.Set Vertex
    unvisitedSet = ss_unvisited_set state
    unvisitedNeighbours = filter (\(_, vertex) -> S.member vertex unvisitedSet) neighboursAndLengths
    lengthVertexPairToVertexCurrentNodeDistancePair (edgeLength, vertex) =
      (vertex, currentNodeDistance &+& (CountableD edgeLength))
    neighboursWithTentativeDistances = fmap lengthVertexPairToVertexCurrentNodeDistancePair unvisitedNeighbours
  pure neighboursWithTentativeDistances

getCurrentNodeDistance :: State StepState Distance
getCurrentNodeDistance = do
  stepState <- St.get
  let
    currentNode = ss_current_node stepState
    distances = ss_distances stepState
    currentNodeDistance = distances M.! currentNode
  pure currentNodeDistance