{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}

module Constrained.Graph where

import Control.Monad
import Data.Foldable
import Data.List (sortOn)
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe
import Data.Set (Set)
import Data.Set qualified as Set
import Prettyprinter

data Graph node = Graph
  { forall node. Graph node -> Map node (Set node)
edges :: !(Map node (Set node))
  , forall node. Graph node -> Map node (Set node)
opEdges :: !(Map node (Set node))
  }
  deriving (Int -> Graph node -> ShowS
[Graph node] -> ShowS
Graph node -> String
(Int -> Graph node -> ShowS)
-> (Graph node -> String)
-> ([Graph node] -> ShowS)
-> Show (Graph node)
forall node. Show node => Int -> Graph node -> ShowS
forall node. Show node => [Graph node] -> ShowS
forall node. Show node => Graph node -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall node. Show node => Int -> Graph node -> ShowS
showsPrec :: Int -> Graph node -> ShowS
$cshow :: forall node. Show node => Graph node -> String
show :: Graph node -> String
$cshowList :: forall node. Show node => [Graph node] -> ShowS
showList :: [Graph node] -> ShowS
Show)

instance Ord node => Semigroup (Graph node) where
  Graph Map node (Set node)
e Map node (Set node)
o <> :: Graph node -> Graph node -> Graph node
<> Graph Map node (Set node)
e' Map node (Set node)
o' =
    Map node (Set node) -> Map node (Set node) -> Graph node
forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph
      ((Set node -> Set node -> Set node)
-> Map node (Set node)
-> Map node (Set node)
-> Map node (Set node)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith Set node -> Set node -> Set node
forall a. Semigroup a => a -> a -> a
(<>) Map node (Set node)
e Map node (Set node)
e')
      ((Set node -> Set node -> Set node)
-> Map node (Set node)
-> Map node (Set node)
-> Map node (Set node)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith Set node -> Set node -> Set node
forall a. Semigroup a => a -> a -> a
(<>) Map node (Set node)
o Map node (Set node)
o')

instance Ord node => Monoid (Graph node) where
  mempty :: Graph node
mempty = Map node (Set node) -> Map node (Set node) -> Graph node
forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph Map node (Set node)
forall a. Monoid a => a
mempty Map node (Set node)
forall a. Monoid a => a
mempty

instance Pretty n => Pretty (Graph n) where
  pretty :: forall ann. Graph n -> Doc ann
pretty Graph n
gr =
    [Doc ann] -> Doc ann
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$
      Doc ann -> [Doc ann] -> [Doc ann]
forall ann. Doc ann -> [Doc ann] -> [Doc ann]
punctuate
        Doc ann
forall ann. Doc ann
hardline
        [ Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
nest Int
4 (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ n -> Doc ann
forall ann. n -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty n
n Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
" <- " Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
brackets ([Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
fillSep ((n -> Doc ann) -> [n] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map n -> Doc ann
forall ann. n -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Set n -> [n]
forall a. Set a -> [a]
Set.toList Set n
ns)))
        | (n
n, Set n
ns) <- Map n (Set n) -> [(n, Set n)]
forall k a. Map k a -> [(k, a)]
Map.toList (Graph n -> Map n (Set n)
forall node. Graph node -> Map node (Set node)
edges Graph n
gr)
        ]

nodes :: Graph node -> Set node
nodes :: forall node. Graph node -> Set node
nodes (Graph Map node (Set node)
e Map node (Set node)
_) = Map node (Set node) -> Set node
forall k a. Map k a -> Set k
Map.keysSet Map node (Set node)
e

opGraph :: Graph node -> Graph node
opGraph :: forall node. Graph node -> Graph node
opGraph (Graph Map node (Set node)
e Map node (Set node)
o) = Map node (Set node) -> Map node (Set node) -> Graph node
forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph Map node (Set node)
o Map node (Set node)
e

subtractGraph :: Ord node => Graph node -> Graph node -> Graph node
subtractGraph :: forall node. Ord node => Graph node -> Graph node -> Graph node
subtractGraph (Graph Map node (Set node)
e Map node (Set node)
o) (Graph Map node (Set node)
e' Map node (Set node)
o') =
  Map node (Set node) -> Map node (Set node) -> Graph node
forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph
    ((Set node -> Set node -> Maybe (Set node))
-> Map node (Set node)
-> Map node (Set node)
-> Map node (Set node)
forall k a b.
Ord k =>
(a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
Map.differenceWith Set node -> Set node -> Maybe (Set node)
forall {a}. Ord a => Set a -> Set a -> Maybe (Set a)
del Map node (Set node)
e Map node (Set node)
e')
    ((Set node -> Set node -> Maybe (Set node))
-> Map node (Set node)
-> Map node (Set node)
-> Map node (Set node)
forall k a b.
Ord k =>
(a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
Map.differenceWith Set node -> Set node -> Maybe (Set node)
forall {a}. Ord a => Set a -> Set a -> Maybe (Set a)
del Map node (Set node)
o Map node (Set node)
o')
  where
    del :: Set a -> Set a -> Maybe (Set a)
del Set a
x Set a
y = Set a -> Maybe (Set a)
forall a. a -> Maybe a
Just (Set a -> Maybe (Set a)) -> Set a -> Maybe (Set a)
forall a b. (a -> b) -> a -> b
$ Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
Set.difference Set a
x Set a
y

dependency :: Ord node => node -> Set node -> Graph node
dependency :: forall node. Ord node => node -> Set node -> Graph node
dependency node
x Set node
xs =
  Map node (Set node) -> Map node (Set node) -> Graph node
forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph
    (node -> Set node -> Map node (Set node)
forall k a. k -> a -> Map k a
Map.singleton node
x Set node
xs)
    ( (Set node -> Set node -> Set node)
-> Map node (Set node)
-> Map node (Set node)
-> Map node (Set node)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith
        Set node -> Set node -> Set node
forall a. Semigroup a => a -> a -> a
(<>)
        (node -> Set node -> Map node (Set node)
forall k a. k -> a -> Map k a
Map.singleton node
x Set node
forall a. Monoid a => a
mempty)
        ([(node, Set node)] -> Map node (Set node)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(node
y, node -> Set node
forall a. a -> Set a
Set.singleton node
x) | node
y <- Set node -> [node]
forall a. Set a -> [a]
Set.toList Set node
xs])
    )

irreflexiveDependencyOn :: Ord node => Set node -> Set node -> Graph node
irreflexiveDependencyOn :: forall node. Ord node => Set node -> Set node -> Graph node
irreflexiveDependencyOn Set node
xs Set node
ys =
  Graph node
deps Graph node -> Graph node -> Graph node
forall a. Semigroup a => a -> a -> a
<> Set node -> Graph node
forall node. Ord node => Set node -> Graph node
noDependencies Set node
ys
  where
    deps :: Graph node
deps =
      Map node (Set node) -> Map node (Set node) -> Graph node
forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph
        ([(node, Set node)] -> Map node (Set node)
forall k a. [(k, a)] -> Map k a
Map.fromDistinctAscList [(node
x, node -> Set node -> Set node
forall a. Ord a => a -> Set a -> Set a
Set.delete node
x Set node
ys) | node
x <- Set node -> [node]
forall a. Set a -> [a]
Set.toList Set node
xs])
        ([(node, Set node)] -> Map node (Set node)
forall k a. [(k, a)] -> Map k a
Map.fromDistinctAscList [(node
a, node -> Set node -> Set node
forall a. Ord a => a -> Set a -> Set a
Set.delete node
a Set node
xs) | node
a <- Set node -> [node]
forall a. Set a -> [a]
Set.toList Set node
ys])

transitiveDependencies :: Ord node => node -> Graph node -> Set node
transitiveDependencies :: forall node. Ord node => node -> Graph node -> Set node
transitiveDependencies node
x (Graph Map node (Set node)
e Map node (Set node)
_) = Set node -> node -> Set node
go (node -> Set node
forall a. a -> Set a
Set.singleton node
x) node
x
  where
    go :: Set node -> node -> Set node
go Set node
seen node
y = Set node
ys Set node -> Set node -> Set node
forall a. Semigroup a => a -> a -> a
<> (node -> Set node) -> Set node -> Set node
forall m a. Monoid m => (a -> m) -> Set a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Set node -> node -> Set node
go (Set node -> node -> Set node) -> Set node -> node -> Set node
forall a b. (a -> b) -> a -> b
$ node -> Set node -> Set node
forall a. Ord a => a -> Set a -> Set a
Set.insert node
y Set node
seen) (Set node -> Set node -> Set node
forall a. Ord a => Set a -> Set a -> Set a
Set.difference Set node
ys Set node
seen)
      where
        ys :: Set node
ys = Set node -> Maybe (Set node) -> Set node
forall a. a -> Maybe a -> a
fromMaybe Set node
forall a. Monoid a => a
mempty (node -> Map node (Set node) -> Maybe (Set node)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup node
y Map node (Set node)
e)

transitiveClosure :: Ord node => Graph node -> Graph node
transitiveClosure :: forall node. Ord node => Graph node -> Graph node
transitiveClosure Graph node
g = (node -> Graph node) -> Set node -> Graph node
forall m a. Monoid m => (a -> m) -> Set a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\node
x -> node -> Set node -> Graph node
forall node. Ord node => node -> Set node -> Graph node
dependency node
x (node -> Graph node -> Set node
forall node. Ord node => node -> Graph node -> Set node
transitiveDependencies node
x Graph node
g)) (Graph node -> Set node
forall node. Graph node -> Set node
nodes Graph node
g)

noDependencies :: Ord node => Set node -> Graph node
noDependencies :: forall node. Ord node => Set node -> Graph node
noDependencies Set node
ns = Map node (Set node) -> Map node (Set node) -> Graph node
forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph Map node (Set node)
nodeMap Map node (Set node)
nodeMap
  where
    nodeMap :: Map node (Set node)
nodeMap = [(node, Set node)] -> Map node (Set node)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ((,Set node
forall a. Monoid a => a
mempty) (node -> (node, Set node)) -> [node] -> [(node, Set node)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set node -> [node]
forall a. Set a -> [a]
Set.toList Set node
ns)

-- | Topsort the graph, returning a cycle
-- (`Left cycle`) on failure.
topsort :: Ord node => Graph node -> Either [node] [node]
topsort :: forall node. Ord node => Graph node -> Either [node] [node]
topsort gr :: Graph node
gr@(Graph Map node (Set node)
e Map node (Set node)
_) = [node] -> Map node (Set node) -> Either [node] [node]
forall {k}. Ord k => [k] -> Map k (Set k) -> Either [node] [k]
go [] Map node (Set node)
e
  where
    go :: [k] -> Map k (Set k) -> Either [node] [k]
go [k]
order Map k (Set k)
g
      | Map k (Set k) -> Bool
forall a. Map k a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Map k (Set k)
g = [k] -> Either [node] [k]
forall a. a -> Either [node] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([k] -> Either [node] [k]) -> [k] -> Either [node] [k]
forall a b. (a -> b) -> a -> b
$ [k] -> [k]
forall a. [a] -> [a]
reverse [k]
order
      | Bool
otherwise = do
          let noDeps :: Set k
noDeps = Map k (Set k) -> Set k
forall k a. Map k a -> Set k
Map.keysSet (Map k (Set k) -> Set k)
-> (Map k (Set k) -> Map k (Set k)) -> Map k (Set k) -> Set k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set k -> Bool) -> Map k (Set k) -> Map k (Set k)
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter Set k -> Bool
forall a. Set a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Map k (Set k) -> Set k) -> Map k (Set k) -> Set k
forall a b. (a -> b) -> a -> b
$ Map k (Set k)
g
              removeNode :: k -> Set k -> f (Set k)
removeNode k
n Set k
ds = Set k -> Set k -> Set k
forall a. Ord a => Set a -> Set a -> Set a
Set.difference Set k
ds Set k
noDeps Set k -> f () -> f (Set k)
forall a b. a -> f b -> f a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool -> f ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ k
n k -> Set k -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set k
noDeps)
          if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Set k -> Bool
forall a. Set a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Set k
noDeps
            then [k] -> Map k (Set k) -> Either [node] [k]
go (Set k -> [k]
forall a. Set a -> [a]
Set.toList Set k
noDeps [k] -> [k] -> [k]
forall a. [a] -> [a] -> [a]
++ [k]
order) ((k -> Set k -> Maybe (Set k)) -> Map k (Set k) -> Map k (Set k)
forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybeWithKey k -> Set k -> Maybe (Set k)
forall {f :: * -> *}. Alternative f => k -> Set k -> f (Set k)
removeNode Map k (Set k)
g)
            else [node] -> Either [node] [k]
forall a b. a -> Either a b
Left ([node] -> Either [node] [k])
-> ([node] -> [node]) -> [node] -> Either [node] [k]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[node]] -> [node]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[node]] -> [node]) -> ([node] -> [[node]]) -> [node] -> [node]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [[node]] -> [[node]]
forall a. Int -> [a] -> [a]
take Int
1 ([[node]] -> [[node]])
-> ([node] -> [[node]]) -> [node] -> [[node]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([node] -> Int) -> [[node]] -> [[node]]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn [node] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([[node]] -> [[node]])
-> ([node] -> [[node]]) -> [node] -> [[node]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([node] -> Bool) -> [[node]] -> [[node]]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> ([node] -> Bool) -> [node] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [node] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) ([[node]] -> [[node]])
-> ([node] -> [[node]]) -> [node] -> [[node]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (node -> [node]) -> [node] -> [[node]]
forall a b. (a -> b) -> [a] -> [b]
map (Graph node -> node -> [node]
forall node. Ord node => Graph node -> node -> [node]
findCycle Graph node
gr) ([node] -> Either [node] [k]) -> [node] -> Either [node] [k]
forall a b. (a -> b) -> a -> b
$ Map node (Set node) -> [node]
forall k a. Map k a -> [k]
Map.keys Map node (Set node)
e

-- | Simple DFS cycle finding
-- TODO: tests for this, currently it can produce a stem with a cycle after it
findCycle :: Ord node => Graph node -> node -> [node]
findCycle :: forall node. Ord node => Graph node -> node -> [node]
findCycle (Graph Map node (Set node)
e Map node (Set node)
_) node
node = [[node]] -> [node]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[node]] -> [node])
-> ([[node]] -> [[node]]) -> [[node]] -> [node]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [[node]] -> [[node]]
forall a. Int -> [a] -> [a]
take Int
1 ([[node]] -> [node]) -> [[node]] -> [node]
forall a b. (a -> b) -> a -> b
$ Set node -> node -> [[node]]
go Set node
forall a. Monoid a => a
mempty node
node
  where
    go :: Set node -> node -> [[node]]
