{-|
Module      : Data.Graph.AdjacencyList.BFS
Description : Breadth-first search on adjacency-list graphs
Copyright   : Thodoris Papakonstantinou, 2017-2026
License     : LGPL-3
Maintainer  : dev@tpapak.com
Stability   : experimental
Portability : POSIX

Breadth-first search (BFS) for directed graphs represented as adjacency lists.
Provides two entry points:

* 'bfs' — BFS on a concrete 'Graph' value
* 'adjBFS' — BFS on an implicit graph given as an @IntMap [Vertex]@ adjacency map

Both produce a 'BFS' record containing the level (distance) of every reachable
vertex, the BFS parent map, and a topological ordering of the visited vertices.

Used by the Tide algorithm ('Data.Graph.AdjacencyList.PushRelabel.Pure') in the
@globalRelabel@ step to compute vertex heights from distances to the source and
sink in the residual graph.
 -}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}


module Data.Graph.AdjacencyList.BFS
  ( -- * BFS result
    BFS (..)
    -- * Running BFS
  , bfs
  , adjBFS
    -- * Utilities
  , spanningTree
  ) where

import Data.List
import Data.Tuple
import Data.Maybe
import qualified Data.IntMap as IM
import qualified Data.IntSet as Set

import Data.Graph.AdjacencyList

-- | Result of a breadth-first search from a single source vertex.
data BFS = BFS { BFS -> IntSet
frontier :: Set.IntSet
                 -- ^ Current frontier (vertices at the deepest explored level).
                 -- Empty when the search is complete.
               , BFS -> IntMap Int
level :: IM.IntMap Int
                 -- ^ Map from vertex to its BFS level (shortest distance from source).
               , BFS -> IntMap Int
parent :: IM.IntMap Vertex
                 -- ^ Map from vertex to its BFS parent.
                 -- The source vertex has no entry.
               , BFS -> Int
maxLevel :: Int
                 -- ^ Maximum level reached during the search.
               , BFS -> [Int]
topSort :: [Vertex]
                 -- ^ Vertices in BFS visit order.
                 -- For DAGs this coincides with a topological sort.
               } deriving (BFS -> BFS -> Bool
(BFS -> BFS -> Bool) -> (BFS -> BFS -> Bool) -> Eq BFS
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BFS -> BFS -> Bool
== :: BFS -> BFS -> Bool
$c/= :: BFS -> BFS -> Bool
/= :: BFS -> BFS -> Bool
Eq, Int -> BFS -> ShowS
[BFS] -> ShowS
BFS -> String
(Int -> BFS -> ShowS)
-> (BFS -> String) -> ([BFS] -> ShowS) -> Show BFS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BFS -> ShowS
showsPrec :: Int -> BFS -> ShowS
$cshow :: BFS -> String
show :: BFS -> String
$cshowList :: [BFS] -> ShowS
showList :: [BFS] -> ShowS
Show)

-- | Initial BFS state with only the source vertex in the frontier.
initialBFS :: Vertex -> BFS
initialBFS :: Int -> BFS
initialBFS Int
s = BFS { frontier :: IntSet
frontier = Int -> IntSet
Set.singleton Int
s
                   , level :: IntMap Int
level = [(Int, Int)] -> IntMap Int
forall a. [(Int, a)] -> IntMap a
IM.fromList [(Int
s,Int
0)]
                   , parent :: IntMap Int
parent= IntMap Int
forall a. IntMap a
IM.empty
                   , maxLevel :: Int
maxLevel = Int
0
                   , topSort :: [Int]
topSort = []
                   }

-- | Run BFS on a 'Graph' from the given source vertex.
--
-- Explores all vertices reachable from @s@ via the graph's 'neighbors'
-- function. Returns a 'BFS' record with levels, parents, and visit order.
--
-- If @s@ is not in the graph's vertex set, returns the initial (empty) BFS.
bfs :: Graph -> Vertex -> BFS
bfs :: Graph -> Int -> BFS
bfs Graph
g Int
s = 
  let vset :: IntSet
vset = [Int] -> IntSet
Set.fromList (Graph -> [Int]
vertices Graph
g)
      sbfs :: BFS
sbfs = Int -> BFS
initialBFS Int
s
      breadthFirstSearch :: BFS -> BFS
breadthFirstSearch BFS
b =
        if IntSet -> Bool
Set.null (BFS -> IntSet
frontier BFS
b) Bool -> Bool -> Bool
|| Bool -> Bool
not (Int -> IntSet -> Bool
Set.member Int
s IntSet
vset)
           then BFS
