    {-|
      Copyright   :  (C) 2018, QBayLogic
      License     :  BSD2 (see the file LICENSE)
      Maintainer  :  Christiaan Baaij <christiaan.baaij@gmail.com>

      Collection of utilities
    -}

    module Clash.Util.Graph
      ( topSort
      , reverseTopSort
      , callGraphBindings
      ) where

    import           Data.Tuple            (swap)
    import           Data.Foldable         (foldlM)
    import qualified Data.IntMap.Strict    as IntMap
    import qualified Data.IntSet           as IntSet

    import           Clash.Core.Var (Id)
    import           Clash.Core.Term (Term)
    import qualified Clash.Data.UniqMap as UniqMap
    import           Clash.Driver.Types (BindingMap, Binding (bindingTerm))
    import           Clash.Normalize.Util (callGraph)

    data Marker
      = Temporary
      | Permanent

    headSafe :: [a] -> Maybe a
    headSafe :: forall a. [a] -> Maybe a
    headSafe [] = Maybe a
    forall a. Maybe a
    Nothing
    headSafe (a
    a:[a]
    _) = a -> Maybe a
    forall a. a -> Maybe a
    Just a
    a

    topSortVisit'
      :: IntMap.IntMap [Int]
      -- ^ Edges
      -> IntSet.IntSet
      -- ^ Unmarked nodes
      -> IntMap.IntMap Marker
      -- ^ Marked nodes
      -> [Int]
      -- ^ Sorted so far
      -> Int
      -- ^ Node to visit
      -> Either String (IntSet.IntSet, IntMap.IntMap Marker, [Int])
    topSortVisit' :: IntMap [Unique]
    -> IntSet
    -> IntMap Marker
    -> [Unique]
    -> Unique
    -> Either String (IntSet, IntMap Marker, [Unique])
    topSortVisit' IntMap [Unique]
    edges IntSet
    unmarked IntMap Marker
    marked [Unique]
    sorted Unique
    node =
      case Unique -> IntMap Marker -> Maybe Marker
    forall a. Unique -> IntMap a -> Maybe a
    IntMap.lookup Unique
    node IntMap Marker
    marked of
        Just Marker
    Permanent -> (IntSet, IntMap Marker, [Unique])
    -> Either String (IntSet, IntMap Marker, [Unique])
    forall a b. b -> Either a b
    Right (IntSet
    unmarked, IntMap Marker
    marked, [Unique]
    sorted)
        Just Marker
    Temporary -> String -> Either String (IntSet, IntMap Marker, [Unique])
    forall a b. a -> Either a b
    Left String
    "cycle detected: cannot topsort cyclic graph"
        Maybe Marker
    Nothing -> do
          let marked' :: IntMap Marker
    marked'   = Unique -> Marker -> IntMap Marker -> IntMap Marker
    forall a. Unique -> a -> IntMap a -> IntMap a
    IntMap.insert Unique
    node Marker
    Temporary IntMap Marker
    marked
          let unmarked' :: IntSet
    unmarked' = Unique -> IntSet -> IntSet
    IntSet.delete Unique
    node IntSet
    unmarked
          let nodeToM :: [Unique]
    nodeToM   = [Unique] -> Unique -> IntMap [Unique] -> [Unique]
    forall a. a -> Unique -> IntMap a -> a
    IntMap.findWithDefault [] Unique
    node IntMap [Unique]
    edges
          (IntSet
    unmarked'', IntMap Marker
    marked'', [Unique]
    sorted'') <-
            ((IntSet, IntMap Marker, [Unique])
     -> Unique -> Either String (IntSet, IntMap Marker, [Unique]))
    -> (IntSet, IntMap Marker, [Unique])
    -> [Unique]
    -> Either String (IntSet, IntMap Marker, [Unique])
    forall (t :: Type -> Type) (m :: Type -> Type) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldlM (IntSet, IntMap Marker, [Unique])
    -> Unique -> Either String (IntSet, IntMap Marker, [Unique])
    visit (IntSet
    unmarked', IntMap Marker
    marked', [Unique]
    sorted) [Unique]
    nodeToM
          let marked''' :: IntMap Marker
    marked''' = Unique -> Marker -> IntMap Marker -> IntMap Marker
    forall a. Unique -> a -> IntMap a -> IntMap a
    IntMap.insert Unique
    node Marker
    Permanent IntMap Marker
    marked''
          (IntSet, IntMap Marker, [Unique])
    -> Either String (IntSet, IntMap Marker, [Unique])
    forall a. a -> Either String a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return (IntSet
    unmarked'', IntMap Marker
    marked''', Unique
    node Unique -> [Unique] -> [Unique]
    forall a. a -> [a] -> [a]
    : [Unique]
    sorted'')
      where
        visit :: (IntSet, IntMap Marker, [Unique])
    -> Unique -> Either String (IntSet, IntMap Marker, [Unique])
    visit (IntSet
    unmarked', IntMap Marker
    marked', [Unique]
    sorted') Unique
    node' =
          IntMap [Unique]
    -> IntSet
    -> IntMap Marker
    -> [Unique]
    -> Unique
    -> Either String (IntSet, IntMap Marker, [Unique])
    topSortVisit' IntMap [Unique]
    edges IntSet
    unmarked' IntMap Marker
    marked' [Unique]
    sorted' Unique
    node'

    topSortVisit
      :: IntMap.IntMap [Int]
      -- ^ Edges
      -> IntSet.IntSet
      -- ^ Unmarked nodes
      -> IntMap.IntMap Marker
      -- ^ Marked nodes
      -> [Int]
      -- ^ Sorted so far
      -> Int
      -- ^ Node to visit
      -> Either String (IntSet.IntSet, IntMap.IntMap Marker, [Int])
    topSortVisit :: IntMap [Unique]
    -> IntSet
    -> IntMap Marker
    -> [Unique]
    -> Unique
    -> Either String (IntSet, IntMap Marker, [Unique])
    topSortVisit IntMap [Unique]
    edges IntSet
    unmarked IntMap Marker
    marked [Unique]
    sorted Unique
    node = do
      (IntSet
    unmarked', IntMap Marker
    marked', [Unique]
    sorted') <-
        IntMap [Unique]
    -> IntSet
    -> IntMap Marker
    -> [Unique]
    -> Unique
    -> Either String (IntSet, IntMap Marker, [Unique])
    topSortVisit' IntMap [Unique]
    edges IntSet
    unmarked IntMap Marker
    marked [Unique]
    sorted Unique
    node

      case [Unique] -> Maybe Unique
    forall a. [a] -> Maybe a
    headSafe (IntSet -> [Unique]
    IntSet.toList IntSet
    unmarked') of
        Maybe Unique
    Nothing    -> (IntSet, IntMap Marker, [Unique])
    -> Either String (IntSet, IntMap Marker, [Unique])
    forall a. a -> Either String a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return (IntSet
    unmarked', IntMap Marker
    marked', [Unique]
    sorted')
        Just Unique
    node' -> IntMap [Unique]
    -> IntSet
    -> IntMap Marker
    -> [Unique]
    -> Unique
    -> Either String (IntSet, IntMap Marker, [Unique])
    topSortVisit IntMap [Unique]
    edges IntSet
    unmarked' IntMap Marker
    marked' [Unique]
    sorted' Unique
    node'

    -- | See: https://en.wikipedia.org/wiki/Topological_sorting. This function
    -- errors if edges mention nodes not mentioned in the node list or if the
    -- given graph contains cycles.
    topSort
      :: [(Int, a)]
      -- ^ Nodes
      -> [(Int, Int)]
      -- ^ Edges
      -> Either String [a]
      -- ^ Error message or topologically sorted nodes
    topSort :: forall a. [(Unique, a)] -> [(Unique, Unique)] -> Either String [a]
    topSort []             []     = [a] -> Either String [a]
    forall a b. b -> Either a b
    Right []
    topSort []             [(Unique, Unique)]
    _edges = String -> Either String [a]
    forall a b. a -> Either a b
    Left String
    "Node list was empty, but edges non-empty"
    topSort nodes :: [(Unique, a)]
    nodes@((Unique, a)
    node:[(Unique, a)]
    _)  [(Unique, Unique)]
    edges = do
      [Unique]
    _ <- ((Unique, Unique) -> Either String Unique)
    -> [(Unique, Unique)] -> Either String [Unique]
    forall (t :: Type -> Type) (m :: Type -> Type) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: Type -> Type) a b.
    Monad m =>
    (a -> m b) -> [a] -> m [b]
    mapM (\(Unique
    n, Unique
    m) -> Unique -> Either String Unique
    checkNode Unique
    n Either String Unique
    -> Either String Unique -> Either String Unique
    forall a b. Either String a -> Either String b -> Either String b
    forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
    >> Unique -> Either String Unique
    checkNode Unique
    m) [(Unique, Unique)]
    edges

      (IntSet
    _, IntMap Marker
    _, [Unique]
    sorted) <-
        IntMap [Unique]
    -> IntSet
    -> IntMap Marker
    -> [Unique]
    -> Unique
    -> Either String (IntSet, IntMap Marker, [Unique])
    topSortVisit IntMap [Unique]
    edges' (IntMap a -> IntSet
    forall a. IntMap a -> IntSet
    IntMap.keysSet IntMap a
    nodes') IntMap Marker
    forall a. IntMap a
    IntMap.empty [] ((Unique, a) -> Unique
    forall a b. (a, b) -> a
    fst (Unique, a)
    node)

      (Unique -> Either String a) -> [Unique] -> Either String [a]
    forall (t :: Type -> Type) (m :: Type -> Type) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: Type -> Type) a b.
    Monad m =>
    (a -> m b) -> [a] -> m [b]
    mapM Unique -> Either String a
    lookup' [Unique]
    sorted
        where
          nodes' :: IntMap a
    nodes' = [(Unique, a)] -> IntMap a
    forall a. [(Unique, a)] -> IntMap a
    IntMap.fromList [(Unique, a)]
    nodes
          edges' :: IntMap [Unique]
    edges' = (IntMap [Unique] -> (Unique, Unique) -> IntMap [Unique])
    -> IntMap [Unique] -> [(Unique, Unique)] -> IntMap [Unique]
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: Type -> Type) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl IntMap [Unique] -> (Unique, Unique) -> IntMap [Unique]
    forall {a}. IntMap [a] -> (Unique, a) -> IntMap [a]
    insert IntMap [Unique]
    forall a. IntMap a
    IntMap.empty [(Unique, Unique)]
    edges

          -- Construction functions for quick lookup of edges from n to m, given n
          insert :: IntMap [a] -> (Unique, a) -> IntMap [a]
    insert IntMap [a]
    im (Unique
    n, a
    m)    = (Maybe [a] -> Maybe [a]) -> Unique -> IntMap [a] -> IntMap [a]
    forall a. (Maybe a -> Maybe a) -> Unique -> IntMap a -> IntMap a
    IntMap.alter (a -> Maybe [a] -> Maybe [a]
    forall {a}. a -> Maybe [a] -> Maybe [a]
    insert' a
    m) Unique
    n IntMap [a]
    im
          insert' :: a -> Maybe [a] -> Maybe [a]
    insert' a
    m Maybe [a]
    Nothing   = [a] -> Maybe [a]
    forall a. a -> Maybe a
    Just [a
    m]
          insert' a
    m (Just [a]
    ms) = [a] -> Maybe [a]
    forall a. a -> Maybe a
    Just (a
    ma -> [a] -> [a]
    forall a. a -> [a] -> [a]
    :[a]
    ms)

          -- Lookup node in nodes map. If not present, yield error
          lookup' :: Unique -> Either String a
    lookup' Unique
    n =
            case Unique -> IntMap a -> Maybe a
    forall a. Unique -> IntMap a -> Maybe a
    IntMap.lookup Unique
    n IntMap a
    nodes' of
              Maybe a
    Nothing
                -> String -> Either String a
    forall a b. a -> Either a b
    Left (String
    "Node " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Unique -> String
    forall a. Show a => a -> String
    show Unique
    n String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    " in edge list, but not in node list.")
              Just a
    n'
                -> a -> Either String a
    forall a b. b -> Either a b
    Right a
    n'

          -- Check if edge is valid (i.e., mentioned nodes are in node list)
          checkNode :: Unique -> Either String Unique
    checkNode Unique
    n
            | Unique -> IntMap a -> Bool
    forall a. Unique -> IntMap a -> Bool
    IntMap.notMember Unique
    n IntMap a
    nodes' =
                String -> Either String Unique
    forall a b. a -> Either a b
    Left (String
    "Node " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Unique -> String
    forall a. Show a => a -> String
    show Unique
    n String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    " in edge list, but not in node list.")
            | Bool
    otherwise =
                Unique -> Either String Unique
    forall a b. b -> Either a b
    Right Unique
    n

    -- | Same as `reverse (topSort nodes edges)` if alternative representations are
    -- considered the same. That is, topSort might produce multiple answers and
    -- still deliver on its promise of yielding a topologically sorted node list.
    -- Likewise, this function promises __one__ of those lists in reverse, but not
    -- necessarily the reverse of topSort itself.
    reverseTopSort
      :: [(Int, a)]
      -- ^ Nodes
      -> [(Int, Int)]
      -- ^ Edges
      -> Either String [a]
      -- ^ Reversely, topologically sorted nodes
    reverseTopSort :: forall a. [(Unique, a)] -> [(Unique, Unique)] -> Either String [a]
    reverseTopSort [(Unique, a)]
    nodes [(Unique, Unique)]
    edges =
      [(Unique, a)] -> [(Unique, Unique)] -> Either String [a]
    forall a. [(Unique, a)] -> [(Unique, Unique)] -> Either String [a]
    topSort [(Unique, a)]
    nodes (((Unique, Unique) -> (Unique, Unique))
    -> [(Unique, Unique)] -> [(Unique, Unique)]
    forall a b. (a -> b) -> [a] -> [b]
    map (Unique, Unique) -> (Unique, Unique)
    forall a b. (a, b) -> (b, a)
    swap [(Unique, Unique)]
    edges)

    -- | Get all the terms corresponding to a call graph
    callGraphBindings
      :: BindingMap
      -- ^ All bindings
      -> Id
      -- ^ Root of the call graph
      -> [Term]
    callGraphBindings :: BindingMap -> Id -> [Term]
    callGraphBindings BindingMap
    bindingsMap Id
    tm =
      (Unique -> Term) -> [Unique] -> [Term]
    forall a b. (a -> b) -> [a] -> [b]
    map (Binding Term -> Term
    forall a. Binding a -> a
    bindingTerm (Binding Term -> Term)
    -> (Unique -> Binding Term) -> Unique -> Term
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Unique -> BindingMap -> Binding Term
    forall a b. Uniquable a => a -> UniqMap b -> b
    `UniqMap.find` BindingMap
    bindingsMap)) (UniqMap (VarEnv Word) -> [Unique]
    forall b. UniqMap b -> [Unique]
    UniqMap.keys UniqMap (VarEnv Word)
    cg)
      where
        cg :: UniqMap (VarEnv Word)
    cg = BindingMap -> Id -> UniqMap (VarEnv Word)
    callGraph BindingMap
    bindingsMap Id
    tm