go Set node
seen node
n
      | node
n node -> Set node -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set node
seen = [[]]
      | Bool
otherwise = do
          node
n' <- [node]
neighbours
          (node
n node -> [node] -> [node]
forall a. a -> [a] -> [a]
:) ([node] -> [node]) -> [[node]] -> [[node]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set node -> node -> [[node]]
go (node -> Set node -> Set node
forall a. Ord a => a -> Set a -> Set a
Set.insert node
n Set node
seen) node
n'
      where
        neighbours :: [node]
neighbours = [node] -> (Set node -> [node]) -> Maybe (Set node) -> [node]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] Set node -> [node]
forall a. Set a -> [a]
Set.toList (Maybe (Set node) -> [node]) -> Maybe (Set node) -> [node]
forall a b. (a -> b) -> a -> b
$ node -> Map node (Set node) -> Maybe (Set node)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup node
n Map node (Set node)
e

dependencies :: Ord node => node -> Graph node -> Set node
dependencies :: forall node. Ord node => node -> Graph node -> Set node
dependencies node
x (Graph Map node (Set node)
e Map node (Set node)
_) = Set node -> Maybe (Set node) -> Set node
forall a. a -> Maybe a -> a
fromMaybe Set node
forall a. Monoid a => a
mempty (node -> Map node (Set node) -> Maybe (Set node)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup node
x Map node (Set node)
e)

