{-# 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
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
showList :: [Graph node] -> ShowS
$cshowList :: forall node. Show node => [Graph node] -> ShowS
show :: Graph node -> String
$cshow :: forall node. Show node => Graph node -> String
showsPrec :: Int -> Graph node -> ShowS
$cshowsPrec :: forall node. Show node => Int -> 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' =
    forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph
      (forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith forall a. Semigroup a => a -> a -> a
(<>) Map node (Set node)
e Map node (Set node)
e')
      (forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith 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 = forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty

instance Pretty n => Pretty (Graph n) where
  pretty :: forall ann. Graph n -> Doc ann
pretty Graph n
gr =
    forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold forall a b. (a -> b) -> a -> b
$
      forall ann. Doc ann -> [Doc ann] -> [Doc ann]
punctuate
        forall ann. Doc ann
hardline
        [ forall ann. Int -> Doc ann -> Doc ann
nest Int
4 forall a b. (a -> b) -> a -> b
$ forall a ann. Pretty a => a -> Doc ann
pretty n
n forall a. Semigroup a => a -> a -> a
<> Doc ann
" <- " forall a. Semigroup a => a -> a -> a
<> forall ann. Doc ann -> Doc ann
brackets (forall ann. [Doc ann] -> Doc ann
fillSep (forall a b. (a -> b) -> [a] -> [b]
map forall a ann. Pretty a => a -> Doc ann
pretty (forall a. Set a -> [a]
Set.toList Set n
ns)))
        | (n
n, Set n
ns) <- forall k a. Map k a -> [(k, a)]
Map.toList (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)
_) = 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) = 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') =
  forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph
    (forall k a b.
Ord k =>
(a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
Map.differenceWith forall {a}. Ord a => Set a -> Set a -> Maybe (Set a)
del Map node (Set node)
e Map node (Set node)
e')
    (forall k a b.
Ord k =>
(a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
Map.differenceWith 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 = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ 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 =
  forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph
    (forall k a. k -> a -> Map k a
Map.singleton node
x Set node
xs)
    ( forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith
        forall a. Semigroup a => a -> a -> a
(<>)
        (forall k a. k -> a -> Map k a
Map.singleton node
x forall a. Monoid a => a
mempty)
        (forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(node
y, forall a. a -> Set a
Set.singleton node
x) | node
y <- 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 forall a. Semigroup a => a -> a -> a
<> forall node. Ord node => Set node -> Graph node
noDependencies Set node
ys
  where
    deps :: Graph node
deps =
      forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph
        (forall k a. [(k, a)] -> Map k a
Map.fromDistinctAscList [(node
x, forall a. Ord a => a -> Set a -> Set a
Set.delete node
x Set node
ys) | node
x <- forall a. Set a -> [a]
Set.toList Set node
xs])
        (forall k a. [(k, a)] -> Map k a
Map.fromDistinctAscList [(node
a, forall a. Ord a => a -> Set a -> Set a
Set.delete node
a Set node
xs) | node
a <- 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 (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 forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Set node -> node -> Set node
go forall a b. (a -> b) -> a -> b
$ forall a. Ord a => a -> Set a -> Set a
Set.insert node
y Set node
seen) (forall a. Ord a => Set a -> Set a -> Set a
Set.difference Set node
ys Set node
seen)
      where
        ys :: Set node
ys = forall a. a -> Maybe a -> a
fromMaybe forall a. Monoid a => a
mempty (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 = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\node
x -> forall node. Ord node => node -> Set node -> Graph node
dependency node
x (forall node. Ord node => node -> Graph node -> Set node
transitiveDependencies node
x Graph node
g)) (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 = 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 = forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ((,forall a. Monoid a => a
mempty) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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)
_) = 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
      | forall (t :: * -> *) a. Foldable t => t a -> Bool
null Map k (Set k)
g = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [a]
reverse [k]
order
      | Bool
otherwise = do
          let noDeps :: Set k
noDeps = forall k a. Map k a -> Set k
Map.keysSet forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter forall (t :: * -> *) a. Foldable t => t a -> Bool
null 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 = forall a. Ord a => Set a -> Set a -> Set a
Set.difference Set k
ds Set k
noDeps forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ k
n forall a. Ord a => a -> Set a -> Bool
`Set.member` Set k
noDeps)
          if Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Bool
null Set k
noDeps
            then [k] -> Map k (Set k) -> Either [node] [k]
go (forall a. Set a -> [a]
Set.toList Set k
noDeps forall a. [a] -> [a] -> [a]
++ [k]
order) (forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybeWithKey forall {f :: * -> *}. Alternative f => k -> Set k -> f (Set k)
removeNode Map k (Set k)
g)
            else forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> [a] -> [a]
take Int
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall (t :: * -> *) a. Foldable t => t a -> Int
length forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Bool
null) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall node. Ord node => Graph node -> node -> [node]
findCycle Graph node
gr) forall a b. (a -> b) -> a -> b
$ 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 = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> [a] -> [a]
take Int
1 forall a b. (a -> b) -> a -> b
$ Set node -> node -> [[node]]
go forall a. Monoid a => a
mempty node
node
  where
    go :: Set node -> node -> [[node]]
go Set node
seen node
n
      | node
n forall a. Ord a => a -> Set a -> Bool
`Set.member` Set node
seen = [[]]
      | Bool
otherwise = do
          node
n' <- [node]
neighbours
          (node
n forall a. a -> [a] -> [a]
:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set node -> node -> [[node]]
go (forall a. Ord a => a -> Set a -> Set a
Set.insert node
n Set node
seen) node
n'
      where
        neighbours :: [node]
neighbours = forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] forall a. Set a -> [a]
Set.toList forall a b. (a -> b) -> a -> b
$ 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)
_) = forall a. a -> Maybe a -> a
fromMaybe forall a. Monoid a => a
mempty (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 forall a. Ord a => a -> Set a -> Bool
`Set.member` 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) =
  forall node.
Map node (Set node) -> Map node (Set node) -> Graph node
Graph
    (forall k a. Ord k => k -> Map k a -> Map k a
Map.delete node
x forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Ord a => a -> Set a -> Set a
Set.delete node
x) Map node (Set node)
e)
    (forall k a. Ord k => k -> Map k a -> Map k a
Map.delete node
x forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Ord a => a -> Set a -> Set a
Set.delete node
x) Map node (Set node)
o)