{-# LANGUAGE ScopedTypeVariables #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Algorithms.Graph.BFS
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--------------------------------------------------------------------------------
module Algorithms.Graph.BFS
  ( bfs
  , bfs'
  ) where

import           Control.Monad.ST.Strict
import qualified Data.Foldable as F
import           Data.Sequence (Seq(..))
import qualified Data.Sequence as Seq
import           Data.Tree
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector.Unboxed.Mutable as UMV
import           Witherable

--------------------------------------------------------------------------------

-- | Runs a BFS from the first vertex in the graph. The graph is given
-- in adjacency list representation.
--
-- running time: \(O(V + E)\)
bfs      :: Foldable f => Int -> V.Vector (v, f Int) -> Tree v
bfs :: Int -> Vector (v, f Int) -> Tree v
bfs Int
s Vector (v, f Int)
gr = (Int -> v) -> Tree Int -> Tree v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((v, f Int) -> v
forall a b. (a, b) -> a
fst ((v, f Int) -> v) -> (Int -> (v, f Int)) -> Int -> v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector (v, f Int)
gr Vector (v, f Int) -> Int -> (v, f Int)
forall a. Vector a -> Int -> a
V.!)) (Tree Int -> Tree v)
-> (Vector (v, f Int) -> Tree Int) -> Vector (v, f Int) -> Tree v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Vector (f Int) -> Tree Int
forall (f :: * -> *).
Foldable f =>
Int -> Vector (f Int) -> Tree Int
bfs' Int
s (Vector (f Int) -> Tree Int)
-> (Vector (v, f Int) -> Vector (f Int))
-> Vector (v, f Int)
-> Tree Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((v, f Int) -> f Int) -> Vector (v, f Int) -> Vector (f Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (v, f Int) -> f Int
forall a b. (a, b) -> b
snd (Vector (v, f Int) -> Tree v) -> Vector (v, f Int) -> Tree v
forall a b. (a -> b) -> a -> b
$ Vector (v, f Int)
gr

-- | Runs a BFS from the first vertex in the graph. The graph is given
-- in adjacency list representation.
--
-- running time: \(O(V + E)\)
bfs'      :: Foldable f => Int -> V.Vector (f Int) -> Tree Int
bfs' :: Int -> Vector (f Int) -> Tree Int
bfs' Int
s Vector (f Int)
gr = Int -> Vector [Int] -> Tree Int
extract Int
s (Vector [Int] -> Tree Int) -> Vector [Int] -> Tree Int
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MVector s [Int])) -> Vector [Int]
forall a. (forall s. ST s (MVector s a)) -> Vector a
V.create
         ((forall s. ST s (MVector s [Int])) -> Vector [Int])
-> (forall s. ST s (MVector s [Int])) -> Vector [Int]
forall a b. (a -> b) -> a -> b
$ do MVector s Bool
st  <- Int -> Bool -> ST s (MVector (PrimState (ST s)) Bool)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UMV.replicate Int
n Bool
False
              MVector s [Int]
out <- Int -> ST s (MVector (PrimState (ST s)) [Int])
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
MV.new Int
n
              MVector s Bool -> MVector s [Int] -> Seq Int -> ST s ()
forall s. MVector s Bool -> MVector s [Int] -> Seq Int -> ST s ()
go0 MVector s Bool
st MVector s [Int]
out (Int
s Int -> Seq Int -> Seq Int
forall a. a -> Seq a -> Seq a
:<| Seq Int
forall a. Monoid a => a
mempty)
              MVector s [Int] -> ST s (MVector s [Int])
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s [Int]
out
  where
    n :: Int
n = Vector (f Int) -> Int
forall a. Vector a -> Int
V.length Vector (f Int)
gr
    go0        :: forall s. UMV.MVector s Bool -> MV.MVector s [Int]
               -> Seq Int -> ST s ()
    go0 :: MVector s Bool -> MVector s [Int] -> Seq Int -> ST s ()
go0 MVector s Bool
st MVector s [Int]
out = Seq Int -> ST s ()
go
      where
        visit :: Int -> ST s (Maybe Int)
visit Int
i = do Bool
b <- MVector (PrimState (ST s)) Bool -> Int -> ST s Bool
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UMV.read MVector s Bool
MVector (PrimState (ST s)) Bool
st Int
i
                     MVector (PrimState (ST s)) Bool -> Int -> Bool -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UMV.write MVector s Bool
MVector (PrimState (ST s)) Bool
st Int
i Bool
True -- mark i as visited
                     Maybe Int -> ST s (Maybe Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Int -> ST s (Maybe Int)) -> Maybe Int -> ST s (Maybe Int)
forall a b. (a -> b) -> a -> b
$ if Bool
b then Maybe Int
forall a. Maybe a
Nothing else Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i

        go :: Seq Int -> ST s ()
        go :: Seq Int -> ST s ()
go = \case
          Seq Int
Empty       -> () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          (Int
u:<|Seq Int
queue) -> do [Int]
ns <- (Int -> ST s (Maybe Int)) -> [Int] -> ST s [Int]
forall (t :: * -> *) (f :: * -> *) a b.
(Witherable t, Applicative f) =>
(a -> f (Maybe b)) -> t a -> f (t b)
wither Int -> ST s (Maybe Int)
visit ([Int] -> ST s [Int]) -> (f Int -> [Int]) -> f Int -> ST s [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f Int -> [Int]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList (f Int -> ST s [Int]) -> f Int -> ST s [Int]
forall a b. (a -> b) -> a -> b
$ Vector (f Int)
gr Vector (f Int) -> Int -> f Int
forall a. Vector a -> Int -> a
V.! Int
u
                            MVector (PrimState (ST s)) [Int] -> Int -> [Int] -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector s [Int]
MVector (PrimState (ST s)) [Int]
out Int
u [Int]
ns -- write that u's children are ns
                            Seq Int -> ST s ()
go (Seq Int
queue Seq Int -> Seq Int -> Seq Int
forall a. Semigroup a => a -> a -> a
<> [Int] -> Seq Int
forall a. [a] -> Seq a
Seq.fromList [Int]
ns)


-- | Give na root index and a vector s.t. v[i] lists the children of
-- node i, builds the acutal tree.
extract     :: Int -> V.Vector [Int] -> Tree Int
extract :: Int -> Vector [Int] -> Tree Int
extract Int
s Vector [Int]
v = Int -> Tree Int
go Int
s
  where
    go :: Int -> Tree Int
go Int
i = Int -> Forest Int -> Tree Int
forall a. a -> Forest a -> Tree a
Node Int
i ((Int -> Tree Int) -> [Int] -> Forest Int
forall a b. (a -> b) -> [a] -> [b]
map Int -> Tree Int
go ([Int] -> Forest Int) -> [Int] -> Forest Int
forall a b. (a -> b) -> a -> b
$ Vector [Int]
v Vector [Int] -> Int -> [Int]
forall a. Vector a -> Int -> a
V.! Int
i)