{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

-- |
-- Module: Data.DiGraph.FloydWarshall
-- Copyright: Copyright © 2018-2021 Kadena LLC.
-- License: MIT
-- Maintainer: Lars Kuhtz <lars@kadena.io>
-- Stability: experimental
--
-- TODO
--
module Data.DiGraph.FloydWarshall
(
-- * Graph Representations
  DenseAdjMatrix
, AdjacencySets

-- * Conversions
, fromAdjacencySets
, toAdjacencySets

-- * FloydWarshall Algorithm
, ShortestPathMatrix(..)
, floydWarshall
, shortestPath
, distance
, diameter

-- * Legacy exports
, distMatrix_
, floydWarshall_
, diameter_
, shortestPaths_
) where

import Control.DeepSeq

import Data.Foldable
import qualified Data.HashMap.Strict as HM
import qualified Data.HashSet as HS
import Data.Massiv.Array as M

import GHC.Generics

import Numeric.Natural

-- -------------------------------------------------------------------------- --
-- Graph Representations

-- | Adjacency matrix representation of a directed graph.
--
type DenseAdjMatrix = Array U Ix2 Int

-- | Adjacency set representation of a directed graph.
--
type AdjacencySets = HM.HashMap Int (HS.HashSet Int)

-- -------------------------------------------------------------------------- --
-- Adjacency Matrix for dense Graphs
--
-- (uses massiv)

-- | Assumes that the input is an directed graph and that the vertex
-- set is a prefix of the natural numbers.
--
fromAdjacencySets :: AdjacencySets -> DenseAdjMatrix
fromAdjacencySets :: AdjacencySets -> DenseAdjMatrix
fromAdjacencySets AdjacencySets
g = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq (forall ix. Index ix => ix -> Sz ix
Sz (Int
n Int -> Int -> Ix2
:. Int
n)) forall {a}. Num a => Ix2 -> a
go
  where
    n :: Int
n = forall k v. HashMap k v -> Int
HM.size AdjacencySets
g
    go :: Ix2 -> a
go (Int
i :. Int
j)
        | (Int, Int) -> Bool
isEdge (Int
i, Int
j) = a
1
        | Bool
otherwise = a
0
    isEdge :: (Int, Int) -> Bool
isEdge (Int
a, Int
b) = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
HS.member Int
b) forall a b. (a -> b) -> a -> b
$ forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup Int
a AdjacencySets
g

-- | Converts an adjacency matrix into a graph in adjacnency set representation.
--
toAdjacencySets :: DenseAdjMatrix -> AdjacencySets
toAdjacencySets :: DenseAdjMatrix -> AdjacencySets
toAdjacencySets = forall ix r e a.
(Index ix, Source r e) =>
(a -> ix -> e -> a) -> a -> Array r ix e -> a
ifoldlS forall {a}.
(Eq a, Num a) =>
AdjacencySets -> Ix2 -> a -> AdjacencySets
f forall a. Monoid a => a
mempty
  where
    f :: AdjacencySets -> Ix2 -> a -> AdjacencySets
f AdjacencySets
a (Int
i :. Int
j) a
x
        | a
x forall a. Eq a => a -> a -> Bool
== a
0 = AdjacencySets
a
        | Bool
otherwise = forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HM.insertWith forall a. Semigroup a => a -> a -> a
(<>) Int
i (forall a. Hashable a => a -> HashSet a
HS.singleton Int
j) AdjacencySets
a

-- -------------------------------------------------------------------------- --
-- Floyd Warshall with Paths

