{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}

module Data.Graph.MinCostFlow where

import Control.Exception
import Control.Monad
import Control.Monad.Primitive
import Control.Monad.ST
import Data.Bits
import Data.Function
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import Data.Word
import Unsafe.Coerce

import Data.Buffer
import Data.Heap.Binary
import My.Prelude (rep)

nothingMCF :: Int
nothingMCF :: Cost
nothingMCF = -Cost
1

type Vertex = Int
type Cost = Int
type Capacity = Int

{- |
Primal Dual /O(FElog V)/

>>> :{
minCostFlow 2 0 1 2 (\builder -> do
    addEdgeMCFB builder 0 1 123 2
    )
:}
(246,2)
>>> :{
minCostFlow 2 0 1 123456789 (\builder -> do
    addEdgeMCFB builder 0 1 123 2
    )
:}
(246,2)
-}
minCostFlow ::
  -- | number of vertices
  Int ->
  -- | source
  Vertex ->
  -- | sink
  Vertex ->
  -- | flow
  Capacity ->
  (forall s. MinCostFlowBuilder s -> ST s ()) ->
  (Cost, Capacity)
minCostFlow :: Cost
-> Cost
-> Cost
-> Cost
-> (forall s. MinCostFlowBuilder s -> ST s ())
-> (Cost, Cost)
minCostFlow Cost
numVertices Cost
src Cost
sink Cost
flow forall s. MinCostFlowBuilder s -> ST s ()
run = (forall s. ST s (Cost, Cost)) -> (Cost, Cost)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Cost, Cost)) -> (Cost, Cost))
-> (forall s. ST s (Cost, Cost)) -> (Cost, Cost)
forall a b. (a -> b) -> a -> b
$ do
  builder <- Cost -> ST s (MinCostFlowBuilder (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Cost -> m (MinCostFlowBuilder (PrimState m))
newMinCostFlowBuilder Cost
numVertices
  run builder
  buildMinCostFlow builder >>= runMinCostFlow src sink flow

data MinCostFlow s = MinCostFlow
  { forall s. MinCostFlow s -> Cost
numVerticesMCF :: !Int
  , forall s. MinCostFlow s -> Cost
numEdgesMCF :: !Int
  , forall s. MinCostFlow s -> Vector Cost
offsetMCF :: U.Vector Int
  , forall s. MinCostFlow s -> Vector Cost
dstMCF :: U.Vector Vertex
  , forall s. MinCostFlow s -> Vector Cost
costMCF :: U.Vector Cost
  , forall s. MinCostFlow s -> MVector s Cost
residualMCF :: UM.MVector s Capacity
  , forall s. MinCostFlow s -> MVector s Cost
potentialMCF :: UM.MVector s Cost
  , forall s. MinCostFlow s -> MVector s Cost
distMCF :: UM.MVector s Cost
  , forall s. MinCostFlow s -> MinBinaryHeap s Word64
heapMCF :: MinBinaryHeap s Word64 -- (Cost, Vertex)
  , forall s. MinCostFlow s -> Vector Cost
revEdgeMCF :: U.Vector Int
  , forall s. MinCostFlow s -> MVector s Cost
prevVertexMCF :: UM.MVector s Vertex
  , forall s. MinCostFlow s -> MVector s Cost
prevEdgeMCF :: UM.MVector s Int
  }

runMinCostFlow ::
  (PrimMonad m) =>
  Vertex ->
  Vertex ->
  Capacity ->
  MinCostFlow (PrimState m) ->
  m (Cost, Capacity)
runMinCostFlow :: forall (m :: * -> *).
PrimMonad m =>
Cost -> Cost -> Cost -> MinCostFlow (PrimState m) -> m (Cost, Cost)
runMinCostFlow Cost
source Cost
sink Cost
flow mcf :: MinCostFlow (PrimState m)
mcf@MinCostFlow{Cost
MVector (PrimState m) Cost
Vector Cost
MinBinaryHeap (PrimState m) Word64
numVerticesMCF :: forall s. MinCostFlow s -> Cost
numEdgesMCF :: forall s. MinCostFlow s -> Cost
offsetMCF :: forall s. MinCostFlow s -> Vector Cost
dstMCF :: forall s. MinCostFlow s -> Vector Cost
costMCF :: forall s. MinCostFlow s -> Vector Cost
residualMCF :: forall s. MinCostFlow s -> MVector s Cost
potentialMCF :: forall s. MinCostFlow s -> MVector s Cost
distMCF :: forall s. MinCostFlow s -> MVector s Cost
heapMCF :: forall s. MinCostFlow s -> MinBinaryHeap s Word64
revEdgeMCF :: forall s. MinCostFlow s -> Vector Cost
prevVertexMCF :: forall s. MinCostFlow s -> MVector s Cost
prevEdgeMCF :: forall s. MinCostFlow s -> MVector s Cost
numVerticesMCF :: Cost
numEdgesMCF :: Cost
offsetMCF :: Vector Cost
dstMCF :: Vector Cost
costMCF :: Vector Cost
residualMCF :: MVector (PrimState m) Cost
potentialMCF :: MVector (PrimState m) Cost
distMCF :: MVector (PrimState m) Cost
heapMCF :: MinBinaryHeap (PrimState m) Word64
revEdgeMCF :: Vector Cost
prevVertexMCF :: MVector (PrimState m) Cost
prevEdgeMCF :: MVector (PrimState m) Cost
..} = Cost -> Cost -> m (Cost, Cost)
forall {m :: * -> *}.
(PrimState m ~ PrimState m, PrimMonad m) =>
Cost -> Cost -> m (Cost, Cost)
go Cost
0 Cost
flow
  where
    go :: Cost -> Cost -> m (Cost, Cost)
go !Cost
res !Cost
f
      | Cost
f Cost -> Cost -> Bool
forall a. Eq a => a -> a -> Bool
== Cost
0 = (Cost, Cost) -> m (Cost, Cost)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Cost
res, Cost
flow)
      | Bool
otherwise = do
          canFlow <- Cost -> Cost -> MinCostFlow (PrimState m) -> m Bool
forall (m :: * -> *).
PrimMonad m =>
Cost -> Cost -> MinCostFlow (PrimState m) -> m Bool
dijkstraMCF Cost
source Cost
sink MinCostFlow (PrimState m)
MinCostFlow (PrimState m)
mcf
          if canFlow
            then do
              rep numVerticesMCF $ \Cost
v -> do
                dv <- MVector (PrimState m) Cost -> Cost -> m Cost
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Cost -> m a
UM.unsafeRead MVector (PrimState m) Cost
MVector (PrimState m) Cost
distMCF Cost
v
                UM.unsafeModify potentialMCF (+ dv) v
              flowed <- updateResidualMCF sink f mcf
              hsink <- UM.unsafeRead potentialMCF sink
              go (hsink * flowed + res) (f - flowed)
            else return (res, flow - f)

-- | cost 48bit / vertex 16bit
encodeMCF :: Cost -> Vertex -> Word64
encodeMCF :: Cost -> Cost -> Word64
encodeMCF Cost
cost Cost
v = Cost -> Word64
forall a b. a -> b
unsafeCoerce (Cost -> Word64) -> Cost -> Word64
forall a b. (a -> b) -> a -> b
$ Cost -> Cost -> Cost
forall a. Bits a => a -> Cost -> a
unsafeShiftL Cost
cost Cost
16 Cost -> Cost -> Cost
forall a. Bits a => a -> a -> a
.|. Cost
v
{-# INLINE encodeMCF #-}

decodeMCF :: Word64 -> (Cost, Vertex)
decodeMCF :: Word64 -> (Cost, Cost)
decodeMCF Word64
costv = (Word64, Word64) -> (Cost, Cost)
forall a b. a -> b
unsafeCoerce (Word64
cost, Word64
v)
  where
    !cost :: Word64
cost = Word64 -> Cost -> Word64
forall a. Bits a => a -> Cost -> a
unsafeShiftR Word64
costv Cost
16
    !v :: Word64
v = Word64
costv Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0xffff
{-# INLINE decodeMCF #-}

dijkstraMCF ::
  (PrimMonad m) =>
  Vertex ->
  Vertex ->
  MinCostFlow (PrimState m) ->
  m Bool
dijkstraMCF :: forall (m :: * -> *).
PrimMonad m =>
Cost -> Cost -> MinCostFlow (PrimState m) -> m Bool
dijkstraMCF Cost
source Cost
sink MinCostFlow{Cost
MVector (PrimState m) Cost
Vector Cost
MinBinaryHeap (PrimState m) Word64
numVerticesMCF :: forall s. MinCostFlow s -> Cost
numEdgesMCF :: forall s. MinCostFlow s -> Cost
offsetMCF :: forall s. MinCostFlow s -> Vector Cost
dstMCF :: forall s. MinCostFlow s -> Vector Cost
costMCF :: forall s. MinCostFlow s -> Vector Cost
residualMCF :: forall s. MinCostFlow s -> MVector s Cost
potentialMCF :: forall s. MinCostFlow s -> MVector s Cost
distMCF :: forall s. MinCostFlow s -> MVector s Cost
heapMCF :: forall s. MinCostFlow s -> MinBinaryHeap s Word64
revEdgeMCF :: forall s. MinCostFlow s -> Vector Cost
prevVertexMCF :: forall s. MinCostFlow s -> MVector s Cost
prevEdgeMCF :: forall s. MinCostFlow s -> MVector s Cost
numVerticesMCF :: Cost
numEdgesMCF :: Cost
offsetMCF :: Vector Cost
dstMCF :: Vector Cost
costMCF :: Vector Cost
residualMCF :: MVector (PrimState m) Cost
potentialMCF :: MVector (PrimState m) Cost
distMCF :: MVector (PrimState m) Cost
heapMCF :: MinBinaryHeap (PrimState m) Word64
revEdgeMCF :: Vector Cost
prevVertexMCF :: MVector (PrimState m) Cost
prevEdgeMCF :: MVector (PrimState m) Cost
..} = do
  MVector (PrimState m) Cost -> Cost -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
UM.set MVector (PrimState m) Cost
distMCF Cost
forall a. Bounded a => a
maxBound
  MVector (PrimState m) Cost -> Cost -> Cost -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Cost -> a -> m ()
UM.unsafeWrite MVector (PrimState m) Cost
distMCF Cost
source Cost
0
  MinBinaryHeap (PrimState m) Word64 -> m ()
forall (m :: * -> *) (f :: * -> *) a.
PrimMonad m =>
BinaryHeap f (PrimState m) a -> m ()
clearBH MinBinaryHeap (PrimState m) Word64
heapMCF
  Word64 -> MinBinaryHeap (PrimState m) Word64 -> m ()
forall (f :: * -> *) a (m :: * -> *).
(OrdVia f a, Unbox a, PrimMonad m) =>
a -> BinaryHeap f (PrimState m) a -> m ()
insertBH (Cost -> Cost -> Word64
encodeMCF Cost
0 Cost
source) MinBinaryHeap (PrimState m) Word64
heapMCF

  (m Bool -> m Bool) -> m Bool
forall a. (a -> a) -> a
fix ((m Bool -> m Bool) -> m Bool) -> (m Bool -> m Bool) -> m Bool
forall a b. (a -> b) -> a -> b
$ \m Bool
loop -> do
    MinBinaryHeap (PrimState m) Word64 -> m (Maybe Word64)
forall (f :: * -> *) a (m :: * -> *).
(OrdVia f a, Unbox a, PrimMonad m) =>
BinaryHeap f (PrimState m) a -> m (Maybe a)
deleteFindTopBH MinBinaryHeap (PrimState m) Word64
heapMCF m (Maybe Word64) -> (Maybe Word64 -> m Bool) -> m Bool
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just Word64
cv -> do
        let (Cost
c, Cost
v) = Word64 -> (Cost, Cost)
decodeMCF Word64
cv
        dv <- MVector (PrimState m) Cost -> Cost -> m Cost
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Cost -> m a
UM.unsafeRead MVector (PrimState m) Cost
distMCF Cost
v
        unless (c > dv) $ do
          let start = Vector Cost -> Cost -> Cost
forall a. Unbox a => Vector a -> Cost -> a
U.unsafeIndex Vector Cost
offsetMCF Cost
v
          let end = Vector Cost -> Cost -> Cost
forall a. Unbox a => Vector a -> Cost -> a
U.unsafeIndex Vector Cost
offsetMCF (Cost
v Cost -> Cost -> Cost
forall a. Num a => a -> a -> a
+ Cost
1)
          U.forM_ (U.generate (end - start) (+ start)) $ \Cost
e -> do
            let nv :: Cost
nv = Vector Cost -> Cost -> Cost
forall a. Unbox a => Vector a -> Cost -> a
U.unsafeIndex Vector Cost
dstMCF Cost
e
            let v2nv :: Cost
v2nv = Vector Cost -> Cost -> Cost
forall a. Unbox a => Vector a -> Cost -> a
U.unsafeIndex Vector Cost
costMCF Cost
e
            cap <- MVector (PrimState m) Cost -> Cost -> m Cost
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Cost -> m a
UM.unsafeRead MVector (PrimState m) Cost
residualMCF Cost
e
            hv <- UM.unsafeRead potentialMCF v
            hnv <- UM.unsafeRead potentialMCF nv
            old <- UM.unsafeRead distMCF nv
            let dnv = Cost
dv Cost -> Cost -> Cost
forall a. Num a => a -> a -> a
+ Cost
v2nv Cost -> Cost -> Cost
forall a. Num a => a -> a -> a
+ Cost
hv Cost -> Cost -> Cost
forall a. Num a => a -> a -> a
- Cost
hnv
            when (cap > 0 && dnv < old) $ do
              UM.unsafeWrite distMCF nv dnv
              UM.unsafeWrite prevVertexMCF nv v
              UM.unsafeWrite prevEdgeMCF nv e
              insertBH (encodeMCF dnv nv) heapMCF
        loop
      Maybe Word64
Nothing -> do
        cost <- MVector (PrimState m) Cost -> Cost -> m Cost
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Cost -> m a
UM.unsafeRead MVector (PrimState m) Cost
distMCF Cost
sink
        return $! cost < maxBound
{-# INLINE dijkstraMCF #-}

updateResidualMCF ::
  (PrimMonad m) =>
  Vertex ->
  Capacity ->
  MinCostFlow (PrimState m) ->
  m Capacity
updateResidualMCF :: forall (m :: * -> *).
PrimMonad m =>
Cost -> Cost -> MinCostFlow (PrimState m) -> m Cost
updateResidualMCF Cost
sink Cost
flow MinCostFlow{Cost
MVector (PrimState m) Cost
Vector Cost
MinBinaryHeap (PrimState m) Word64
numVerticesMCF :: forall s. MinCostFlow s -> Cost
numEdgesMCF :: forall s. MinCostFlow s -> Cost
offsetMCF :: forall s. MinCostFlow s -> Vector Cost
dstMCF :: forall s. MinCostFlow s -> Vector Cost
costMCF :: forall s. MinCostFlow s -> Vector Cost
residualMCF :: forall s. MinCostFlow s -> MVector s Cost
potentialMCF :: forall s. MinCostFlow s -> MVector s Cost
distMCF :: forall s. MinCostFlow s -> MVector s Cost
heapMCF :: forall s. MinCostFlow s -> MinBinaryHeap s Word64
revEdgeMCF :: forall s. MinCostFlow s -> Vector Cost
prevVertexMCF :: forall s. MinCostFlow s -> MVector s Cost
prevEdgeMCF :: forall s. MinCostFlow s -> MVector s Cost
numVerticesMCF :: Cost
numEdgesMCF :: Cost
offsetMCF :: Vector Cost
dstMCF :: Vector Cost
costMCF :: Vector Cost
residualMCF :: MVector (PrimState m) Cost
potentialMCF :: MVector (PrimState m) Cost
distMCF :: MVector (PrimState m) Cost
heapMCF :: MinBinaryHeap (PrimState m) Word64
revEdgeMCF :: Vector Cost
prevVertexMCF :: MVector (PrimState m) Cost
prevEdgeMCF :: MVector (PrimState m) Cost
..} = Cost -> Cost -> (Cost -> m Cost) -> m Cost
forall {m :: * -> *} {b}.
(PrimState m ~ PrimState m, PrimMonad m) =>
Cost -> Cost -> (Cost -> m b) -> m b
go Cost
sink Cost
flow Cost -> m Cost
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
  where
    go :: Cost -> Cost -> (Cost -> m b) -> m b
go !Cost
v !Cost
f Cost -> m b
k = do
      pv <- MVector (PrimState m) Cost -> Cost -> m Cost
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Cost -> m a
UM.unsafeRead MVector (PrimState m) Cost
MVector (PrimState m) Cost
prevVertexMCF Cost
v
      if pv < 0
        then k f
        else do
          pv2v <- UM.unsafeRead prevEdgeMCF v
          f' <- UM.unsafeRead residualMCF pv2v
          go pv (min f f') $ \Cost
nf -> do
            MVector (PrimState m) Cost -> (Cost -> Cost) -> Cost -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Cost -> m ()
UM.unsafeModify MVector (PrimState m) Cost
MVector (PrimState m) Cost
residualMCF (Cost -> Cost -> Cost
forall a. Num a => a -> a -> a
subtract Cost
nf) Cost
pv2v
            MVector (PrimState m) Cost -> (Cost -> Cost) -> Cost -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Cost -> m ()
UM.unsafeModify MVector (PrimState m) Cost
MVector (PrimState m) Cost
residualMCF (Cost -> Cost -> Cost
forall a. Num a => a -> a -> a
+ Cost
nf) (Vector Cost -> Cost -> Cost
forall a. Unbox a => Vector a -> Cost -> a
U.unsafeIndex Vector Cost
revEdgeMCF Cost
pv2v)
            Cost -> m b
k Cost
nf
{-# INLINE updateResidualMCF #-}

data MinCostFlowBuilder s = MinCostFlowBuilder
  { forall s. MinCostFlowBuilder s -> Cost
numVerticesMCFB :: !Int
  , forall s. MinCostFlowBuilder s -> MVector s Cost
inDegreeMCFB :: UM.MVector s Int
  , forall s. MinCostFlowBuilder s -> Buffer s (Cost, Cost, Cost, Cost)
edgesMCFB :: Buffer s (Vertex, Vertex, Cost, Capacity)
  -- ^ default buffer size: /1024 * 1024/
  }

newMinCostFlowBuilder ::
  (PrimMonad m) =>
  Int ->
  m (MinCostFlowBuilder (PrimState m))
newMinCostFlowBuilder :: forall (m :: * -> *).
PrimMonad m =>
Cost -> m (MinCostFlowBuilder (PrimState m))
newMinCostFlowBuilder Cost
n =
  Cost
-> MVector (PrimState m) Cost
-> Buffer (PrimState m) (Cost, Cost, Cost, Cost)
-> MinCostFlowBuilder (PrimState m)
forall s.
Cost
-> MVector s Cost
-> Buffer s (Cost, Cost, Cost, Cost)
-> MinCostFlowBuilder s
MinCostFlowBuilder Cost
n
    (MVector (PrimState m) Cost
 -> Buffer (PrimState m) (Cost, Cost, Cost, Cost)
 -> MinCostFlowBuilder (PrimState m))
-> m (MVector (PrimState m) Cost)
-> m (Buffer (PrimState m) (Cost, Cost, Cost, Cost)
      -> MinCostFlowBuilder (PrimState m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Cost -> Cost -> m (MVector (PrimState m) Cost)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Cost -> a -> m (MVector (PrimState m) a)
UM.replicate Cost
n Cost
0
    m (Buffer (PrimState m) (Cost, Cost, Cost, Cost)
   -> MinCostFlowBuilder (PrimState m))
-> m (Buffer (PrimState m) (Cost, Cost, Cost, Cost))
-> m (MinCostFlowBuilder (PrimState m))
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Cost -> m (Buffer (PrimState m) (Cost, Cost, Cost, Cost))
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Cost -> m (Buffer (PrimState m) a)
newBuffer (Cost
1024 Cost -> Cost -> Cost
forall a. Num a => a -> a -> a
* Cost
1024)

-- | /cost >= 0/
addEdgeMCFB ::
  (PrimMonad m) =>
  MinCostFlowBuilder (PrimState m) ->
  Vertex ->
  Vertex ->
  Cost ->
  Capacity ->
  m ()
addEdgeMCFB :: forall (m :: * -> *).
PrimMonad m =>
MinCostFlowBuilder (PrimState m)
-> Cost -> Cost -> Cost -> Cost -> m ()
addEdgeMCFB MinCostFlowBuilder{Cost
MVector (PrimState m) Cost
Buffer (PrimState m) (Cost, Cost, Cost, Cost)
numVerticesMCFB :: forall s. MinCostFlowBuilder s -> Cost
inDegreeMCFB :: forall s. MinCostFlowBuilder s -> MVector s Cost
edgesMCFB :: forall s. MinCostFlowBuilder s -> Buffer s (Cost, Cost, Cost, Cost)
numVerticesMCFB :: Cost
inDegreeMCFB :: MVector (PrimState m) Cost
edgesMCFB :: Buffer (PrimState m) (Cost, Cost, Cost, Cost)
..} Cost
src Cost
dst Cost
cost Cost
capacity =
  Bool -> m () -> m ()
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Cost
cost Cost -> Cost -> Bool
forall a. Ord a => a -> a -> Bool
>= Cost
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    MVector (PrimState m) Cost -> (Cost -> Cost) -> Cost -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Cost -> m ()
UM.unsafeModify MVector (PrimState m) Cost
inDegreeMCFB (Cost -> Cost -> Cost
forall a. Num a => a -> a -> a
+ Cost
1) Cost
src
    MVector (PrimState m) Cost -> (Cost -> Cost) -> Cost -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Cost -> m ()
UM.unsafeModify MVector (PrimState m) Cost
inDegreeMCFB (Cost -> Cost -> Cost
forall a. Num a => a -> a -> a
+ Cost
1) Cost
dst
    (Cost, Cost, Cost, Cost)
-> Buffer (PrimState m) (Cost, Cost, Cost, Cost) -> m ()
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
a -> Buffer (PrimState m) a -> m ()
pushBack (Cost
src, Cost
dst, Cost
cost, Cost
capacity) Buffer (PrimState m) (Cost, Cost, Cost, Cost)
edgesMCFB

buildMinCostFlow ::
  (PrimMonad m) =>
  MinCostFlowBuilder (PrimState m) ->
  m (MinCostFlow (PrimState m))
buildMinCostFlow :: forall (m :: * -> *).
PrimMonad m =>
MinCostFlowBuilder (PrimState m) -> m (MinCostFlow (PrimState m))
buildMinCostFlow MinCostFlowBuilder{Cost
MVector (PrimState m) Cost
Buffer (PrimState m) (Cost, Cost, Cost, Cost)
numVerticesMCFB :: forall s. MinCostFlowBuilder s -> Cost
inDegreeMCFB :: forall s. MinCostFlowBuilder s -> MVector s Cost
edgesMCFB :: forall s. MinCostFlowBuilder s -> Buffer s (Cost, Cost, Cost, Cost)
numVerticesMCFB :: Cost
inDegreeMCFB :: MVector (PrimState m) Cost
edgesMCFB :: Buffer (PrimState m) (Cost, Cost, Cost, Cost)
..} = do
  offsetMCF <- (Cost -> Cost -> Cost) -> Cost -> Vector Cost -> Vector Cost
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> a) -> a -> Vector b -> Vector a
U.scanl' Cost -> Cost -> Cost
forall a. Num a => a -> a -> a
(+) Cost
0 (Vector Cost -> Vector Cost) -> m (Vector Cost) -> m (Vector Cost)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Cost -> m (Vector Cost)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector (PrimState m) Cost
inDegreeMCFB
  let numVerticesMCF = Cost
numVerticesMCFB
  let numEdgesMCF = Vector Cost -> Cost
forall a. Unbox a => Vector a -> a
U.last Vector Cost
offsetMCF

  moffset <- U.thaw offsetMCF
  mdstMCF <- UM.replicate numEdgesMCF nothingMCF
  mcostMCF <- UM.replicate numEdgesMCF 0
  mrevEdgeMCF <- UM.replicate numEdgesMCF nothingMCF
  residualMCF <- UM.replicate numEdgesMCF 0

  edges <- unsafeFreezeBuffer edgesMCFB
  U.forM_ edges $ \(Cost
src, Cost
dst, Cost
cost, Cost
capacity) -> do
    srcOffset <- MVector (PrimState m) Cost -> Cost -> m Cost
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Cost -> m a
UM.unsafeRead MVector (PrimState m) Cost
moffset Cost
src
    dstOffset <- UM.unsafeRead moffset dst
    UM.unsafeModify moffset (+ 1) src
    UM.unsafeModify moffset (+ 1) dst

    UM.unsafeWrite mdstMCF srcOffset dst
    UM.unsafeWrite mdstMCF dstOffset src
    UM.unsafeWrite mcostMCF srcOffset cost
    UM.unsafeWrite mcostMCF dstOffset (-cost)
    UM.unsafeWrite mrevEdgeMCF srcOffset dstOffset
    UM.unsafeWrite mrevEdgeMCF dstOffset srcOffset
    UM.unsafeWrite residualMCF srcOffset capacity

  dstMCF <- U.unsafeFreeze mdstMCF
  costMCF <- U.unsafeFreeze mcostMCF
  potentialMCF <- UM.replicate numVerticesMCF 0
  distMCF <- UM.replicate numVerticesMCF 0
  heapMCF <- newMinBinaryHeap (numEdgesMCF + 1)
  revEdgeMCF <- U.unsafeFreeze mrevEdgeMCF
  prevVertexMCF <- UM.replicate numVerticesMCF nothingMCF
  prevEdgeMCF <- UM.replicate numVerticesMCF nothingMCF
  return MinCostFlow{..}