{-|
Module      : Data.Graph.AdjacencyList.PushRelabel.Internal
Description : Residual graph types and primitive operations for the Tide algorithm
Copyright   : Thodoris Papakonstantinou, 2017-2026
License     : LGPL-3
Maintainer  : dev@tpapak.com
Stability   : experimental
Portability : POSIX

Internal definitions for the Tide push-pull-relabel max-flow algorithm.

This module defines:

* 'ResidualGraph' — the mutable state threaded through each tide iteration,
  containing vertex heights, excesses, edge flows, and the set of overflowing
  vertices grouped by level.
* 'ResidualVertex' and 'ResidualEdge' — per-vertex and per-edge state.
* 'NeighborsMap' — an @IntMap@-based adjacency structure that maps each vertex
  to its forward and reverse neighbors with O(log V) edge-index lookup
  (replacing the original O(log E) @Map Edge Int@ lookup).
* Primitive operations: 'push', 'pull', 'updateHeight', 'updateExcess',
  'updateEdge', 'residualDistances'.

The 'topologyChanged' flag tracks whether any edge crossed a saturation
boundary (became saturated or unsaturated) during push\/pull.  When the
flag is 'False', the next tide can skip @globalRelabel@ — an optimization
that yields 1.25--1.61x speedup in practice.
 -}

