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]
-> IntSet.IntSet
-> IntMap.IntMap Marker
-> [Int]
-> Int
-> 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]
-> IntSet.IntSet
-> IntMap.IntMap Marker
-> [Int]
-> Int
-> 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'
topSort
:: [(Int, a)]
-> [(Int, Int)]
-> Either String [a]
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
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' :: 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'
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
reverseTopSort
:: [(Int, a)]
-> [(Int, Int)]
-> Either String [a]
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)
callGraphBindings
:: BindingMap
-> Id
-> [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