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
7 import Control.Monad.ST
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)
19 data (Graph gr, Real w) => BFState s gr a b w = BFState {
22 bfsMPaths :: STArray s Node (Maybe (BFPath b w)),
23 bfsChanged :: ArrayQueue s,
24 -- Sentinel in the queue that ends a pass.
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
32 negativeCycleCheck :: (Graph gr, Real w) => BFState s gr a b w ->
34 negativeCycleCheck state path =
35 let getNodes (BFPath _ dst from) = dst : case from of
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)
44 offerPath :: (Graph gr, Real w) => BFState s gr a b w ->
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
51 Just (BFPath oldLen _ _) -> newLen < oldLen
55 writeArray (bfsMPaths state) dest (Just newPath)
56 -- Mark dest as changed.
57 aqEnqueue (bfsChanged state) dest
59 else nop -- Don't update anything.
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
68 endPass :: (Graph gr, Real w) => BFState s gr a b w -> ST s ()
70 qEmpty <- aqIsEmpty (bfsChanged state)
72 then nop -- No nodes to visit on the next pass. We're done.
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!"
80 -- Re-enqueue the pass ender.
81 aqEnqueue (bfsChanged state) (bfsPassEnder state)
82 search state -- Continue with the next pass.
84 search :: (Graph gr, Real w) => BFState s gr a b w -> ST s ()
86 Just node <- aqDequeue (bfsChanged state)
87 if node == bfsPassEnder state
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.
96 bellmanFord :: (Graph gr, Real w) =>
97 (b -> w) -> -- Edge label -> weight
98 Node -> -- Source node
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.
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)
112 return (bfsMPaths state)