{-# LANGUAGE BangPatterns #-}

module Data.Graph.AdjacencyList.PushRelabel.Internal
  ( -- * Re-exports from Network
    Network (..)
  , Capacity (..)
  , Capacities (..)
  , Flow 
    -- * Residual graph types
  , ResidualGraph (..)
  , ResidualVertex (..)
  , ResidualVertices
  , ResidualEdge (..)
  , ResidualEdges
  , NeighborsMap
  , Overflowing (..)
    -- * Vertex property types
  , Height
  , Excess
  , Level
    -- * Initialization
  , initializeResidualGraph
    -- * Vertex property accessors
  , level
  , excess
  , height
    -- * Edge property accessors
  , edgeCapacity
  , edgeFlow
  , resEdgeIndex
    -- * Flow queries
  , netFlow
  , inflow
  , outflow
  , sourceEdgesCapacity
    -- * Push and pull operations
  , push
  , pull
    -- * State updates
  , updateHeight
  , updateExcess
  , updateEdge
    -- * Overflowing vertex tracking
  , getOverflowing
    -- * Network reconstruction
  , networkFromResidual
    -- * Residual BFS (for @globalRelabel@)
  , residualDistances
    -- * Min-cut
  , stCut
  ) where

import Data.List
import Data.Maybe
import qualified Data.Map.Lazy as M
import qualified Data.IntMap.Lazy as IM
import qualified Data.IntSet as Set

import Data.Graph.AdjacencyList
import Data.Graph.AdjacencyList.Network
import qualified Data.Graph.AdjacencyList.BFS as BFS

-- | Vertex height in the push-relabel framework.
-- For source-side vertices: @height = |V| + distance_from_source@.
-- For sink-side vertices: @height = distance_from_sink@.
type Height = Int

-- | Vertex excess: @inflow - outflow@.  Positive excess means the vertex
-- is overflowing and needs to push or pull flow.
type Excess = Capacity

-- | Level: the shortest-path distance from the source in the /original/
-- (not residual) graph.  Constant throughout the algorithm.
-- Determines the ordering of vertices in globalPush (left fold, ascending)
-- and globalPull (right fold, descending).
type Level = Int

-- | Per-vertex state in the residual graph.
--
-- @ResidualVertex v l h x@ stores:
--
-- * @v@ — vertex identifier
-- * @l@ — level (BFS distance from source in original graph, constant)
-- * @h@ — height (updated by @globalRelabel@ each tide)
-- * @x@ — excess flow (updated by push\/pull operations)
data ResidualVertex = ResidualVertex !Vertex !Level !Height !Excess
  deriving (ResidualVertex -> ResidualVertex -> Bool
(ResidualVertex -> ResidualVertex -> Bool)
-> (ResidualVertex -> ResidualVertex -> Bool) -> Eq ResidualVertex
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ResidualVertex -> ResidualVertex -> Bool
== :: ResidualVertex -> ResidualVertex -> Bool
$c/= :: ResidualVertex -> ResidualVertex -> Bool
/= :: ResidualVertex -> ResidualVertex -> Bool
Eq)
instance Show ResidualVertex where
  show :: ResidualVertex -> String
show (ResidualVertex Vertex
v Vertex
l Vertex
h Excess
x) =
    String
"RVertex " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Vertex -> String
forall a. Show a => a -> String
show Vertex
v String -> ShowS
forall a. [a] -> [a] -> [a]
++  String
" level: " String -> ShowS
forall a. [a] -> [a] -> [a]
++
      Vertex -> String
forall a. Show a => a -> String
show Vertex
l String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" height: " String -> ShowS
forall a. [a] -> [a] -> [a]
++
      Vertex -> String
forall a. Show a => a -> String
show Vertex
h String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" excess: " String -> ShowS
forall a. [a] -> [a] -> [a]
++
      Double -> String
forall a. Show a => a -> String
show (Excess -> Double
forall a. Fractional a => Excess -> a
fromRational Excess
x :: Double)

-- | Map from vertex id to its 'ResidualVertex' state.
type ResidualVertices = IM.IntMap ResidualVertex

-- | Per-edge state: original edge, capacity, and current flow (preflow).
--
-- @ResidualEdge e c f@: edge @e@ with capacity @c@ and flow @f@.
-- A forward residual edge exists when @f < c@; a backward residual edge
-- exists when @f > 0@.
data ResidualEdge = ResidualEdge Edge Capacity Flow
  deriving (ResidualEdge -> ResidualEdge -> Bool
(ResidualEdge -> ResidualEdge -> Bool)
-> (ResidualEdge -> ResidualEdge -> Bool) -> Eq ResidualEdge
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ResidualEdge -> ResidualEdge -> Bool
== :: ResidualEdge -> ResidualEdge -> Bool
$c/= :: ResidualEdge -> ResidualEdge -> Bool
/= :: ResidualEdge -> ResidualEdge -> Bool
Eq)
instance Show ResidualEdge where
  show :: ResidualEdge -> String
show (ResidualEdge Edge
e Excess
c Excess
f) =
    String
"REdge " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Edge -> String
forall a. Show a => a -> String
show Edge
e 
      String -> ShowS
forall a. [a] -> [a] -> [a]
++  String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++
      Double -> String
forall a. Show a => a -> String
show (Excess -> Double
forall a. Fractional a => Excess -> a
fromRational Excess
c :: Double)
      String -> ShowS
forall a. [a] -> [a] -> [a]
++  String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++
      Double -> String
forall a. Show a => a -> String
show (Excess -> Double
forall a. Fractional a => Excess -> a
fromRational Excess
f :: Double)
-- | Map from edge index to its 'ResidualEdge' state.
type ResidualEdges = IM.IntMap ResidualEdge

-- | For each vertex, maps forward neighbors and reverse neighbors
-- to their edge indices in the graph's 'EdgeMap'.
--
-- @NeighborsMap ! v = (fwdMap, revMap)@ where:
--
-- * @fwdMap ! w@ = index of edge @(v, w)@ (forward neighbor)
-- * @revMap ! u@ = index of edge @(u, v)@ (reverse neighbor)
--
-- This provides O(log degree) edge-index lookup, replacing the original
-- O(log E) lookup via @Map Edge Int@.
type NeighborsMap = IM.IntMap (IM.IntMap Int, IM.IntMap Int)

-- | Overflowing vertices grouped by level.
-- Keys are levels (BFS distance from source); values are sets of
-- vertices at that level with positive excess.
--
-- This structure determines the iteration order for globalPush
-- (ascending level = left fold) and globalPull (descending level = right fold).
type Overflowing = IM.IntMap Set.IntSet

-- | The residual graph: the complete mutable state of the Tide algorithm.
--
-- Threaded through each tide iteration.  Contains the underlying network,
-- per-vertex and per-edge state, the neighbor map for O(log V) edge lookup,
-- overflowing vertex sets, step counter, and the topology-change flag.
data ResidualGraph = 
  ResidualGraph { ResidualGraph -> Network
network :: !Network
                  -- ^ The original flow network.
                , ResidualGraph -> ResidualVertices
netVertices :: !ResidualVertices
                  -- ^ Per-vertex state (level, height, excess).
                , ResidualGraph -> ResidualEdges
netEdges :: !ResidualEdges 
                  -- ^ Per-edge state (capacity, flow).
                , ResidualGraph -> NeighborsMap
netNeighborsMap :: !NeighborsMap 
                  -- ^ Adjacency map for O(log V) edge-index lookup.
                , ResidualGraph -> Overflowing
overflowing :: !Overflowing
                  -- ^ Overflowing vertices grouped by level.
                , ResidualGraph -> Vertex
steps :: !Int
                  -- ^ Number of completed tide iterations.
                , ResidualGraph -> Bool
topologyChanged :: !Bool
                  -- ^ Whether any edge crossed a saturation boundary
                  -- (became saturated or unsaturated) during the
                  -- most recent push\/pull phase.  When 'False',
                  -- the next tide can skip @globalRelabel@.
                }
   deriving (Vertex -> ResidualGraph -> ShowS
[ResidualGraph] -> ShowS
ResidualGraph -> String
(Vertex -> ResidualGraph -> ShowS)
-> (ResidualGraph -> String)
-> ([ResidualGraph] -> ShowS)
-> Show ResidualGraph
forall a.
(Vertex -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Vertex -> ResidualGraph -> ShowS
showsPrec :: Vertex -> ResidualGraph -> ShowS
$cshow :: ResidualGraph -> String
show :: ResidualGraph -> String
$cshowList :: [ResidualGraph] -> ShowS
showList :: [ResidualGraph] -> ShowS
Show,ResidualGraph -> ResidualGraph -> Bool
(ResidualGraph -> ResidualGraph -> Bool)
-> (ResidualGraph -> ResidualGraph -> Bool) -> Eq ResidualGraph
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ResidualGraph -> ResidualGraph -> Bool
== :: ResidualGraph -> ResidualGraph -> Bool
$c/= :: ResidualGraph -> ResidualGraph -> Bool
/= :: ResidualGraph -> ResidualGraph -> Bool
Eq)

-- | Build the initial 'ResidualGraph' from a 'Network'.
--
-- Saturates all edges leaving the source (setting their flow equal to
-- capacity), sets the source height to @|V|@, and initializes the
-- overflowing set with all vertices that received flow from the source.
--
-- The 'topologyChanged' flag is set to 'True' so the first tide always
-- runs @globalRelabel@.
initializeResidualGraph :: Network -> ResidualGraph
initializeResidualGraph :: Network -> ResidualGraph
initializeResidualGraph Network
net = 
  let vs :: ResidualVertices
vs = Network -> ResidualVertices
initializeVertices Network
net
      es :: ResidualEdges
es = Network -> ResidualEdges
initializeEdges Network
net
      neimap :: NeighborsMap
neimap = Graph -> NeighborsMap
getNetNeighborsMap (Graph -> NeighborsMap) -> Graph -> NeighborsMap
forall a b. (a -> b) -> a -> b
$ Network -> Graph
graph Network
net 
   in ResidualGraph { network :: Network
network = Network
net
                    , netVertices :: ResidualVertices
netVertices = ResidualVertices
vs 
                    , netEdges :: ResidualEdges
netEdges = ResidualEdges
es 
                    , netNeighborsMap :: NeighborsMap
netNeighborsMap = NeighborsMap
neimap
                    , overflowing :: Overflowing
overflowing = 
                      let ovfs :: IntSet
ovfs = ResidualVertices -> IntSet
getOverflowing ResidualVertices
vs
                          bfs :: BFS
bfs = Graph -> Vertex -> BFS
BFS.bfs (Network -> Graph
graph Network
net) (Network -> Vertex
source Network
net)
                          maxLevel :: Vertex
maxLevel = BFS -> Vertex
BFS.maxLevel BFS
bfs
                          fl :: Vertex -> Vertex
fl Vertex
v = 
                            let (ResidualVertex Vertex
_ Vertex
l Vertex
_ Excess
_) = 
                                  Maybe ResidualVertex -> ResidualVertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe ResidualVertex -> ResidualVertex)
-> Maybe ResidualVertex -> ResidualVertex
forall a b. (a -> b) -> a -> b
$ Vertex -> ResidualVertices -> Maybe ResidualVertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v ResidualVertices
vs
                             in Vertex
l
                       in (Overflowing -> Vertex -> Overflowing)
-> Overflowing -> IntSet -> Overflowing
forall a. (a -> Vertex -> a) -> a -> IntSet -> a
Set.foldl' 
                            (\Overflowing
ac Vertex
v -> 
                               (IntSet -> IntSet) -> Vertex -> Overflowing -> Overflowing
forall a. (a -> a) -> Vertex -> IntMap a -> IntMap a
IM.adjust (\IntSet
ps -> Vertex -> IntSet -> IntSet
Set.insert Vertex
v IntSet
ps) (Vertex -> Vertex
fl Vertex
v) Overflowing
ac
                            ) ([(Vertex, IntSet)] -> Overflowing
forall a. [(Vertex, a)] -> IntMap a
IM.fromList ([Vertex] -> [IntSet] -> [(Vertex, IntSet)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Vertex
1..Vertex
maxLevel] (IntSet -> [IntSet]
forall a. a -> [a]
repeat IntSet
Set.empty))) IntSet
ovfs
                    , steps :: Vertex
steps = Vertex
0
                     , topologyChanged :: Bool
topologyChanged = Bool
True
                     } 

-- | Build the 'NeighborsMap' from a 'Graph'.
--
-- For each vertex @v@, computes:
--
-- * Forward map: @neighbor -> edgeIndex@ for edges @(v, neighbor)@
-- * Reverse map: @neighbor -> edgeIndex@ for edges @(neighbor, v)@
getNetNeighborsMap :: Graph -> NeighborsMap
getNetNeighborsMap :: Graph -> NeighborsMap
getNetNeighborsMap Graph
g =
  let revgraph :: Graph
revgraph = Graph -> Graph
reverseGraph Graph
g
      neis :: Vertex -> (IntMap Vertex, IntMap Vertex)
neis Vertex
v = 
        let fwd :: IntMap Vertex
fwd = [(Vertex, Vertex)] -> IntMap Vertex
forall a. [(Vertex, a)] -> IntMap a
IM.fromList 
                    [ (Vertex
n, Maybe Vertex -> Vertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ Graph -> Edge -> Maybe Vertex
edgeIndex Graph
g (Vertex -> Vertex -> Edge
Edge Vertex
v Vertex
n)) 
                    | Vertex
n <- Graph -> Neighbors
neighbors Graph
g Vertex
v ]
            rev :: IntMap Vertex
rev = [(Vertex, Vertex)] -> IntMap Vertex
forall a. [(Vertex, a)] -> IntMap a
IM.fromList 
                    [ (Vertex
n, Maybe Vertex -> Vertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ Graph -> Edge -> Maybe Vertex
edgeIndex Graph
g (Vertex -> Vertex -> Edge
Edge Vertex
n Vertex
v)) 
                    | Vertex
n <- Graph -> Neighbors
neighbors Graph
revgraph Vertex
v ]
         in (IntMap Vertex
fwd, IntMap Vertex
rev)
   in (NeighborsMap -> Vertex -> NeighborsMap)
-> NeighborsMap -> [Vertex] -> NeighborsMap
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' 
        (\NeighborsMap
ac Vertex
v -> Vertex
-> (IntMap Vertex, IntMap Vertex) -> NeighborsMap -> NeighborsMap
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
v (Vertex -> (IntMap Vertex, IntMap Vertex)
neis Vertex
v) NeighborsMap
ac) 
        NeighborsMap
forall a. IntMap a
IM.empty (Graph -> [Vertex]
vertices Graph
g)

-- | Look up forward and reverse neighbor maps for a vertex.
netNeighbors :: NeighborsMap 
             -> Vertex 
             -> (IM.IntMap Int, IM.IntMap Int) 
netNeighbors :: NeighborsMap -> Vertex -> (IntMap Vertex, IntMap Vertex)
netNeighbors NeighborsMap
nm Vertex
v = 
  Maybe (IntMap Vertex, IntMap Vertex)
-> (IntMap Vertex, IntMap Vertex)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (IntMap Vertex, IntMap Vertex)
 -> (IntMap Vertex, IntMap Vertex))
