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

import           Data.Tuple            (swap)
import           Data.Foldable         (foldlM)
#if MIN_VERSION_ghc(9,8,4)
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)

#if MIN_VERSION_ghc(9,8,4)
type IntMap = IntMap.Word64Map
type IntSet = IntSet.Word64Set
#endif

data Marker
  = Temporary
  | Permanent

headSafe :: [a] -> Maybe a
headSafe :: [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 -> 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 (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 [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
  (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 (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 :: [(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)
mapM (\(Unique
n, Unique
m) -> Unique -> Either String Unique
checkNode Unique
n Either String Unique
-> Either String Unique -> Either String Unique
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)
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 (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
  :: [(Unique, a)]
  -- ^ Nodes
  -> [(Unique, Unique)]
  -- ^ Edges
  -> Either String [a]
  -- ^ Reversely, topologically sorted nodes
reverseTopSort :: [(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