-- | Shortest path matrix of a graph.
--
newtype ShortestPathMatrix = ShortestPathMatrix (Array U Ix2 (Double, Int))
    deriving (Int -> ShortestPathMatrix -> ShowS
[ShortestPathMatrix] -> ShowS
ShortestPathMatrix -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ShortestPathMatrix] -> ShowS
$cshowList :: [ShortestPathMatrix] -> ShowS
show :: ShortestPathMatrix -> String
$cshow :: ShortestPathMatrix -> String
showsPrec :: Int -> ShortestPathMatrix -> ShowS
$cshowsPrec :: Int -> ShortestPathMatrix -> ShowS
Show, ShortestPathMatrix -> ShortestPathMatrix -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c/= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
== :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c== :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
Eq, Eq ShortestPathMatrix
ShortestPathMatrix -> ShortestPathMatrix -> Bool
ShortestPathMatrix -> ShortestPathMatrix -> Ordering
ShortestPathMatrix -> ShortestPathMatrix -> ShortestPathMatrix
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ShortestPathMatrix -> ShortestPathMatrix -> ShortestPathMatrix
$cmin :: ShortestPathMatrix -> ShortestPathMatrix -> ShortestPathMatrix
max :: ShortestPathMatrix -> ShortestPathMatrix -> ShortestPathMatrix
$cmax :: ShortestPathMatrix -> ShortestPathMatrix -> ShortestPathMatrix
>= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c>= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
> :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c> :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
<= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c<= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
< :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c< :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
compare :: ShortestPathMatrix -> ShortestPathMatrix -> Ordering
$ccompare :: ShortestPathMatrix -> ShortestPathMatrix -> Ordering
Ord, forall x. Rep ShortestPathMatrix x -> ShortestPathMatrix
forall x. ShortestPathMatrix -> Rep ShortestPathMatrix x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep ShortestPathMatrix x -> ShortestPathMatrix
$cfrom :: forall x. ShortestPathMatrix -> Rep ShortestPathMatrix x
Generic)
    deriving newtype (ShortestPathMatrix -> ()
forall a. (a -> ()) -> NFData a
rnf :: ShortestPathMatrix -> ()
$crnf :: ShortestPathMatrix -> ()
NFData)

-- | Shortest path computation for integral matrixes.
--
floydWarshall :: Unbox a => Real a => Array U Ix2 a -> ShortestPathMatrix
floydWarshall :: forall a. (Unbox a, Real a) => Array U Ix2 a -> ShortestPathMatrix
floydWarshall = Array U Ix2 (Double, Int) -> ShortestPathMatrix
ShortestPathMatrix
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array U Ix2 (Double, Int) -> Array U Ix2 (Double, Int)
floydWarshallInternal
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r e r' ix.
(Manifest r e, Load r' ix e) =>
r -> Array r' ix e -> Array r ix e
computeAs U
U
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a r.
(Real a, Source r a) =>
Array r Ix2 a -> Array D Ix2 (Double, Int)
intDistMatrix

-- | Compute a shortest path between two vertices of a graph from the shortest
-- path matrix of the graph.
--
shortestPath
    :: ShortestPathMatrix
    -> Int
    -> Int
    -> Maybe [Int]
shortestPath :: ShortestPathMatrix -> Int -> Int -> Maybe [Int]
shortestPath (ShortestPathMatrix Array U Ix2 (Double, Int)
m) Int
src Int
trg
    | forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
M.isEmpty DenseAdjMatrix
mat = forall a. Maybe a
Nothing
    | Bool -> Bool
not (forall ix. Index ix => Sz ix -> ix -> Bool
M.isSafeIndex (forall r ix e. Size r => Array r ix e -> Sz ix
size Array U Ix2 (Double, Int)
m) (Int
src Int -> Int -> Ix2
:. Int
trg)) = forall a. Maybe a
Nothing
    | (DenseAdjMatrix
mat forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
src Int -> Int -> Ix2
:. Int
trg)) forall a. Eq a => a -> a -> Bool
== (-Int
1) = forall a. Maybe a
Nothing
    | Bool
otherwise = Int -> Int -> Maybe [Int]
go Int
src Int
trg
  where
    mat :: DenseAdjMatrix