-> Maybe (IntMap Vertex, IntMap Vertex)
-> (IntMap Vertex, IntMap Vertex)
forall a b. (a -> b) -> a -> b
$ Vertex -> NeighborsMap -> Maybe (IntMap Vertex, IntMap Vertex)
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v NeighborsMap
nm

-- | O(log degree) edge index lookup via 'NeighborsMap'.
--
-- Looks up the edge index of @(u, v)@ by finding @v@ in the forward
-- neighbor map of @u@.  Returns 'Nothing' if the edge does not exist.
resEdgeIndex :: NeighborsMap -> Edge -> Maybe Int
resEdgeIndex :: NeighborsMap -> Edge -> Maybe Vertex
resEdgeIndex NeighborsMap
nm (Edge Vertex
u Vertex
v) = do
  (IntMap Vertex
fwd, IntMap Vertex
_) <- Vertex -> NeighborsMap -> Maybe (IntMap Vertex, IntMap Vertex)
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
u NeighborsMap
nm
  Vertex -> IntMap Vertex -> Maybe Vertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v IntMap Vertex
fwd

sourceEdges :: Network -> [(Edge,Capacity)]
sourceEdges :: Network -> [(Edge, Excess)]
sourceEdges Network
net = 
  let g :: Graph
g = Network -> Graph
graph Network
net
      cs :: Capacities
cs = Network -> Capacities
capacities Network
net
      s :: Vertex
s = Network -> Vertex
source Network
net
      cap :: Vertex -> Excess
cap Vertex
v = Maybe Excess -> Excess
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Excess -> Excess) -> Maybe Excess -> Excess
forall a b. (a -> b) -> a -> b
$ Edge -> Capacities -> Maybe Excess
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Vertex -> Vertex -> Edge
Edge Vertex
s Vertex
v) Capacities
cs
    in (Vertex -> (Edge, Excess)) -> [Vertex] -> [(Edge, Excess)]
forall a b. (a -> b) -> [a] -> [b]
map (\Vertex
v -> ((Vertex -> Vertex -> Edge
Edge Vertex
s Vertex
v), Vertex -> Excess
cap Vertex
v )) (Graph -> Neighbors
neighbors Graph
g Vertex
s) 

-- | Total capacity of all edges leaving the source.
-- This is an upper bound on the maximum flow.
sourceEdgesCapacity :: Network -> Capacity
sourceEdgesCapacity :: Network -> Excess
sourceEdgesCapacity Network
net = 
  let ses :: [(Edge, Excess)]
ses = Network -> [(Edge, Excess)]
sourceEdges Network
net
   in [Excess] -> Excess
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Excess] -> Excess) -> [Excess] -> Excess
forall a b. (a -> b) -> a -> b
$ ((Edge, Excess) -> Excess) -> [(Edge, Excess)] -> [Excess]
forall a b. (a -> b) -> [a] -> [b]
map (Edge, Excess) -> Excess
forall a b. (a, b) -> b
snd [(Edge, Excess)]
ses

-- | Initialize vertex state: set source height to @|V|@, saturate source
-- edges (giving excess to source neighbors), set all other heights to 0.
initializeVertices :: Network -> ResidualVertices
initializeVertices :: Network -> ResidualVertices
initializeVertices Network
net =
  let g :: Graph
g = Network -> Graph
graph Network
net
      cs :: Capacities
cs = Network -> Capacities
capacities Network
net
      s :: Vertex
s = Network -> Vertex
source Network
net
      t :: Vertex
t = Network -> Vertex
sink Network
net
      sh :: Vertex
sh = Vertex -> Vertex
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vertex -> Vertex) -> Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ Graph -> Vertex
numVertices Graph
g
      ses :: [(Edge, Excess)]
ses = Network -> [(Edge, Excess)]
sourceEdges Network
net
      vs :: [Vertex]
vs = Graph -> [Vertex]
vertices (Graph -> [Vertex]) -> Graph -> [Vertex]
forall a b. (a -> b) -> a -> b
$ Network -> Graph
graph Network
net
      flevels :: IntMap Vertex
flevels = BFS -> IntMap Vertex
BFS.level (BFS -> IntMap Vertex) -> BFS -> IntMap Vertex
forall a b. (a -> b) -> a -> b
$ Graph -> Vertex -> BFS
BFS.bfs (Network -> Graph
graph Network
net) (Network -> Vertex
source Network
net)
      fl :: Vertex -> Vertex
fl Vertex
v = Maybe Vertex -> Vertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ Vertex -> IntMap Vertex -> Maybe Vertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v IntMap Vertex
flevels
      zvs :: ResidualVertices
zvs = [(Vertex, ResidualVertex)] -> ResidualVertices
forall a. [(Vertex, a)] -> IntMap a
IM.fromList ([(Vertex, ResidualVertex)] -> ResidualVertices)
-> [(Vertex, ResidualVertex)] -> ResidualVertices
forall a b. (a -> b) -> a -> b
$ 
        [Vertex] -> [ResidualVertex] -> [(Vertex, ResidualVertex)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Graph -> [Vertex]
vertices Graph
g) ((Vertex -> ResidualVertex) -> [Vertex] -> [ResidualVertex]
forall a b. (a -> b) -> [a] -> [b]
map (\Vertex
v -> 
          Vertex -> Vertex -> Vertex -> Excess -> ResidualVertex
ResidualVertex Vertex
v (Vertex -> Vertex
fl Vertex
v) Vertex
0 Excess
0) ([Vertex] -> [ResidualVertex]) -> [Vertex] -> [ResidualVertex]
forall a b. (a -> b) -> a -> b
$ Graph -> [Vertex]
vertices Graph
g)
      (Excess
sx, ResidualVertices
nvs) = ((Excess, ResidualVertices)
 -> (Edge, Excess) -> (Excess, ResidualVertices))
-> (Excess, ResidualVertices)
-> [(Edge, Excess)]
-> (Excess, ResidualVertices)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\(Excess
cx,ResidualVertices
ac) (Edge
e,Excess
c) -> 
        let v :: Vertex
v = Edge -> Vertex
to Edge
e
         in (Excess
cxExcess -> Excess -> Excess
forall a. Num a => a -> a -> a
-Excess
c, (ResidualVertex -> ResidualVertex)
-> Vertex -> ResidualVertices -> ResidualVertices
forall a. (a -> a) -> Vertex -> IntMap a -> IntMap a
IM.adjust (ResidualVertex -> ResidualVertex -> ResidualVertex
forall a b. a -> b -> a
const (Vertex -> Vertex -> Vertex -> Excess -> ResidualVertex
ResidualVertex Vertex
v (Vertex -> Vertex
fl Vertex
v) Vertex
0 Excess
c)) Vertex
v ResidualVertices
ac)) (Excess
0, ResidualVertices
zvs) [(Edge, Excess)]
ses
   in Vertex -> ResidualVertex -> ResidualVertices -> ResidualVertices
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
s (Vertex -> Vertex -> Vertex -> Excess -> ResidualVertex
ResidualVertex Vertex
s Vertex
0 Vertex
sh Excess
sx) ResidualVertices
nvs

-- | Initialize edge state: saturate source edges, set all others to zero flow.
initializeEdges :: Network -> ResidualEdges
initializeEdges :: Network -> ResidualEdges
initializeEdges Network
net =
  let g :: Graph
