{-# LANGUAGE FlexibleContexts #-}
{-|
== Floyd-Warshall algorithm

Finds shortest paths between all pairs of vertices in a directed weighted graph.

Sources:

* https://en.wikipedia.org/wiki/Floyd%E2%80%93Warshall_algorithm
* Robert W. Floyd, "Algorithm 97: Shortest Path", Communications of the ACM, Vol. 5, No. 6, 1962
  https://dl.acm.org/doi/10.1145/367766.368168

-}

{-
Implementation notes:
* loop l r is used instead of for_ [l..r] because [l..r] gets let-floated out of the nested loops,
  preventing it from being optimizing away.
-}

module FloydWarshall
    ( WEdge(..)
    , Weight
    , floydWarshallFromEdges
    , floydWarshall
    ) where

import Control.Monad
import Data.Array.ST
import Data.Array.Unboxed
import Data.Graph

import Misc ( modifyArray )

type Weight = Int
data WEdge = WEdge !Vertex !Vertex !Weight deriving (WEdge -> WEdge -> Bool
(WEdge -> WEdge -> Bool) -> (WEdge -> WEdge -> Bool) -> Eq WEdge
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: WEdge -> WEdge -> Bool
$c/= :: WEdge -> WEdge -> Bool
== :: WEdge -> WEdge -> Bool
$c== :: WEdge -> WEdge -> Bool
Eq, Int -> WEdge -> ShowS
[WEdge] -> ShowS
WEdge -> String
(Int -> WEdge -> ShowS)
-> (WEdge -> String) -> ([WEdge] -> ShowS) -> Show WEdge
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [WEdge] -> ShowS
$cshowList :: [WEdge] -> ShowS
show :: WEdge -> String
$cshow :: WEdge -> String
showsPrec :: Int -> WEdge -> ShowS
$cshowsPrec :: Int -> WEdge -> ShowS
Show)

-- | Runs the algorithm on the graph represented by the given list of edges. There should be no negative
-- cycles. O(|V|^3). Queries take O(1).
floydWarshallFromEdges :: Bounds -> [WEdge] -> Vertex -> Vertex -> Maybe Weight
floydWarshallFromEdges :: Bounds -> [WEdge] -> Int -> Int -> Maybe Int
floydWarshallFromEdges (Int
l, Int
r) [WEdge]
es = Int -> Int -> Maybe Int
forall (f :: * -> *). Alternative f => Int -> Int -> f Int
qry where
    qry :: Int -> Int -> f Int
qry Int
u Int
v = let x :: Int
x = UArray Bounds Int
daUArray Bounds Int -> Bounds -> Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!(Int
u, Int
v) in Int
x Int -> f () -> f Int
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool -> f ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
forall a. Bounded a => a
maxBound)
    da :: UArray Bounds Int
da = (forall s. ST s (STUArray s Bounds Int)) -> UArray Bounds Int
forall i e. (forall s. ST s (STUArray s i e)) -> UArray i e
runSTUArray ((forall s. ST s (STUArray s Bounds Int)) -> UArray Bounds Int)
-> (forall s. ST s (STUArray s Bounds Int)) -> UArray Bounds Int
forall a b. (a -> b) -> a -> b
$ do
        STUArray s Bounds Int
d <- (Bounds, Bounds) -> Int -> ST s (STUArray s Bounds Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray ((Int
l, Int
l), (Int
r, Int
r)) Int
forall a. Bounded a => a
maxBound
        [WEdge] -> (WEdge -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [WEdge]
es ((WEdge -> ST s ()) -> ST s ()) -> (WEdge -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(WEdge Int
u Int
v Int
w) -> STUArray s Bounds Int -> Bounds -> (Int -> Int) -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray STUArray s Bounds Int
d (Int
u, Int
v) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
w)
        [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
l..Int
r] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> STUArray s Bounds Int -> Bounds -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Bounds Int
d (Int
i, Int
i) Int
0
        STUArray s Bounds Int -> ST s ()
forall (a :: * -> * -> *) (m :: * -> *).
MArray a Int m =>
a Bounds Int -> m ()
floydWarshall STUArray s Bounds Int
d
        STUArray s Bounds Int -> ST s (STUArray s Bounds Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure STUArray s Bounds Int
d

-- | Runs the algorithm on the adjacency matrix of the graph given as a mutable array. O(|V|^3). 
floydWarshall :: MArray a Weight m => a (Vertex, Vertex) Weight -> m ()
floydWarshall :: a Bounds Int -> m ()
floydWarshall a Bounds Int
d = do
    ((Int
l, Int
_), (Int
r, Int
_)) <- a Bounds Int -> m (Bounds, Bounds)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
getBounds a Bounds Int
d
    Int -> Int -> (Int -> m ()) -> m ()
forall (f :: * -> *) t a.
(Ord t, Monad f, Num t) =>
t -> t -> (t -> f a) -> f ()
loop Int
l Int
r ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
k ->
        Int -> Int -> (Int -> m ()) -> m ()
forall (f :: * -> *) t a.
(Ord t, Monad f, Num t) =>
t -> t -> (t -> f a) -> f ()
loop Int
l Int
r ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
            Int
ik <- a Bounds Int -> Bounds -> m Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray a Bounds Int
d (Int
i, Int
k)
            Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ik Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
forall a. Bounded a => a
maxBound) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> (Int -> m ()) -> m ()
forall (f :: * -> *) t a.
(Ord t, Monad f, Num t) =>
t -> t -> (t -> f a) -> f ()
loop Int
l Int
r ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
                Int
kj <- a Bounds Int -> Bounds -> m Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray a Bounds Int
d (Int
k, Int
j)
                Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
kj Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
forall a. Bounded a => a
maxBound) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
                    a Bounds Int -> Bounds -> (Int -> Int) -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray a Bounds Int
d (Int
i, Int
j) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
ik Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kj))
  where
    loop :: t -> t -> (t -> f a) -> f ()
loop t
l t
r t -> f a
f = t -> f ()
go t
l where go :: t -> f ()
go t
i = Bool -> f () -> f ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (t
i t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
r) (f () -> f ()) -> f () -> f ()
forall a b. (a -> b) -> a -> b
$ t -> f a
f t
i f a -> f () -> f ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> t -> f ()
go (t
i t -> t -> t
forall a. Num a => a -> a -> a
+ t
1)

--------------------------------------------------------------------------------
-- For tests

-- Allows specialization across modules
{-# INLINABLE floydWarshall #-}