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