g = Network -> Graph
graph Network
net
      cs :: Capacities
cs = Network -> Capacities
capacities Network
net
      s :: Vertex
s = Network -> Vertex
source Network
net
      t :: Vertex
t = Network -> Vertex
sink Network
net
      inites :: ResidualEdges
inites = [(Vertex, ResidualEdge)] -> ResidualEdges
forall a. [(Vertex, a)] -> IntMap a
IM.fromList ([(Vertex, ResidualEdge)] -> ResidualEdges)
-> [(Vertex, ResidualEdge)] -> ResidualEdges
forall a b. (a -> b) -> a -> b
$ ((Edge, Excess) -> (Vertex, ResidualEdge))
-> [(Edge, Excess)] -> [(Vertex, ResidualEdge)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Edge
e,Excess
c) -> (Maybe Vertex -> Vertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ Graph -> Edge -> Maybe Vertex
edgeIndex Graph
g Edge
e, Edge -> Excess -> Excess -> ResidualEdge
ResidualEdge Edge
e Excess
c Excess
0)) (Capacities -> [(Edge, Excess)]
forall k a. Map k a -> [(k, a)]
M.toList Capacities
cs)
      ses :: [(Edge, Excess)]
ses = Network -> [(Edge, Excess)]
sourceEdges Network
net
   in  (ResidualEdges -> (Edge, Excess) -> ResidualEdges)
-> ResidualEdges -> [(Edge, Excess)] -> ResidualEdges
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\ResidualEdges
ac (Edge
e,Excess
c) -> Vertex -> ResidualEdge -> ResidualEdges -> ResidualEdges
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert (Maybe Vertex -> Vertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ Graph -> Edge -> Maybe Vertex
edgeIndex Graph
g Edge
e) (Edge -> Excess -> Excess -> ResidualEdge
ResidualEdge Edge
e Excess
c Excess
c) ResidualEdges
ac) ResidualEdges
inites [(Edge, Excess)]
ses 

-- | Collect all vertices with positive excess.
getOverflowing :: IM.IntMap ResidualVertex -> Set.IntSet
getOverflowing :: ResidualVertices -> IntSet
getOverflowing ResidualVertices
nvs = 
  let xv :: ResidualVertex -> Excess
xv (ResidualVertex Vertex
v Vertex
_ Vertex
_ Excess
x) = Excess
x
      vv :: ResidualVertex -> Vertex
vv (ResidualVertex Vertex
v Vertex
_ Vertex
_ Excess
x) = Vertex
v
   in [Vertex] -> IntSet
Set.fromList ([Vertex] -> IntSet) -> [Vertex] -> IntSet
forall a b. (a -> b) -> a -> b
$ ((Vertex, Vertex) -> Vertex) -> [(Vertex, Vertex)] -> [Vertex]
forall a b. (a -> b) -> [a] -> [b]
map (Vertex, Vertex) -> Vertex
forall a b. (a, b) -> b
snd ([(Vertex, Vertex)] -> [Vertex]) -> [(Vertex, Vertex)] -> [Vertex]
forall a b. (a -> b) -> a -> b
$ IntMap Vertex -> [(Vertex, Vertex)]
forall a. IntMap a -> [(Vertex, a)]
IM.toList ((ResidualVertex -> Vertex) -> ResidualVertices -> IntMap Vertex
forall a b. (a -> b) -> IntMap a -> IntMap b
IM.map (\ResidualVertex
nv -> ResidualVertex -> Vertex
vv ResidualVertex
nv) ((ResidualVertex -> Bool) -> ResidualVertices -> ResidualVertices
forall a. (a -> Bool) -> IntMap a -> IntMap a
IM.filter (\ResidualVertex
nv -> ResidualVertex -> Excess
xv ResidualVertex
nv Excess -> Excess -> Bool
forall a. Ord a => a -> a -> Bool
> Excess
0) ResidualVertices
nvs))

-- | Push flow along a /forward/ edge @(u, v)@.
--
-- Preconditions (checked, returns 'Nothing' if not met):
--
-- * @height(u) = height(v) + 1@ (flow goes downhill)
-- * Residual capacity @c - f > 0@ (edge is not saturated)
-- * @excess(u) > 0@ (source vertex has excess to push)
--
-- Pushes @min(excess(u), c - f)@ units of flow.
-- Updates the 'topologyChanged' flag if the edge becomes saturated.
push :: ResidualGraph -> Edge -> Maybe ResidualGraph
push :: ResidualGraph -> Edge -> Maybe ResidualGraph
push ResidualGraph
g Edge
e =  
  let u :: Vertex
u = Edge -> Vertex
from Edge
e
      v :: Vertex
v = Edge -> Vertex
to Edge
e
      hu :: Vertex
hu = ResidualGraph -> Vertex -> Vertex
height ResidualGraph
g Vertex
u
      hv :: Vertex
hv = ResidualGraph -> Vertex -> Vertex
height ResidualGraph
g Vertex
v 
      xu :: Excess
xu = ResidualGraph -> Vertex -> Excess
excess ResidualGraph
g Vertex
u 
      xv :: Excess
xv = ResidualGraph -> Vertex -> Excess
excess ResidualGraph
g Vertex
v
      c :: Excess
c = ResidualGraph -> Edge -> Excess
edgeCapacity ResidualGraph
g Edge
e
      f :: Excess
f = ResidualGraph -> Edge -> Excess
edgeFlow ResidualGraph
g Edge
e
      nvs :: ResidualVertices
nvs = ResidualGraph -> ResidualVertices
netVertices ResidualGraph
g
      xf :: Excess
xf = Excess -> Excess -> Excess
forall a. Ord a => a -> a -> a
min Excess
xu (Excess
c Excess -> Excess -> Excess
forall a. Num a => a -> a -> a
- Excess
f)
   in if (Vertex
hu Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
hv Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1) Bool -> Bool -> Bool
&& Excess
xf Excess -> Excess -> Bool
forall a. Ord a => a -> a -> Bool
> Excess
0
         then
           let g' :: ResidualGraph
g' = ((ResidualGraph -> ResidualGraph)
 -> ResidualGraph -> ResidualGraph)
-> ResidualGraph
-> [ResidualGraph -> ResidualGraph]
-> ResidualGraph
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ResidualGraph -> ResidualGraph
f ResidualGraph
ac -> ResidualGraph -> ResidualGraph
f ResidualGraph
ac) ResidualGraph
g
                      [ (\ResidualGraph
nt -> ResidualGraph -> Edge -> Excess -> ResidualGraph
updateEdge ResidualGraph
nt Edge
e (Excess
f Excess -> Excess -> Excess
forall a. Num a => a -> a -> a
+ Excess
xf))
                      , (\ResidualGraph
nt -> ResidualGraph -> Vertex -> Excess -> ResidualGraph
updateExcess ResidualGraph
nt Vertex
u (Excess
xu Excess -> Excess -> Excess
forall a. Num a => a -> a -> a
- Excess
xf))
                      , (\ResidualGraph
nt -> ResidualGraph -> Vertex -> Excess -> ResidualGraph
updateExcess ResidualGraph
nt Vertex
v (Excess
xv Excess -> Excess -> Excess
forall a. Num a => a -> a -> a
+ Excess
xf))
                      ]
            in ResidualGraph -> Maybe ResidualGraph
forall a. a -> Maybe a
Just ResidualGraph
g'
         else Maybe ResidualGraph
forall a. Maybe a
Nothing 

-- | Pull flow along a /reverse/ edge @(u, v)@.
--
-- This is the dual of 'push': it decreases flow on edge @(u, v)@ by moving
-- excess from @v@ back to @u@.
--
-- Preconditions (checked, returns 'Nothing' if not met):
--
-- * @height(v) = height(u) + 1@ (pull goes uphill in the forward direction)
-- * @flow(u, v) > 0@ (there is flow to pull back)
-- * @excess(v) > 0@ (pulling vertex has excess)
--
-- Pulls @min(excess(v), flow)@ units.
-- Updates the 'topologyChanged' flag if the edge becomes zero-flow.
pull :: ResidualGraph -> Edge -> Maybe ResidualGraph
pull :: ResidualGraph -> Edge -> Maybe ResidualGraph
pull ResidualGraph
g Edge
e  = 
  let u :: Vertex
