Commit | Line | Data |
---|---|---|
5a07db44 | 1 | module BellmanFord (bellmanFord, BFPath(BFPath)) where |
d7d9561e MM |
2 | import Data.Graph.Inductive.Graph |
3 | import Data.Graph.Inductive.Internal.Queue | |
5a07db44 MM |
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 | |
d7d9561e | 11 | |
5a07db44 MM |
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 | |
d7d9561e MM |
29 | } |
30 | ||
5a07db44 MM |
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. | |
d7d9561e | 60 | |
5a07db44 MM |
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 | |
d7d9561e | 67 | |
5a07db44 MM |
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. | |
d7d9561e | 83 | |
5dcb8125 | 84 | search :: (Graph gr, Real w) => BFState s gr a b w -> ST s () |
5a07db44 MM |
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. | |
d7d9561e | 95 | |
5a07db44 MM |
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 | ) |