content: module Jikka.Common.Graph where
  
  import Control.Monad
  import Control.Monad.ST
  import Data.List (nub)
  import Data.STRef
  import qualified Data.Vector as V
  import qualified Data.Vector.Mutable as MV
  
  type Graph = V.Vector [Int]
  
  makeReversedDigraph :: Graph -> Graph
  makeReversedDigraph g = runST $ do
    let n = V.length g
    h <- MV.replicate n []
    forM_ [0 .. n - 1] $ \x -> do
      forM_ (g V.! x) $ \y -> do
        MV.modify h (x :) y
    V.freeze h
  
  makeInducedDigraph :: Graph -> V.Vector Int -> Graph
  makeInducedDigraph g f = runST $ do
    let n = V.length g
    let k = if V.null f then 0 else V.maximum f + 1
    h <- MV.replicate k []
    forM_ [0 .. n - 1] $ \x -> do
      forM_ (g V.! x) $ \y -> do
        MV.modify h ((f V.! y) :) (f V.! x)
    forM_ [0 .. k - 1] $ \a -> do
      MV.modify h nub a
    V.freeze h
  
  -- | `decomposeToStronglyConnectedComponents` does SCC in \(O(V + E)\) using Kosaraju's algorithm.
  -- It takes a digraph \(G = (V, E)\) as an adjacent list \(g : V \to V^{\lt \omega}\), and returns an mapping \(f : V \to V'\) for the SCC DAG \(G' = (V', E')\).
  -- The indices of vertices of the SCC DAG are topologically sorted.
  decomposeToStronglyConnectedComponents :: Graph -> V.Vector Int
  decomposeToStronglyConnectedComponents g = runST $ do
    let n = V.length g
    let unless' used x f = do
          usedX <- MV.read used x
          unless usedX $ do
            f
    -- The first DFS
    let order = topologicalSort g
    -- DFS on the reversed graph
    let gRev = makeReversedDigraph g
    componentOf <- MV.replicate n (-1)
    size <- newSTRef 0
    used <- MV.replicate n False
    let go x = do
          MV.write used x True
          forM_ (gRev V.! x) $ \y -> do
            unless' used y $ do
              go y
    V.forM_ order $ \x -> do
      unless' used x $ do
        go x
        modifySTRef' size succ
    V.freeze componentOf
  
  -- | `topologicalSort` does topological sort in \(O(V + E)\) using Tarjan's algorithm.
  -- The input is an adjacent list of a DAG.
  topologicalSort :: Graph -> V.Vector Int
  topologicalSort g = runST $ do
    let n = V.length g
    let unless' used x f = do
          usedX <- MV.read used x
          unless usedX $ do
            f
    order <- newSTRef []
    used <- MV.replicate n False
    let go x = do
          MV.write used x True
          forM_ (g V.! x) $ \y -> do
            unless' used y $ do
              go y
          modifySTRef' order (x :)
    forM_ [0 .. n - 1] $ \x -> do
      unless' used x $ do
        go x
    V.fromList <$> readSTRef order
  