u   = Edge -> Vertex
from Edge
e
      v :: Vertex
v   = Edge -> Vertex
to Edge
e
      hu :: Vertex
hu  = ResidualGraph -> Vertex -> Vertex
height ResidualGraph
g Vertex
u
      hv :: Vertex
hv  = ResidualGraph -> Vertex -> Vertex
height ResidualGraph
g Vertex
v 
      xu :: Excess
xu  = ResidualGraph -> Vertex -> Excess
excess ResidualGraph
g Vertex
u 
      xv :: Excess
xv  = ResidualGraph -> Vertex -> Excess
excess ResidualGraph
g Vertex
v
      c :: Excess
c   = ResidualGraph -> Edge -> Excess
edgeCapacity ResidualGraph
g Edge
e
      f :: Excess
f   = ResidualGraph -> Edge -> Excess
edgeFlow ResidualGraph
g Edge
e
      nvs :: ResidualVertices
nvs = ResidualGraph -> ResidualVertices
netVertices ResidualGraph
g
      xf :: Excess
xf  = Excess -> Excess -> Excess
forall a. Ord a => a -> a -> a
min Excess
xv Excess
f
   in if (Vertex
hv Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
hu Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1) Bool -> Bool -> Bool
&& Excess
xf Excess -> Excess -> Bool
forall a. Ord a => a -> a -> Bool
> Excess
0 
         then
           let g' :: ResidualGraph
g' = ((ResidualGraph -> ResidualGraph)
 -> ResidualGraph -> ResidualGraph)
-> ResidualGraph
-> [ResidualGraph -> ResidualGraph]
-> ResidualGraph
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ResidualGraph -> ResidualGraph
f ResidualGraph
ac -> ResidualGraph -> ResidualGraph
f ResidualGraph
ac) ResidualGraph
g
                     [ (\ResidualGraph
nt -> ResidualGraph -> Edge -> Excess -> ResidualGraph
updateEdge ResidualGraph
nt Edge
e (Excess
f Excess -> Excess -> Excess
forall a. Num a => a -> a -> a
- Excess
xf))
                     , (\ResidualGraph
nt -> ResidualGraph -> Vertex -> Excess -> ResidualGraph
updateExcess ResidualGraph
nt Vertex
u (Excess
xu Excess -> Excess -> Excess
forall a. Num a => a -> a -> a
+ Excess
xf))
                     , (\ResidualGraph
nt -> ResidualGraph -> Vertex -> Excess -> ResidualGraph
updateExcess ResidualGraph
nt Vertex
v (Excess
xv Excess -> Excess -> Excess
forall a. Num a => a -> a -> a
- Excess
xf))
                     ]
            in ResidualGraph -> Maybe ResidualGraph
forall a. a -> Maybe a
Just ResidualGraph
g'
         else Maybe ResidualGraph
forall a. Maybe a
Nothing 

-- | Update the height of a vertex.  Source and sink heights are never modified.
updateHeight :: ResidualGraph -> Vertex -> Height -> ResidualGraph
updateHeight :: ResidualGraph -> Vertex -> Vertex -> ResidualGraph
updateHeight ResidualGraph
g Vertex
v Vertex
nh =
  let netvs :: ResidualVertices
netvs = ResidualGraph -> ResidualVertices
netVertices ResidualGraph
g
      !nv :: ResidualVertex
nv = Maybe ResidualVertex -> ResidualVertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe ResidualVertex -> ResidualVertex)
-> Maybe ResidualVertex -> ResidualVertex
forall a b. (a -> b) -> a -> b
$ Vertex -> ResidualVertices -> Maybe ResidualVertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v ResidualVertices
netvs
      !x :: Excess
x = ResidualGraph -> Vertex -> Excess
excess ResidualGraph
g Vertex
v
      !l :: Vertex
l = ResidualGraph -> Vertex -> Vertex
level ResidualGraph
g Vertex
v
      !s :: Vertex
