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

      Collection of utilities
    -}

    {-# LANGUAGE CPP #-}

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

    #if MIN_VERSION_ghc(9,8,4) || (MIN_VERSION_ghc(9,6,7) && !MIN_VERSION_ghc(9,8,0))
    #define UNIQUE_IS_WORD64
    #endif

    import           Data.Tuple            (swap)
    import           Data.Foldable         (foldlM)
    #ifdef UNIQUE_IS_WORD64
    import qualified GHC.Data.Word64Map.Strict as IntMap
    import qualified GHC.Data.Word64Set        as IntSet
    #else
    import           Data.IntMap.Strict    (IntMap)
    import qualified Data.IntMap.Strict    as IntMap
    import           Data.IntSet           (IntSet)
    import qualified Data.IntSet           as IntSet
    #endif

    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)
    import           Clash.Unique (Unique)

    #ifdef UNIQUE_IS_WORD64
    type IntMap = IntMap.Word64Map
    type IntSet = IntSet.Word64Set
    #endif

    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 [Unique]
      -- ^ Edges
      -> IntSet
      -- ^ Unmarked nodes
      -> IntMap Marker
      -- ^ Marked nodes
      -> [Unique]
      -- ^ Sorted so far
      -> Unique
      -- ^ Node to visit
      -> Either String (IntSet, IntMap Marker, [Unique])
    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 -> Word64Map 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 -> Word64Map a -> Word64Map 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 -> Word64Map a -> a
    IntMap.findWithDefault [] Unique
    node IntMap [Unique]
    edges
          (unmarked'', marked'', 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''' = Unique -> Marker -> IntMap Marker -> IntMap Marker
    forall a. Unique -> a -> Word64Map a -> Word64Map a
    IntMap.insert Unique
    node Marker
    Permanent IntMap Marker
    marked''
          return (unmarked'', marked''', node : 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 [Unique]
      -- ^ Edges
      -> IntSet
      -- ^ Unmarked nodes
      -> IntMap Marker
      -- ^ Marked nodes
      -> [Unique]
      -- ^ Sorted so far
      -> Unique
      -- ^ Node to visit
      -> Either String (IntSet, IntMap Marker, [Unique])
    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
      (unmarked', marked', 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 headSafe (IntSet.toList 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
      :: [(Unique, a)]
      -- ^ Nodes
      -> [(Unique, Unique)]
      -- ^ 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) -> 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

      (_, _, sorted) <-
        topSortVisit edges' (IntMap.keysSet nodes') IntMap.empty [] (fst node)

      mapM lookup' sorted
        where
          nodes' :: Word64Map a
    nodes' = [(Unique, a)] -> Word64Map a
    forall a. [(Unique, a)] -> Word64Map 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}. Word64Map [a] -> (Unique, a) -> Word64Map [a]
    insert IntMap [Unique]
    forall a. Word64Map a
    IntMap.empty [(Unique, Unique)]
    edges

          -- Construction functions for quick lookup of edges from n to m, given n
          insert :: Word64Map [a] -> (Unique, a) -> Word64Map [a]
    insert Word64Map [a]
    im (Unique
    n, a
    m)    = (Maybe [a] -> Maybe [a])
    -> Unique -> Word64Map [a] -> Word64Map [a]
    forall a.
    (Maybe a -> Maybe a) -> Unique -> Word64Map a -> Word64Map a
    IntMap.alter (a -> Maybe [a] -> Maybe [a]
    forall {a}. a -> Maybe [a] -> Maybe [a]
    insert' a
    m) Unique
    n Word64Map [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 -> Word64Map a -> Maybe a
    forall a. Unique -> Word64Map a -> Maybe a
    IntMap.lookup Unique
    n Word64Map 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 -> Word64Map a -> Bool
    forall a. Unique -> Word64Map a -> Bool
    IntMap.notMember Unique
    n Word64Map 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
      :: [(Unique, a)]
      -- ^ Nodes
      -> [(Unique, Unique)]
      -- ^ 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