b { topSort = reverse (topSort b) }
           else
             let oldLevel :: Int
oldLevel = BFS -> Int
maxLevel BFS
b
                 newLevel :: Int
newLevel = Int
oldLevel Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                 oldLevels :: IntMap Int
oldLevels = BFS -> IntMap Int
level BFS
b
                 oldFrontiers :: IntSet
oldFrontiers = BFS -> IntSet
frontier BFS
b
                 -- Collect (neighbor, parent) pairs; use IntMap to deduplicate
                 -- and keep only newly discovered vertices in one pass
                 newParMap :: IntMap Int
newParMap = (IntMap Int -> Int -> IntMap Int)
-> IntMap Int -> IntSet -> IntMap Int
forall a. (a -> Int -> a) -> a -> IntSet -> a
Set.foldl'
                   (\IntMap Int
acc Int
v ->
                     (IntMap Int -> Int -> IntMap Int)
-> IntMap Int -> [Int] -> IntMap Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap Int
acc' Int
n ->
                       if Int -> IntMap Int -> Bool
forall a. Int -> IntMap a -> Bool
IM.member Int
n IntMap Int
oldLevels Bool -> Bool -> Bool
|| Int -> IntMap Int -> Bool
forall a. Int -> IntMap a -> Bool
IM.member Int
n IntMap Int
acc'
                         then IntMap Int
acc'
                         else Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
n Int
v IntMap Int
acc'
                     ) IntMap Int
acc (Graph -> Neighbors
neighbors Graph
g Int
v)
                   ) IntMap Int
forall a. IntMap a
IM.empty IntSet
oldFrontiers
                 newFrontiers :: IntSet
newFrontiers = IntMap Int -> IntSet
forall a. IntMap a -> IntSet
IM.keysSet IntMap Int
newParMap
                 newParents :: IntMap Int
newParents = IntMap Int -> IntMap Int -> IntMap Int
forall a. IntMap a -> IntMap a -> IntMap a
IM.union (BFS -> IntMap Int
parent BFS
b) IntMap Int
newParMap
                 newLevels :: IntMap Int
newLevels = (IntMap Int -> Int -> IntMap Int)
-> IntMap Int -> IntSet -> IntMap Int
forall a. (a -> Int -> a) -> a -> IntSet -> a
Set.foldl' 
                           (\IntMap Int
ac Int
v -> Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
v Int
newLevel IntMap Int
ac) 
                           IntMap Int
oldLevels IntSet
newFrontiers
                 -- Prepend frontier to topSort (reversed at the end)
                 newTopSort :: [Int]
newTopSort = ([Int] -> Neighbors) -> [Int] -> IntSet -> [Int]
forall a. (a -> Int -> a) -> a -> IntSet -> a
Set.foldl' ((Int -> [Int] -> [Int]) -> [Int] -> Neighbors
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) (BFS -> [Int]
topSort BFS
b) IntSet
oldFrontiers
                 bbfs :: BFS
bbfs = BFS -> BFS
breadthFirstSearch (BFS
b { frontier = newFrontiers
                                              , level = newLevels 
                                              , parent = newParents
                                              , maxLevel = newLevel
                                              , topSort = newTopSort
                                            })
               in BFS
bbfs
   in BFS -> BFS
breadthFirstSearch BFS
sbfs

-- | Run BFS on an implicit graph defined by an adjacency map.
--
-- @adjBFS neimap s@ performs BFS from vertex @s@ where the neighbors of
-- each vertex are given by @neimap :: IntMap [Vertex]@.  Vertices not
-- present in @neimap@ are treated as having no outgoing edges.
--
-- This is used by 'Data.Graph.AdjacencyList.PushRelabel.Internal.residualDistances'
-- to run BFS on the residual graph (whose edge set changes each tide)
-- without constructing a full 'Graph' value.
adjBFS :: IM.IntMap [Vertex] -> Vertex -> BFS
adjBFS :: IntMap [Int] -> Int -> BFS
adjBFS IntMap [Int]
neimap Int
s = let b :: BFS
b = BFS -> BFS
breadthFirstSearch BFS
sbfs
                  in BFS
b { topSort = reverse (topSort b) }
  where neighbors :: Neighbors
neighbors Int
v = case Int -> IntMap [Int] -> Maybe [Int]
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
v IntMap [Int]
neimap of
                        Maybe [Int]
