content: {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE LinearTypes #-}
  {-# LANGUAGE MagicHash #-}
  {-# OPTIONS_GHC -Wno-name-shadowing #-}
  {-# OPTIONS_GHC -Wno-unused-matches #-}
  
  module Simple.TopSort where
  
  import Data.Bifunctor.Linear (second)
  import qualified Data.Functor.Linear as Data
  import Data.HashMap.Mutable.Linear (HashMap)
  import qualified Data.HashMap.Mutable.Linear as HMap
  import Data.Maybe.Linear (catMaybes)
  import Data.Unrestricted.Linear
  import qualified Prelude.Linear as Linear
  
  -- # The topological sort of a DAG
  -------------------------------------------------------------------------------
  
  type Node = Int
  
  type InDegGraph = HashMap Node ([Node], Int)
  
  topsort :: [(Node, [Node])] -> [Node]
  topsort = reverse . postOrder . fmap (\(n, nbrs) -> (n, (nbrs, 0)))
    where
      postOrder :: [(Node, ([Node], Int))] -> [Node]
      postOrder [] = []
      postOrder (xs) =
        let nodes = map fst xs
         in unur Linear.$
              HMap.empty (length xs * 2) Linear.$
                \hm -> postOrderHM nodes (HMap.insertAll xs hm)
  
  postOrderHM :: [Node] -> InDegGraph %1 -> Ur [Node]
  postOrderHM nodes dag =
    case findSources nodes (computeInDeg nodes dag) of
      (dag, Ur sources) -> pluckSources sources [] dag
    where
      -- O(V + N)
      computeInDeg :: [Node] -> InDegGraph %1 -> InDegGraph
      computeInDeg nodes dag = Linear.foldl incChildren dag (map Ur nodes)
  
      -- Increment in-degree of all neighbors
      incChildren :: InDegGraph %1 -> Ur Node %1 -> InDegGraph
      incChildren dag (Ur node) =
        case HMap.lookup node dag of
          (Ur Nothing, dag) -> dag
          (Ur (Just (xs, i)), dag) -> incNodes (move xs) dag
        where
          incNodes :: Ur [Node] %1 -> InDegGraph %1 -> InDegGraph
          incNodes (Ur ns) dag = Linear.foldl incNode dag (map Ur ns)
  
          incNode :: InDegGraph %1 -> Ur Node %1 -> InDegGraph
          incNode dag (Ur node) =
            case HMap.lookup node dag of
              (Ur Nothing, dag') -> dag'
              (Ur (Just (n, d)), dag') ->
                HMap.insert node (n, d + 1) dag'
  
  -- HMap.alter dag (\(Just (n,d)) -> Just (n,d+1)) node
  
  -- pluckSources sources postOrdSoFar dag
  pluckSources :: [Node] -> [Node] -> InDegGraph %1 -> Ur [Node]
  pluckSources [] postOrd dag = lseq dag (move postOrd)
  pluckSources (s : ss) postOrd dag =
    case HMap.lookup s dag of
      (Ur Nothing, dag) -> pluckSources ss (s : postOrd) dag
      (Ur (Just (xs, i)), dag) ->
        case walk xs dag of
          (dag', Ur newSrcs) ->
            pluckSources (newSrcs ++ ss) (s : postOrd) dag'
    where
      -- decrement degree of children, save newly made sources
      walk :: [Node] -> InDegGraph %1 -> (InDegGraph, Ur [Node])
      walk children dag =
        second (Data.fmap catMaybes) (mapAccum decDegree children dag)
  
      -- Decrement the degree of a node, save it if it is now a source
      decDegree :: Node -> InDegGraph %1 -> (InDegGraph, Ur (Maybe Node))
      decDegree node dag =
        case HMap.lookup node dag of
          (Ur Nothing, dag') -> (dag', Ur Nothing)
          (Ur (Just (n, d)), dag') ->
            checkSource node (HMap.insert node (n, d - 1) dag')
  
  -- Given a list of nodes, determines which are sources
  findSources :: [Node] -> InDegGraph %1 -> (InDegGraph, Ur [Node])
  findSources nodes dag =
    second (Data.fmap catMaybes) (mapAccum checkSource nodes dag)
  
  -- | Check if a node is a source, and if so return it
  checkSource :: Node -> InDegGraph %1 -> (InDegGraph, Ur (Maybe Node))
  checkSource node dag =
    case HMap.lookup node dag of
      (Ur Nothing, dag) -> (dag, Ur Nothing)
      (Ur (Just (xs, 0)), dag) -> (dag, Ur (Just node))
      (Ur (Just (xs, n)), dag) -> (dag, Ur Nothing)
  
  mapAccum ::
    (a -> b %1 -> (b, Ur c)) -> [a] -> b %1 -> (b, Ur [c])
  mapAccum f [] b = (b, Ur [])
  mapAccum f (x : xs) b =
    case mapAccum f xs b of
      (b, Ur cs) -> second (Data.fmap (: cs)) (f x b)
  