mat = forall r e r' ix.
(Manifest r e, Load r' ix e) =>
r -> Array r' ix e -> Array r ix e
M.computeAs U
U forall a b. (a -> b) -> a -> b
$ forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
M.map forall a b. (a, b) -> b
snd Array U Ix2 (Double, Int)
m
    go :: Int -> Int -> Maybe [Int]
go Int
a Int
b
        | Int
a forall a. Eq a => a -> a -> Bool
== Int
b = forall (m :: * -> *) a. Monad m => a -> m a
return []
        | Bool
otherwise = do
            Int
n <- forall ix r e.
(Index ix, Manifest r e) =>
Array r ix e -> ix -> Maybe e
M.index DenseAdjMatrix
mat (Int
a Int -> Int -> Ix2
:. Int
b)
            (:) Int
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> Maybe [Int]
go Int
n Int
b

-- | Compute the distance between two vertices of a graph from the shortest path
-- matrix of the graph.
--
distance :: ShortestPathMatrix -> Int -> Int -> Maybe Double
distance :: ShortestPathMatrix -> Int -> Int -> Maybe Double
distance (ShortestPathMatrix Array U Ix2 (Double, Int)
m) Int
src Int
trg
    | forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
M.isEmpty Array U Ix2 (Double, Int)
m = forall a. Maybe a
Nothing
    | Bool
otherwise = forall a. RealFrac a => a -> Maybe a
toDistance forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall ix r e.
(Index ix, Manifest r e) =>
Array r ix e -> ix -> Maybe e
M.index Array U Ix2 (Double, Int)
m (Int
src Int -> Int -> Ix2
:. Int
trg)

-- | Compute the diameter of a graph from the shortest path matrix of the graph.
--
diameter :: ShortestPathMatrix -> Maybe Double
diameter :: ShortestPathMatrix -> Maybe Double
diameter (ShortestPathMatrix Array U Ix2 (Double, Int)
m)
    | forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
M.isEmpty Array U Ix2 (Double, Int)
m = forall a. a -> Maybe a
Just Double
0
    | Bool
otherwise = forall a. RealFrac a => a -> Maybe a
toDistance forall a b. (a -> b) -> a -> b
$ forall r ix e.
(HasCallStack, Shape r ix, Source r e, Ord e) =>
Array r ix e -> e
maximum' forall a b. (a -> b) -> a -> b
$ forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
M.map forall a b. (a, b) -> a
fst Array U Ix2 (Double, Int)
m

-- -------------------------------------------------------------------------- --
-- Internal

toDistance :: RealFrac a => a -> Maybe a
toDistance :: forall a. RealFrac a => a -> Maybe a
toDistance a
x
    | a
x forall a. Eq a => a -> a -> Bool
== a
1forall a. Fractional a => a -> a -> a
/a
0 = forall a. Maybe a
Nothing
    | Bool
otherwise = forall a. a -> Maybe a
Just a
x

-- | Distance matrix for int inputs.
--
intDistMatrix
    :: Real a
#if MIN_VERSION_massiv(1,0,0)
    => Source r a
#else
    => Source r Ix2 a
#endif
    => Array r Ix2 a
    -> Array M.D Ix2 (Double, Int)
intDistMatrix :: forall a r.
(Real a, Source r a) =>
Array r Ix2 a -> Array D Ix2 (Double, Int)
intDistMatrix = forall r ix e a.
(Index ix, Source r e) =>
(ix -> e -> a) -> Array r ix e -> Array D ix a
M.imap forall {a} {a}. (Real a, Fractional a) => Ix2 -> a -> (a, Int)
go
  where
    go :: Ix2 -> a -> (a, Int)
go (Int
x :. Int
y) a
e
        | Int
x forall a. Eq a => a -> a -> Bool
== Int
y = (a
0, Int
y)
        | a
e forall a. Ord a => a -> a -> Bool
> a
0 = (forall a b. (Real a, Fractional b) => a -> b
realToFrac a
e, Int
y)
        | Bool
otherwise = (a
1forall a. Fractional a => a -> a -> a
/a
0, -Int
1)

-- | Floyd-Warshall With Path Matrix
--
-- TODO: use a mutable array?
-- TODO: implement Dijkstra's algorithm for adj matrix representation.
--
floydWarshallInternal
    :: Array U Ix2 (Double, Int)
    -> Array U Ix2 (Double,Int)
floydWarshallInternal :: Array U Ix2 (Double, Int) -> Array U Ix2 (Double, Int)
floydWarshallInternal Array U Ix2 (Double, Int)
a = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Array U Ix2 (Double, Int) -> Int -> Array U Ix2 (Double, Int)
go Array U Ix2 (Double, Int)
a [Int
0..Int
nforall a. Num a => a -> a -> a
-Int
1]
  where
    Sz (Int
n :. Int
_) = forall r ix e. Size r => Array r ix e -> Sz ix
size Array U Ix2 (Double, Int)
a

    go :: Array U Ix2 (Double, Int) -> Int -> Array U Ix2 (Double,Int)
    go :: Array U Ix2 (Double, Int) -> Int -> Array U Ix2 (Double, Int)
go Array U Ix2 (Double, Int)
c Int
k = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq (forall ix. Index ix => ix -> Sz ix
Sz (Int
n Int -> Int -> Ix2
:. Int
n)) forall a b. (a -> b) -> a -> b
$ \(Int
x :. Int
y) ->
        let
            !xy :: Double
xy = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$! Array U Ix2 (Double, Int)
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
y)
            !xk :: Double
xk = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$! Array U Ix2 (Double, Int)
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
k)
            !ky :: Double
ky = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$! Array U Ix2 (Double, Int)
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
k Int -> Int -> Ix2
:. Int
y)
            !nxy :: Int
nxy = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$! Array U Ix2 (Double, Int)
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
y)
            !nxk :: Int
nxk = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$! Array U Ix2 (Double, Int)
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
k)
        in if Double
xy forall a. Ord a => a -> a -> Bool
> Double
xk forall a. Num a => a -> a -> a
+ Double
ky then (Double
xk forall a. Num a => a -> a -> a
+ Double
ky, Int
nxk) else (Double
xy, Int
nxy)