s = Network -> Vertex
source (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
g
      !t :: Vertex
t = Network -> Vertex
sink (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
g
      !nnetv :: ResidualVertices
nnetv = (ResidualVertex -> Maybe ResidualVertex)
-> Vertex -> ResidualVertices -> ResidualVertices
forall a. (a -> Maybe a) -> Vertex -> IntMap a -> IntMap a
IM.update (\ResidualVertex
_ -> ResidualVertex -> Maybe ResidualVertex
forall a. a -> Maybe a
Just (Vertex -> Vertex -> Vertex -> Excess -> ResidualVertex
ResidualVertex Vertex
v Vertex
l Vertex
nh Excess
x)) Vertex
v ResidualVertices
netvs
  in if Vertex
v Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
t Bool -> Bool -> Bool
|| Vertex
v Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
s 
        then ResidualGraph
g
        else ResidualGraph
g { netVertices = nnetv }

-- | Update the excess of a vertex and maintain the 'overflowing' index.
--
-- When excess transitions between zero and non-zero, the vertex is
-- added to or removed from the 'Overflowing' map at its level.
-- Source and sink are excluded from the overflowing set.
updateExcess :: ResidualGraph -> Vertex -> Excess -> ResidualGraph
updateExcess :: ResidualGraph -> Vertex -> Excess -> ResidualGraph
updateExcess ResidualGraph
g Vertex
v Excess
nx =
  let netvs :: ResidualVertices
netvs = ResidualGraph -> ResidualVertices
netVertices ResidualGraph
g
      nv :: ResidualVertex
nv = Maybe ResidualVertex -> ResidualVertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe ResidualVertex -> ResidualVertex)
-> Maybe ResidualVertex -> ResidualVertex
forall a b. (a -> b) -> a -> b
$ Vertex -> ResidualVertices -> Maybe ResidualVertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v ResidualVertices
netvs
      h :: Vertex
h = ResidualGraph -> Vertex -> Vertex
height ResidualGraph
g Vertex
v
      l :: Vertex
l = ResidualGraph -> Vertex -> Vertex
level ResidualGraph
g Vertex
v
      ovfs :: Overflowing
ovfs = ResidualGraph -> Overflowing
overflowing ResidualGraph
g
      s :: Vertex
s = Network -> Vertex
source (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
g
      t :: Vertex
t = Network -> Vertex
sink (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
g
      newovfs :: Overflowing
newovfs = 
        if Vertex
v Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
s Bool -> Bool -> Bool
|| Vertex
v Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
t
           then Overflowing
ovfs
           else
             let ovfs' :: Overflowing
ovfs' = (IntSet -> Maybe IntSet) -> Vertex -> Overflowing -> Overflowing
forall a. (a -> Maybe a) -> Vertex -> IntMap a -> IntMap a
IM.update (\IntSet
lvs -> 
                         let lset :: IntSet
lset = Vertex -> IntSet -> IntSet
Set.delete Vertex
v IntSet
lvs
                          in if IntSet -> Bool
Set.null IntSet
lset
                                      then Maybe IntSet
forall a. Maybe a
Nothing 
                                      else IntSet -> Maybe IntSet
forall a. a -> Maybe a
Just IntSet
lset) Vertex
l Overflowing
ovfs
              in if Excess
nx Excess -> Excess -> Bool
forall a. Eq a => a -> a -> Bool
== Excess
0
                then 
                  Overflowing
ovfs'
                else 
                  let mlset :: Maybe IntSet
mlset = Vertex -> Overflowing -> Maybe IntSet
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
l Overflowing
ovfs'
                   in case Maybe IntSet
mlset of 
                        Maybe IntSet
Nothing -> Vertex -> IntSet -> Overflowing -> Overflowing
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
l (Vertex -> IntSet
Set.singleton Vertex
v) Overflowing
ovfs'
                        Just IntSet
lset -> (IntSet -> IntSet) -> Vertex -> Overflowing -> Overflowing
forall a. (a -> a) -> Vertex -> IntMap a -> IntMap a
IM.adjust (Vertex -> IntSet -> IntSet
Set.insert Vertex
v) Vertex
l Overflowing
ovfs'
   in if Vertex
v Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
t then ResidualGraph
g
                else ResidualGraph
g { netVertices = IM.insert v (ResidualVertex v l h nx) netvs
                       , overflowing = newovfs
                       } 

-- | Update the flow on an edge and track topology changes.
--
-- A topology change occurs when a forward residual edge appears or
-- disappears (flow crosses the capacity boundary) or a backward residual
-- edge appears or disappears (flow crosses zero).
-- The 'topologyChanged' flag is set to 'True' (OR-ed) if such a change occurs.
updateEdge :: ResidualGraph -> Edge -> Flow -> ResidualGraph
updateEdge :: ResidualGraph -> Edge -> Excess -> ResidualGraph
updateEdge ResidualGraph
g Edge
e Excess
f =
  let es :: ResidualEdges
es = ResidualGraph -> ResidualEdges
netEdges ResidualGraph
g
      eid :: Vertex
eid = Maybe Vertex -> Vertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ NeighborsMap -> Edge -> Maybe Vertex
resEdgeIndex (ResidualGraph -> NeighborsMap
netNeighborsMap ResidualGraph
g) Edge
e
      (ResidualEdge Edge
e' Excess
c Excess
f') = Maybe ResidualEdge -> ResidualEdge
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe ResidualEdge -> ResidualEdge)
-> Maybe ResidualEdge -> ResidualEdge
forall a b. (a -> b) -> a -> b
$ Vertex -> ResidualEdges -> Maybe ResidualEdge
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
eid ResidualEdges
es
      -- Detect if edge crossed a saturation boundary:
      -- forward edge exists iff flow < capacity
      -- backward edge exists iff flow > 0
      !fwdBefore :: Bool
fwdBefore = Excess
f' Excess -> Excess -> Bool
forall a. Ord a => a -> a -> Bool
< Excess
c
      !fwdAfter :: Bool
fwdAfter  = Excess
f Excess -> Excess -> Bool
forall a. Ord a => a -> a -> Bool
< Excess
c
      !bwdBefore :: Bool
bwdBefore = Excess
f' Excess -> Excess -> Bool
forall a. Ord a => a -> a -> Bool
> Excess
0
      !bwdAfter :: Bool
bwdAfter  = Excess
f Excess -> Excess -> Bool
forall a. Ord a => a -> a -> Bool
> Excess
0
      !changed :: Bool
changed   = (Bool
fwdBefore Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
/= Bool
fwdAfter) Bool -> Bool -> Bool
|| (Bool
bwdBefore Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
/= Bool
bwdAfter)
   in ResidualGraph
g { netEdges = IM.adjust (const (ResidualEdge e c f)) eid es
        , topologyChanged = topologyChanged g || changed
        }

-- | Net flow into the sink.  This is the current flow value of the network.
-- At termination, this equals the maximum flow.
netFlow :: ResidualGraph -> Flow
netFlow :: ResidualGraph -> Excess
netFlow ResidualGraph
g = ResidualGraph -> Vertex -> Excess
inflow ResidualGraph
g (Network -> Vertex
sink (ResidualGraph -> Network
network ResidualGraph
g))

-- | Height of a vertex.
height :: ResidualGraph -> Vertex -> Height
height :: ResidualGraph -> Vertex -> Vertex
height ResidualGraph
rg Vertex
v =
  let g :: Graph
g = Network -> Graph
graph (Network -> Graph) -> Network -> Graph
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      s :: Vertex
s = Network -> Vertex
source (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      t :: Vertex
t = Network -> Vertex
sink (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      nvs :: Integer
nvs = Vertex -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vertex -> Integer) -> Vertex -> Integer
forall a b. (a -> b) -> a -> b
$ Graph -> Vertex
numVertices Graph
g
      (ResidualVertex Vertex
nv Vertex
l Vertex
h Excess
x) = Maybe ResidualVertex -> ResidualVertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe ResidualVertex -> ResidualVertex)
-> Maybe ResidualVertex -> ResidualVertex
forall a b. (a -> b) -> a -> b
$ Vertex -> ResidualVertices -> Maybe ResidualVertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v (ResidualGraph -> ResidualVertices
netVertices ResidualGraph
rg)
   in Vertex
h

-- | Excess of a vertex.
excess :: ResidualGraph -> Vertex -> Excess
excess :: ResidualGraph -> Vertex -> Excess
excess ResidualGraph
rg Vertex
v =
  let g :: Graph
g = Network -> Graph
graph (Network -> Graph) -> Network -> Graph
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      s :: Vertex
s = Network -> Vertex
source (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      t :: Vertex
t = Network -> Vertex
sink (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      nvs :: Integer
nvs = Vertex -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vertex -> Integer) -> Vertex -> Integer
forall a b. (a -> b) -> a -> b
$ Graph -> Vertex
numVertices Graph
g
      (ResidualVertex Vertex
nv Vertex
l Vertex
h Excess
x) = Maybe ResidualVertex -> ResidualVertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe ResidualVertex -> ResidualVertex)
-> Maybe ResidualVertex -> ResidualVertex
forall a b. (a -> b) -> a -> b
$ Vertex -> ResidualVertices -> Maybe ResidualVertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v (ResidualGraph -> ResidualVertices
netVertices ResidualGraph
rg)
   in Excess
x

-- | Level of a vertex (shortest distance from source in original graph).
level :: ResidualGraph -> Vertex -> Level
level :: ResidualGraph -> Vertex -> Vertex
level ResidualGraph
rg Vertex
v =
  let g :: Graph
g = Network -> Graph
graph (Network -> Graph) -> Network -> Graph
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      s :: Vertex
s = Network -> Vertex
source (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      t :: Vertex
t = Network -> Vertex
sink (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      nvs :: Integer
nvs = Vertex -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vertex -> Integer) -> Vertex -> Integer
forall a b. (a -> b) -> a -> b
$ Graph -> Vertex
numVertices Graph
g
      (ResidualVertex Vertex
nv Vertex
l Vertex
h Excess
x) = Maybe ResidualVertex -> ResidualVertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe ResidualVertex -> ResidualVertex)
-> Maybe ResidualVertex -> ResidualVertex
forall a b. (a -> b) -> a -> b
$ Vertex -> ResidualVertices -> Maybe ResidualVertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v (ResidualGraph -> ResidualVertices
netVertices ResidualGraph
rg)
   in Vertex
l

-- | Capacity of an edge.
edgeCapacity :: ResidualGraph -> Edge -> Capacity
edgeCapacity :: ResidualGraph -> Edge -> Excess
edgeCapacity ResidualGraph
g Edge
e = let (ResidualEdge Edge
ne Excess
c Excess
f) = Maybe ResidualEdge -> ResidualEdge
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe ResidualEdge -> ResidualEdge)
-> Maybe ResidualEdge -> ResidualEdge
forall a b. (a -> b) -> a -> b
$ Vertex -> ResidualEdges -> Maybe ResidualEdge
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup (Maybe Vertex -> Vertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ NeighborsMap -> Edge -> Maybe Vertex
resEdgeIndex (ResidualGraph -> NeighborsMap
netNeighborsMap ResidualGraph
g) Edge
e) (ResidualGraph -> ResidualEdges
netEdges ResidualGraph
g)
                    in Excess
c 

-- | Current flow on an edge.
edgeFlow :: ResidualGraph -> Edge -> Flow
edgeFlow :: ResidualGraph -> Edge -> Excess
edgeFlow ResidualGraph
g Edge
e = let (ResidualEdge Edge
ne Excess
c Excess
f) = Maybe ResidualEdge -> ResidualEdge
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe ResidualEdge -> ResidualEdge)
-> Maybe ResidualEdge -> ResidualEdge
forall a b. (a -> b) -> a -> b
$ Vertex -> ResidualEdges -> Maybe ResidualEdge
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup (Maybe Vertex -> Vertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ NeighborsMap -> Edge -> Maybe Vertex
resEdgeIndex (ResidualGraph -> NeighborsMap
netNeighborsMap ResidualGraph
g) Edge
e) (ResidualGraph -> ResidualEdges
netEdges ResidualGraph
g)
                in Excess
f 

-- | Total flow into a vertex (sum of flows on incoming edges).
inflow :: ResidualGraph -> Vertex -> Flow
inflow :: ResidualGraph -> Vertex -> Excess
inflow ResidualGraph
g Vertex
v =
  let (IntMap Vertex
_, IntMap Vertex
revMap) = NeighborsMap -> Vertex -> (IntMap Vertex, IntMap Vertex)
netNeighbors (ResidualGraph -> NeighborsMap
netNeighborsMap ResidualGraph
g) Vertex
v 
      reds :: [Edge]
reds = (Vertex -> Edge) -> [Vertex] -> [Edge]
forall a b. (a -> b) -> [a] -> [b]
map (\Vertex
n -> (Vertex, Vertex) -> Edge
fromTuple (Vertex
n,Vertex
v)) ([Vertex] -> [Edge]) -> [Vertex] -> [Edge]
forall a b. (a -> b) -> a -> b
$ IntMap Vertex -> [Vertex]
forall a. IntMap a -> [Vertex]
IM.keys IntMap Vertex
revMap
   in (Excess -> Edge -> Excess) -> Excess -> [Edge] -> Excess
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Excess
ac Edge
e -> (Excess
ac Excess -> Excess -> Excess
forall a. Num a => a -> a -> a
+ ResidualGraph -> Edge -> Excess
edgeFlow ResidualGraph
g Edge
e)) Excess
0 [Edge]
reds 