Nothing -> []
                        Just [Int]
ns -> [Int]
ns
        sbfs :: BFS
sbfs = Int -> BFS
initialBFS Int
s
        breadthFirstSearch :: BFS -> BFS
breadthFirstSearch BFS
b
          | IntSet -> Bool
Set.null (BFS -> IntSet
frontier BFS
b) = BFS
b
          | Bool
otherwise = BFS
bbfs
            where oldLevel :: Int
oldLevel = BFS -> Int
maxLevel BFS
b
                  newLevel :: Int
newLevel = Int
oldLevel Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                  oldLevels :: IntMap Int
oldLevels = BFS -> IntMap Int
level BFS
b
                  oldFrontiers :: IntSet
oldFrontiers = BFS -> IntSet
frontier BFS
b
                  -- Collect new vertices; use IntMap to deduplicate
                  newParMap :: IntMap Int
newParMap = (IntMap Int -> Int -> IntMap Int)
-> IntMap Int -> IntSet -> IntMap Int
forall a. (a -> Int -> a) -> a -> IntSet -> a
Set.foldl'
                    (\IntMap Int
acc Int
v ->
                      (IntMap Int -> Int -> IntMap Int)
-> IntMap Int -> [Int] -> IntMap Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap Int
acc' Int
n ->
                        if Int -> IntMap Int -> Bool
forall a. Int -> IntMap a -> Bool
IM.member Int
n IntMap Int
oldLevels Bool -> Bool -> Bool
|| Int -> IntMap Int -> Bool
forall a. Int -> IntMap a -> Bool
IM.member Int
n IntMap Int
acc'
                          then IntMap Int
acc'
                          else Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
n Int
v IntMap Int
acc'
                      ) IntMap Int
acc (Neighbors
neighbors Int
v)
                    ) IntMap Int
forall a. IntMap a
IM.empty IntSet
oldFrontiers
                  newFrontiers :: IntSet
newFrontiers = IntMap Int -> IntSet
forall a. IntMap a -> IntSet
IM.keysSet IntMap Int
newParMap
                  newParents :: IntMap Int
newParents = IntMap Int -> IntMap Int -> IntMap Int
forall a. IntMap a -> IntMap a -> IntMap a
IM.union (BFS -> IntMap Int
parent BFS
b) IntMap Int
newParMap
                  newLevels :: IntMap Int
newLevels = (IntMap Int -> Int -> IntMap Int)
-> IntMap Int -> IntSet -> IntMap Int
forall a. (a -> Int -> a) -> a -> IntSet -> a
Set.foldl' 
                                 (\IntMap Int
ac Int
v -> Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
v Int
newLevel IntMap Int
ac) 
                                 IntMap Int
oldLevels IntSet
newFrontiers
                  newTopSort :: [Int]
newTopSort = ([Int] -> Neighbors) -> [Int] -> IntSet -> [Int]
forall a. (a -> Int -> a) -> a -> IntSet -> a
Set.foldl' ((Int -> [Int] -> [Int]) -> [Int] -> Neighbors
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) (BFS -> [Int]
topSort BFS
b) IntSet
oldFrontiers
                  bbfs :: BFS
bbfs = BFS -> BFS
breadthFirstSearch (BFS
b { frontier = newFrontiers
                             , level = newLevels 
                             , parent = newParents
                             , maxLevel = newLevel
                             , topSort = newTopSort
                             })

-- | Extract the BFS spanning tree as a list of edges.
--
-- Each edge @(parent, child)@ in the returned list corresponds to one
-- entry in the 'parent' map.
spanningTree :: BFS -> [Edge]
spanningTree :: BFS -> [Edge]
spanningTree BFS
b = 
  ((Int, Int) -> Edge) -> [(Int, Int)] -> [Edge]
forall a b. (a -> b) -> [a] -> [b]
map ((Int, Int) -> Edge
fromTuple ((Int, Int) -> Edge)
-> ((Int, Int) -> (Int, Int)) -> (Int, Int) -> Edge
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Int) -> (Int, Int)
forall a b. (a, b) -> (b, a)
swap) ([(Int, Int)] -> [Edge]) -> [(Int, Int)] -> [Edge]
forall a b. (a -> b) -> a -> b
$ IntMap Int -> [(Int, Int)]
forall a. IntMap a -> [(Int, a)]
IM.toList (IntMap Int -> [(Int, Int)]) -> IntMap Int -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ BFS -> IntMap Int
parent BFS
b