dependsOn :: Ord node => node -> node -> Graph node -> Bool
dependsOn :: forall node. Ord node => node -> node -> Graph node -> Bool
dependsOn node
x node
y Graph node
g = node
y node -> Set node -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` node -> Graph node -> Set node
forall node. Ord node => node -> Graph node -> Set node
dependencies node
x Graph node
g

deleteNode :: Ord node => node -> Graph node -> Graph node
deleteNode :: forall node. Ord node => node -> Graph node -> Graph node
deleteNode node
x (Graph Map node (Set node)
e Map node (Set node)
o) =
  Map node (Set node) -> Map node (Set node) -> Graph node
forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph
    (node -> Map node (Set node) -> Map node (Set node)
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete node
x (Map node (Set node) -> Map node (Set node))
-> Map node (Set node) -> Map node (Set node)
forall a b. (a -> b) -> a -> b
$ (Set node -> Set node)
-> Map node (Set node) -> Map node (Set node)
forall a b. (a -> b) -> Map node a -> Map node b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (node -> Set node -> Set node
forall a. Ord a => a -> Set a -> Set a
Set.delete node
x) Map node (Set node)
e)
    (node -> Map node (Set node) -> Map node (Set node)
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete node
x (Map node (Set node) -> Map node (Set node))
-> Map node (Set node) -> Map node (Set node)
forall a b. (a -> b) -> a -> b
$ (Set node -> Set node)
-> Map node (Set node) -> Map node (Set node)
forall a b. (a -> b) -> Map node a -> Map node b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (node -> Set node -> Set node
forall a. Ord a => a -> Set a -> Set a
Set.delete node
x) Map node (Set node)
o)