program/BellmanFord.hs

   1 module BellmanFord (bellmanFord, BFPath(BFPath)) where
   2 import Data.Graph.Inductive.Graph
   3 import Data.Graph.Inductive.Internal.Queue
   4 import Data.Array.IArray
   5 import Data.Array.ST
   6 import Data.STRef
   7 import Control.Monad.ST
   8 import ArrayQueue
   9 import MonadStuff
  10 import Data.List
  11
  12 -- Path to a node
  13 data Real w => BFPath b w = BFPath {
  14         pLen  :: w,                    -- Total distance at the end of the path
  15         pDest :: Node,                 -- Destination of the path
  16         pFrom :: Maybe (b, BFPath b w) -- Last edge and remaining path (Nothing for source)
  17 } deriving Show
  18
  19 data (Graph gr, Real w) => BFState s gr a b w = BFState {
  20         bfsGraph :: gr a b,
  21         bfsEdgeWt :: b -> w,
  22         bfsMPaths :: STArray s Node (Maybe (BFPath b w)),
  23         bfsChanged :: ArrayQueue s,
  24         -- Sentinel in the queue that ends a pass.
  25         bfsPassEnder :: Node,
  26         -- The length of the paths we are currently considering.
  27         -- If this reaches the number of nodes in the graph, we must have a negative cycle.
  28         bfsPass :: STRef s Int
  29 }
  30
  31 {--
  32 negativeCycleCheck :: (Graph gr, Real w) => BFState s gr a b w ->
  33         BFPath b w -> ST s ()
  34 negativeCycleCheck state path =
  35         let getNodes (BFPath _ dst from) = dst : case from of
  36                 Nothing -> []
  37                 Just (_, p1) -> getNodes p1 in
  38         let nodes = getNodes path in
  39         if length (nub nodes) < length nodes
  40         then error ("Negative cycle detected: " ++ show path)
  41         else nop
  42 --}
  43
  44 offerPath :: (Graph gr, Real w) => BFState s gr a b w ->
  45         BFPath b w -> ST s ()
  46 offerPath state newPath@(BFPath newLen dest _) = do
  47         oldMPath <- readArray (bfsMPaths state) dest
  48         -- Is newPath the first path to dest or shorter than the existing one?
  49         let adoptPath = case oldMPath of
  50                 Nothing -> True
  51                 Just (BFPath oldLen _ _) -> newLen < oldLen
  52         if adoptPath
  53                 then do
  54                         -- Save the new path.
  55                         writeArray (bfsMPaths state) dest (Just newPath)
  56                         -- Mark dest as changed.
  57                         aqEnqueue (bfsChanged state) dest
  58                         nop
  59                 else nop -- Don't update anything.
  60
  61 processEdge :: (Graph gr, Real w) => BFState s gr a b w ->
  62         BFPath b w -> LEdge b -> ST s ()
  63 processEdge state path1@(BFPath len1 _ _) (_, dest, label) = do
  64         let edgeLen = (bfsEdgeWt state) label
  65         let newPath = BFPath (len1 + edgeLen) dest (Just (label, path1))
  66         offerPath state newPath
  67
  68 endPass :: (Graph gr, Real w) => BFState s gr a b w -> ST s ()
  69 endPass state = do
  70         qEmpty <- aqIsEmpty (bfsChanged state)
  71         if qEmpty
  72                 then nop -- No nodes to visit on the next pass.  We're done.
  73                 else do
  74                         -- Increment the pass number.
  75                         modifySTRef (bfsPass state) (+ 1)
  76                         p <- readSTRef (bfsPass state)
  77                         if p == noNodes (bfsGraph state)
  78                                 then error "BellmanFord: Negative cycle detected!"
  79                                 else do
  80                                         -- Re-enqueue the pass ender.
  81                                         aqEnqueue (bfsChanged state) (bfsPassEnder state)
  82                                         search state -- Continue with the next pass.
  83
  84 search :: (Graph gr, Real w) => BFState s gr a b w -> ST s ()
  85 search state = do
  86         Just node <- aqDequeue (bfsChanged state)
  87         if node == bfsPassEnder state
  88                 then endPass state
  89                 else do
  90                         mPath1 <- readArray (bfsMPaths state) node
  91                         let Just path1 = mPath1
  92                         sequence $ map (processEdge state path1) $
  93                                 out (bfsGraph state) node
  94                         search state -- Keep going.
  95
  96 bellmanFord :: (Graph gr, Real w) =>
  97         (b -> w) ->                     -- Edge label -> weight
  98         Node ->                         -- Source node
  99         gr a b ->                       -- Graph
 100         Array Node (Maybe (BFPath b w)) -- ! node -> maybe a path
 101 bellmanFord edgeWt source theGraph = runSTArray (do
 102                 mPaths <- newArray (nodeRange theGraph) Nothing
 103                 let (nlo, nhi) = nodeRange theGraph
 104                 let passEnder = nlo - 1
 105                 changed <- newArrayQueue (passEnder, nhi) -- Queue can contain the pass ender.
 106                 pass <- newSTRef 0
 107                 let state = BFState theGraph edgeWt mPaths changed passEnder pass
 108                 -- Pass is 0, and the queue contains a node that was offered a path with 0 edges.
 109                 offerPath state (BFPath 0 source Nothing)
 110                 aqEnqueue (bfsChanged state) (bfsPassEnder state)
 111                 search state
 112                 return (bfsMPaths state)
 113         )