-- | Total flow out of a vertex (sum of flows on outgoing edges).
outflow :: ResidualGraph -> Vertex -> Flow
outflow :: ResidualGraph -> Vertex -> Excess
outflow ResidualGraph
g Vertex
v =
  let (IntMap Vertex
fwdMap, IntMap Vertex
_) = NeighborsMap -> Vertex -> (IntMap Vertex, IntMap Vertex)
netNeighbors (ResidualGraph -> NeighborsMap
netNeighborsMap ResidualGraph
g) Vertex
v 
      reds :: [Edge]
reds = (Vertex -> Edge) -> [Vertex] -> [Edge]
forall a b. (a -> b) -> [a] -> [b]
map (\Vertex
n -> (Vertex, Vertex) -> Edge
fromTuple (Vertex
v,Vertex
n)) ([Vertex] -> [Edge]) -> [Vertex] -> [Edge]
forall a b. (a -> b) -> a -> b
$ IntMap Vertex -> [Vertex]
forall a. IntMap a -> [Vertex]
IM.keys IntMap Vertex
fwdMap
   in (Excess -> Edge -> Excess) -> Excess -> [Edge] -> Excess
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Excess
ac Edge
e -> (Excess
ac Excess -> Excess -> Excess
forall a. Num a => a -> a -> a
+ ResidualGraph -> Edge -> Excess
edgeFlow ResidualGraph
g Edge
e)) Excess
0 [Edge]
reds 

-- | Reconstruct the 'Network' with final edge flows from the residual graph.
-- Called when the algorithm terminates.
networkFromResidual :: ResidualGraph -> Network
networkFromResidual :: ResidualGraph -> Network
networkFromResidual ResidualGraph
resg =
  let net :: Network
net = ResidualGraph -> Network
network ResidualGraph
resg
      es :: [Edge]
es = Graph -> [Edge]
edges (Graph -> [Edge]) -> Graph -> [Edge]
forall a b. (a -> b) -> a -> b
$ Network -> Graph
graph (Network -> Graph) -> Network -> Graph
forall a b. (a -> b) -> a -> b
$ Network
net
      flow' :: Capacities
flow' = [(Edge, Excess)] -> Capacities
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Edge, Excess)] -> Capacities) -> [(Edge, Excess)] -> Capacities
forall a b. (a -> b) -> a -> b
$ (Edge -> (Edge, Excess)) -> [Edge] -> [(Edge, Excess)]
forall a b. (a -> b) -> [a] -> [b]
map (\Edge
e -> (Edge
e, ResidualGraph -> Edge -> Excess
edgeFlow ResidualGraph
resg Edge
e) ) [Edge]
es
   in Network
net {flow = flow'}

-- | Compute distances from source and sink in the residual graph via BFS.
--
-- Returns @(sourceDists, sinkDists)@ where:
--
-- * @sourceDists@: @IntMap@ from vertex to BFS distance from source
--   (traversing edges with residual capacity > 0 in reverse, and edges
--   with flow > 0 forward)
-- * @sinkDists@: @IntMap@ from vertex to BFS distance from sink
--   (traversing edges with residual capacity > 0 forward, and edges
--   with flow > 0 in reverse)
--
-- Used by @globalRelabel@ to set vertex heights:
-- source-side vertices get @height = |V| + dist_from_source@,
-- sink-side vertices get @height = dist_from_sink@.
residualDistances :: ResidualGraph -> (IM.IntMap Int, IM.IntMap Int)
residualDistances :: ResidualGraph -> (IntMap Vertex, IntMap Vertex)
residualDistances ResidualGraph
rg = 
  let es :: [ResidualEdge]
es = ((Vertex, ResidualEdge) -> ResidualEdge)
-> [(Vertex, ResidualEdge)] -> [ResidualEdge]
forall a b. (a -> b) -> [a] -> [b]
map (Vertex, ResidualEdge) -> ResidualEdge
forall a b. (a, b) -> b
snd (ResidualEdges -> [(Vertex, ResidualEdge)]
forall a. IntMap a -> [(Vertex, a)]
IM.toList (ResidualEdges -> [(Vertex, ResidualEdge)])
-> ResidualEdges -> [(Vertex, ResidualEdge)]
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> ResidualEdges
netEdges ResidualGraph
rg)
      -- forward residual edges (flow < capacity)
      tres :: [ResidualEdge]
tres = (ResidualEdge -> Bool) -> [ResidualEdge] -> [ResidualEdge]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(ResidualEdge Edge
e Excess
c Excess
f) -> Excess
f Excess -> Excess -> Bool
forall a. Ord a => a -> a -> Bool
< Excess
c) [ResidualEdge]
es
      -- backward residual edges (flow > 0)
      tbes :: [ResidualEdge]
tbes = (ResidualEdge -> Bool) -> [ResidualEdge] -> [ResidualEdge]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(ResidualEdge Edge
e Excess
c Excess
f) -> Excess
f Excess -> Excess -> Bool
forall a. Ord a => a -> a -> Bool
> Excess
0) [ResidualEdge]
es
      tfsatnbs :: IntMap [Vertex]
tfsatnbs = (IntMap [Vertex] -> ResidualEdge -> IntMap [Vertex])
-> IntMap [Vertex] -> [ResidualEdge] -> IntMap [Vertex]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap [Vertex]
ac (ResidualEdge Edge
e Excess
c Excess
f) -> 
        let u :: Vertex
u = Edge -> Vertex
from Edge
e
            v :: Vertex
v = Edge -> Vertex
to Edge
e 
            mns :: Maybe [Vertex]
mns = Vertex -> IntMap [Vertex] -> Maybe [Vertex]
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v IntMap [Vertex]
ac 
         in case Maybe [Vertex]
mns of 
               Maybe [Vertex]
Nothing -> Vertex -> [Vertex] -> IntMap [Vertex] -> IntMap [Vertex]
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
v [Vertex
u] IntMap [Vertex]
ac
               Just [Vertex]
ns -> Vertex -> [Vertex] -> IntMap [Vertex] -> IntMap [Vertex]
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
v (Vertex
uVertex -> [Vertex] -> [Vertex]
forall a. a -> [a] -> [a]
:[Vertex]
ns) IntMap [Vertex]
ac
             ) IntMap [Vertex]
