{-# 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 s gr = fmap (fst . (gr V.!)) . bfs' s . fmap snd $ 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' s gr = extract s $ V.create
         $ do st  <- UMV.replicate n False
              out <- MV.new n
              go0 st out (s :<| mempty)
              pure out
  where
    n = V.length gr
    go0        :: forall s. UMV.MVector s Bool -> MV.MVector s [Int]
               -> Seq Int -> ST s ()
    go0 st out = go
      where
        visit i = do b <- UMV.read st i
                     UMV.write st i True -- mark i as visited
                     pure $ if b then Nothing else Just i

        go :: Seq Int -> ST s ()
        go = \case
          Empty       -> pure ()
          (u:<|queue) -> do ns <- wither visit . F.toList $ gr V.! u
                            MV.write out u ns -- write that u's children are ns
                            go (queue <> Seq.fromList 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 s v = go s
  where
    go i = Node i (map go $ v V.! i)