-- -------------------------------------------------------------------------- --
-- Floyd Warshall Without Paths (more efficient, by factor of 2)

-- | Floyd Warshall Without Paths (more efficient, by factor of 2).
--
distMatrix_
#if MIN_VERSION_massiv(1,0,0)
    :: Source r Int
#else
    :: Source r Ix2 Int
#endif
    => Array r Ix2 Int
    -> Array M.D Ix2 Double
distMatrix_ :: forall r. Source r Int => Array r Ix2 Int -> Array D Ix2 Double
distMatrix_ = forall r ix e a.
(Index ix, Source r e) =>
(ix -> e -> a) -> Array r ix e -> Array D ix a
M.imap forall {a} {a}. (Real a, Fractional a) => Ix2 -> a -> a
go
  where
    go :: Ix2 -> a -> a
go (Int
x :. Int
y) a
e
        | Int
x forall a. Eq a => a -> a -> Bool
== Int
y = a
0
        | a
e forall a. Ord a => a -> a -> Bool
> a
0 = forall a b. (Real a, Fractional b) => a -> b
realToFrac a
e
        | Bool
otherwise = a
1forall a. Fractional a => a -> a -> a
/a
0

-- | Floyd Warshall Without Paths (more efficient, by factor of 2).
--
-- TODO: use a mutable array?
-- TODO: implement Dijkstra's algorithm for adj matrix representation.
--
floydWarshall_
    :: Array U Ix2 Double
    -> Array U Ix2 Double
floydWarshall_ :: Array U Ix2 Double -> Array U Ix2 Double
floydWarshall_ Array U Ix2 Double
a = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Array U Ix2 Double -> Int -> Array U Ix2 Double
go Array U Ix2 Double
a [Int
0..Int
nforall a. Num a => a -> a -> a
-Int
1]
  where
    Sz (Int
n :. Int
_) = forall r ix e. Size r => Array r ix e -> Sz ix
size Array U Ix2 Double
a

    go :: Array U Ix2 Double -> Int -> Array U Ix2 Double
    go :: Array U Ix2 Double -> Int -> Array U Ix2 Double
go Array U Ix2 Double
c Int
k = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq (forall ix. Index ix => ix -> Sz ix
Sz (Int
n Int -> Int -> Ix2
:. Int
n)) forall a b. (a -> b) -> a -> b
$ \(Int
x :. Int
y) ->
        let
            !xy :: Double
xy = Array U Ix2 Double
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
y)
            !xk :: Double
xk = Array U Ix2 Double
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
k)
            !ky :: Double
ky = Array U Ix2 Double
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
k Int -> Int -> Ix2
:. Int
y)
        in if Double
xy forall a. Ord a => a -> a -> Bool
> Double
xk forall a. Num a => a -> a -> a
+ Double
ky then Double
xk forall a. Num a => a -> a -> a
+ Double
ky else Double
xy

-- | Shortest path matrix.
--
-- All entries of the result matrix are either whole numbers or @Infinity@.
--
shortestPaths_ :: Array U Ix2 Int -> Array U Ix2 Double
shortestPaths_ :: DenseAdjMatrix -> Array U Ix2 Double
shortestPaths_ = Array U Ix2 Double -> Array U Ix2 Double
floydWarshall_ forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r e r' ix.
(Manifest r e, Load r' ix e) =>
r -> Array r' ix e -> Array r ix e
computeAs U
U forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r. Source r Int => Array r Ix2 Int -> Array D Ix2 Double
distMatrix_

-- | Diameter of a graph.
--
diameter_ :: Array U Ix2 Int -> Maybe Natural
diameter_ :: DenseAdjMatrix -> Maybe Natural
diameter_ DenseAdjMatrix
g
    | forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
M.isEmpty DenseAdjMatrix
g = forall a. a -> Maybe a
Just Natural
0
    | Bool
otherwise = let x :: Natural
x = forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ forall r ix e.
(HasCallStack, Shape r ix, Source r e, Ord e) =>
Array r ix e -> e
maximum' forall a b. (a -> b) -> a -> b
$ DenseAdjMatrix -> Array U Ix2 Double
shortestPaths_ DenseAdjMatrix
g
        in if Natural
x forall a. Eq a => a -> a -> Bool
== forall a b. (RealFrac a, Integral b) => a -> b
round (Double
1forall a. Fractional a => a -> a -> a
/Double
0 :: Double) then forall a. Maybe a
Nothing else forall a. a -> Maybe a
Just Natural
x