forall a. IntMap a
IM.empty [ResidualEdge]
tres
      tsatnbs :: IntMap [Vertex]
tsatnbs = (IntMap [Vertex] -> ResidualEdge -> IntMap [Vertex])
-> IntMap [Vertex] -> [ResidualEdge] -> IntMap [Vertex]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap [Vertex]
ac (ResidualEdge Edge
e Excess
c Excess
f) -> 
        let u :: Vertex
u = Edge -> Vertex
from Edge
e
            v :: Vertex
v = Edge -> Vertex
to Edge
e 
            mns :: Maybe [Vertex]
mns = Vertex -> IntMap [Vertex] -> Maybe [Vertex]
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
u IntMap [Vertex]
ac 
         in case Maybe [Vertex]
mns of 
               Maybe [Vertex]
Nothing -> Vertex -> [Vertex] -> IntMap [Vertex] -> IntMap [Vertex]
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
u [Vertex
v] IntMap [Vertex]
ac
               Just [Vertex]
ns -> Vertex -> [Vertex] -> IntMap [Vertex] -> IntMap [Vertex]
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
u (Vertex
vVertex -> [Vertex] -> [Vertex]
forall a. a -> [a] -> [a]
:[Vertex]
ns) IntMap [Vertex]
ac
             ) IntMap [Vertex]
tfsatnbs [ResidualEdge]
tbes
      sfsatnbs :: IntMap [Vertex]
sfsatnbs = (IntMap [Vertex] -> ResidualEdge -> IntMap [Vertex])
-> IntMap [Vertex] -> [ResidualEdge] -> IntMap [Vertex]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap [Vertex]
ac (ResidualEdge Edge
e Excess
c Excess
f) -> 
        let u :: Vertex
u = Edge -> Vertex
from Edge
e
            v :: Vertex
v = Edge -> Vertex
to Edge
e 
            mns :: Maybe [Vertex]
mns = Vertex -> IntMap [Vertex] -> Maybe [Vertex]
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
u IntMap [Vertex]
ac 
         in case Maybe [Vertex]
mns of 
               Maybe [Vertex]
Nothing -> Vertex -> [Vertex] -> IntMap [Vertex] -> IntMap [Vertex]
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
u [Vertex
v] IntMap [Vertex]
ac
               Just [Vertex]
ns -> Vertex -> [Vertex] -> IntMap [Vertex] -> IntMap [Vertex]
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
u (Vertex
vVertex -> [Vertex] -> [Vertex]
forall a. a -> [a] -> [a]
:[Vertex]
ns) IntMap [Vertex]
ac
             ) IntMap [Vertex]
forall a. IntMap a
IM.empty [ResidualEdge]
tres
      ssatnbs :: IntMap [Vertex]
ssatnbs = (IntMap [Vertex] -> ResidualEdge -> IntMap [Vertex])
-> IntMap [Vertex] -> [ResidualEdge] -> IntMap [Vertex]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap [Vertex]
ac (ResidualEdge Edge
e Excess
c Excess
f) -> 
        let u :: Vertex
u = Edge -> Vertex
from Edge
e
            v :: Vertex
v = Edge -> Vertex
to Edge
e 
            mns :: Maybe [Vertex]
mns = Vertex -> IntMap [Vertex] -> Maybe [Vertex]
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v IntMap [Vertex]
ac 
         in case Maybe [Vertex]
mns of 
               Maybe [Vertex]
Nothing -> Vertex -> [Vertex] -> IntMap [Vertex] -> IntMap [Vertex]
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
v [Vertex
u] IntMap [Vertex]
ac
               Just [Vertex]
ns -> Vertex -> [Vertex] -> IntMap [Vertex] -> IntMap [Vertex]
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
v (Vertex
uVertex -> [Vertex] -> [Vertex]
forall a. a -> [a] -> [a]
:[Vertex]
ns) IntMap [Vertex]
ac
             ) IntMap [Vertex]
sfsatnbs [ResidualEdge]
tbes
      tlvs :: IntMap Vertex
tlvs = BFS -> IntMap Vertex
BFS.level (BFS -> IntMap Vertex) -> BFS -> IntMap Vertex
forall a b. (a -> b) -> a -> b
$ IntMap [Vertex] -> Vertex -> BFS
BFS.adjBFS IntMap [Vertex]
tsatnbs Vertex
t
      slvs :: IntMap Vertex
slvs = BFS -> IntMap Vertex
BFS.level (BFS -> IntMap Vertex) -> BFS -> IntMap Vertex
forall a b. (a -> b) -> a -> b
$ IntMap [Vertex] -> Vertex -> BFS
BFS.adjBFS IntMap [Vertex]
ssatnbs Vertex
s
    in (IntMap Vertex
slvs, IntMap Vertex
tlvs)
  where
    g :: Graph
g = Network -> Graph
graph (Network -> Graph) -> Network -> Graph
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
    s :: Vertex
s = Network -> Vertex
source (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
    t :: Vertex
t = Network -> Vertex
sink (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg

-- | Compute the source-sink minimum cut from the residual graph.
--
-- Returns @(S, T)@ where @S@ is the set of vertices reachable from the
-- source in the residual graph (excluding source and sink) and @T@ is
-- the complement.  By the max-flow min-cut theorem, the total capacity
-- of edges crossing from @S@ to @T@ equals the maximum flow.
stCut :: ResidualGraph -> ([Vertex],[Vertex])
stCut :: ResidualGraph -> ([Vertex], [Vertex])
stCut ResidualGraph
rg = 
  let !resdis :: (IntMap Vertex, IntMap Vertex)
resdis = ResidualGraph -> (IntMap Vertex, IntMap Vertex)
residualDistances ResidualGraph
rg
      ts :: IntSet
ts = Vertex -> IntSet -> IntSet
Set.delete Vertex
s (IntSet -> IntSet) -> IntSet -> IntSet
forall a b. (a -> b) -> a -> b
$ Vertex -> IntSet -> IntSet
Set.delete Vertex
t (IntSet -> IntSet) -> IntSet -> IntSet
forall a b. (a -> b) -> a -> b
$ [Vertex] -> IntSet
Set.fromList ([Vertex] -> IntSet) -> [Vertex] -> IntSet
forall a b. (a -> b) -> a -> b
$ ((Vertex, Vertex) -> Vertex) -> [(Vertex, Vertex)] -> [Vertex]
forall a b. (a -> b) -> [a] -> [b]
map (Vertex, Vertex) -> Vertex
forall a b. (a, b) -> a
fst (IntMap Vertex -> [(Vertex, Vertex)]
forall a. IntMap a -> [(Vertex, a)]
IM.toList ((IntMap Vertex, IntMap Vertex) -> IntMap Vertex
forall a b. (a, b) -> b
snd (IntMap Vertex, IntMap Vertex)
resdis))
      g :: Graph
g = Network -> Graph
graph (Network -> Graph) -> Network -> Graph
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      s :: Vertex
s = Network -> Vertex
source (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      t :: Vertex
t = Network -> Vertex
sink (Network -> Vertex) -> Network -> Vertex
forall a b. (a -> b) -> a -> b
$ ResidualGraph -> Network
network ResidualGraph
rg
      vs :: IntSet
vs = Vertex -> IntSet -> IntSet
Set.delete Vertex
s (IntSet -> IntSet) -> IntSet -> IntSet
forall a b. (a -> b) -> a -> b
$ Vertex -> IntSet -> IntSet
Set.delete Vertex
t (IntSet -> IntSet) -> IntSet -> IntSet
forall a b. (a -> b) -> a -> b
$ [Vertex] -> IntSet
Set.fromList ([Vertex] -> IntSet) -> [Vertex] -> IntSet
forall a b. (a -> b) -> a -> b
$ Graph -> [Vertex]
vertices Graph
g
      ss :: IntSet
ss = IntSet -> IntSet -> IntSet
Set.difference IntSet
vs IntSet
ts
   in (IntSet -> [Vertex]
Set.toList IntSet
ss, IntSet -> [Vertex]
Set.toList IntSet
ts)