| 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 | ) |