{-|
  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