{-# LANGUAGE LambdaCase #-}

module Data.Graph.Tree.DFS where

import Control.Monad
import Data.Function
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

import Data.Buffer
import Data.Graph.Sparse

shortestPath :: (U.Unbox w, Num w) => SparseGraph w -> Vertex -> U.Vector w
shortestPath :: forall w. (Unbox w, Num w) => SparseGraph w -> Vertex -> Vector w
shortestPath SparseGraph w
gr Vertex
root = (forall s. ST s (MVector s w)) -> Vector w
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
U.create ((forall s. ST s (MVector s w)) -> Vector w)
-> (forall s. ST s (MVector s w)) -> Vector w
forall a b. (a -> b) -> a -> b
$ do
  let n :: Vertex
n = SparseGraph w -> Vertex
forall w. SparseGraph w -> Vertex
numVerticesSG SparseGraph w
gr
  MVector s w
dist <- Vertex -> ST s (MVector (PrimState (ST s)) w)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> m (MVector (PrimState m) a)
UM.unsafeNew Vertex
n
  MVector (PrimState (ST s)) w -> Vertex -> w -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Vertex -> a -> m ()
UM.unsafeWrite MVector s w
MVector (PrimState (ST s)) w
dist Vertex
root w
0
  Buffer s Vertex
stack <- Vertex -> ST s (Buffer (PrimState (ST s)) Vertex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vertex -> m (Buffer (PrimState m) a)
newBufferAsStack Vertex
n
  MVector s Vertex
parent <- Vertex -> ST s (MVector (PrimState (ST s)) Vertex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> m (MVector (PrimState m) a)
UM.unsafeNew Vertex
n

  Vector (Vertex, Vertex, w)
-> ((Vertex, Vertex, w) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
U.forM_ (SparseGraph w
gr SparseGraph w -> Vertex -> Vector (Vertex, Vertex, w)
forall w.
Unbox w =>
SparseGraph w -> Vertex -> Vector (Vertex, Vertex, w)
`iadjW` Vertex
root) (((Vertex, Vertex, w) -> ST s ()) -> ST s ())
-> ((Vertex, Vertex, w) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Vertex
ei, Vertex
v, w
d) -> do
    Vertex -> Buffer (PrimState (ST s)) Vertex -> ST s ()
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
a -> Buffer (PrimState m) a -> m ()
pushBack Vertex
ei Buffer s Vertex
Buffer (PrimState (ST s)) Vertex
stack
    MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Vertex -> a -> m ()
UM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
parent Vertex
v Vertex
root
    MVector (PrimState (ST s)) w -> Vertex -> w -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Vertex -> a -> m ()
UM.unsafeWrite MVector s w
MVector (PrimState (ST s)) w
dist Vertex
v w
d

  (ST s () -> ST s ()) -> ST s ()
forall a. (a -> a) -> a
fix ((ST s () -> ST s ()) -> ST s ())
-> (ST s () -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ST s ()
loop ->
    Buffer (PrimState (ST s)) Vertex -> ST s (Maybe Vertex)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Buffer (PrimState m) a -> m (Maybe a)
popBack Buffer s Vertex
Buffer (PrimState (ST s)) Vertex
stack ST s (Maybe Vertex) -> (Maybe Vertex -> ST s ()) -> ST s ()
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just Vertex
ei -> do
        let v :: Vertex
v = SparseGraph w -> Vector Vertex
forall w. SparseGraph w -> Vector Vertex
adjacentSG SparseGraph w
gr Vector Vertex -> Vertex -> Vertex
forall a. Unbox a => Vector a -> Vertex -> a
`U.unsafeIndex` Vertex
ei
        Vertex
pv <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Vertex -> m a
UM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
parent Vertex
v
        w
dv <- MVector (PrimState (ST s)) w -> Vertex -> ST s w
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Vertex -> m a
UM.unsafeRead MVector s w
MVector (PrimState (ST s)) w
dist Vertex
v
        Vector (Vertex, Vertex, w)
-> ((Vertex, Vertex, w) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
U.forM_ (SparseGraph w
gr SparseGraph w -> Vertex -> Vector (Vertex, Vertex, w)
forall w.
Unbox w =>
SparseGraph w -> Vertex -> Vector (Vertex, Vertex, w)
`iadjW` Vertex
v) (((Vertex, Vertex, w) -> ST s ()) -> ST s ())
-> ((Vertex, Vertex, w) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Vertex
nei, Vertex
nv, w
d) -> do
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Vertex
pv Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
/= Vertex
nv) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            Vertex -> Buffer (PrimState (ST s)) Vertex -> ST s ()
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
a -> Buffer (PrimState m) a -> m ()
pushBack Vertex
nei Buffer s Vertex
Buffer (PrimState (ST s)) Vertex
stack
            MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Vertex -> a -> m ()
UM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
parent Vertex
nv Vertex
v
            MVector (PrimState (ST s)) w -> Vertex -> w -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Vertex -> a -> m ()
UM.unsafeWrite MVector s w
MVector (PrimState (ST s)) w
dist Vertex
nv (w -> ST s ()) -> w -> ST s ()
forall a b. (a -> b) -> a -> b
$ w
dv w -> w -> w
forall a. Num a => a -> a -> a
+ w
d
        ST s ()
loop
      Maybe Vertex
Nothing -> () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  MVector s w -> ST s (MVector s w)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return MVector s w
dist

diameter :: (U.Unbox w, Ord w, Num w) => SparseGraph w -> w
diameter :: forall w. (Unbox w, Ord w, Num w) => SparseGraph w -> w
diameter SparseGraph w
tree =
  Vector w -> w
forall a. (Unbox a, Ord a) => Vector a -> a
U.maximum
    (Vector w -> w) -> (Vector w -> Vector w) -> Vector w -> w
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SparseGraph w -> Vertex -> Vector w
forall w. (Unbox w, Num w) => SparseGraph w -> Vertex -> Vector w
shortestPath SparseGraph w
tree
    (Vertex -> Vector w)
-> (Vector w -> Vertex) -> Vector w -> Vector w
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector w -> Vertex
forall a. (Unbox a, Ord a) => Vector a -> Vertex
U.maxIndex
    (Vector w -> w) -> Vector w -> w
forall a b. (a -> b) -> a -> b
$ SparseGraph w -> Vertex -> Vector w
forall w. (Unbox w, Num w) => SparseGraph w -> Vertex -> Vector w
shortestPath SparseGraph w
tree Vertex
0

height :: (U.Unbox w, Ord w, Num w) => SparseGraph w -> U.Vector w
height :: forall w. (Unbox w, Ord w, Num w) => SparseGraph w -> Vector w
height SparseGraph w
tree = (w -> w -> w) -> Vector w -> Vector w -> Vector w
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
U.zipWith w -> w -> w
forall a. Ord a => a -> a -> a
max Vector w
fromS Vector w
fromT
  where
    s :: Vertex
s = Vector w -> Vertex
forall a. (Unbox a, Ord a) => Vector a -> Vertex
U.maxIndex (Vector w -> Vertex) -> Vector w -> Vertex
forall a b. (a -> b) -> a -> b
$ SparseGraph w -> Vertex -> Vector w
forall w. (Unbox w, Num w) => SparseGraph w -> Vertex -> Vector w
shortestPath SparseGraph w
tree Vertex
0
    fromS :: Vector w
fromS = SparseGraph w -> Vertex -> Vector w
forall w. (Unbox w, Num w) => SparseGraph w -> Vertex -> Vector w
shortestPath SparseGraph w
tree Vertex
s
    t :: Vertex
t = Vector w -> Vertex
forall a. (Unbox a, Ord a) => Vector a -> Vertex
U.maxIndex Vector w
fromS
    fromT :: Vector w
fromT = SparseGraph w -> Vertex -> Vector w
forall w. (Unbox w, Num w) => SparseGraph w -> Vertex -> Vector w
shortestPath SparseGraph w
tree Vertex
t