{-# LANGUAGE OverloadedStrings #-}

-- |
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Graph utilities shared by multiple aspects of scenarios
module Swarm.Util.Graph (
  isAcyclicGraph,
  findCycle,
  failOnCyclicGraph,
) where

import Control.Monad (forM_)
import Control.Monad.ST
import Data.Array ((!))
import Data.Array.ST
import Data.Graph (SCC (..), Vertex, graphFromEdges)
import Data.IntSet (IntSet)
import Data.IntSet qualified as IS
import Data.Text (Text)
import Data.Text qualified as T
import Swarm.Util

isAcyclicGraph :: [SCC a] -> Bool
isAcyclicGraph :: forall a. [SCC a] -> Bool
isAcyclicGraph =
  (SCC a -> Bool) -> [SCC a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all SCC a -> Bool
forall {vertex}. SCC vertex -> Bool
isAcyclicVertex
 where
  isAcyclicVertex :: SCC vertex -> Bool
isAcyclicVertex = \case
    AcyclicSCC vertex
_ -> Bool
True
    SCC vertex
_ -> Bool
False

-- | Keep track of the current search path in a DFS, both as a set of
--   vertices (for fast membership testing) and as a reversed list of
--   vertices visited along the current path, in order.
--
--   Note this is different than just keeping track of which vertices
--   have been visited at all; visited vertices remain visited when
--   DFS backtracks, but the DFSPath gets shorter again.
data DFSPath = DFSPath IntSet [Vertex]

emptyDFSPath :: DFSPath
emptyDFSPath :: DFSPath
emptyDFSPath = IntSet -> [Vertex] -> DFSPath
DFSPath IntSet
IS.empty []

appendPath :: DFSPath -> Vertex -> DFSPath
appendPath :: DFSPath -> Vertex -> DFSPath
appendPath (DFSPath IntSet
s [Vertex]
p) Vertex
v = IntSet -> [Vertex] -> DFSPath
DFSPath (Vertex -> IntSet -> IntSet
IS.insert Vertex
v IntSet
s) (Vertex
v Vertex -> [Vertex] -> [Vertex]
forall a. a -> [a] -> [a]
: [Vertex]
p)

-- | Find a cycle in a directed graph (if any exist) via DFS.
--
-- >>> findCycle [("a", 0, [0])]
-- Just ["a"]
-- >>> findCycle [("a", 0, [1]), ("b", 1, [])]
-- Nothing
-- >>> findCycle [("a", 0, [1]), ("b", 1, [0])]
-- Just ["a","b"]
-- >>> findCycle [("a", 0, [1]), ("b", 1, [2]), ("c", 2, [1])]
-- Just ["b","c"]
-- >>> findCycle [("a",3,[1]), ("b",1,[0,3]), ("c",2,[1]), ("d",0,[])]
-- Just ["b","a"]
-- >>> findCycle [("a",3,[]), ("b",1,[0,3]), ("c",2,[1]), ("d",0,[])]
-- Nothing
-- >>> findCycle [("a",3,[1]), ("b",1,[0,3]), ("c",2,[1]), ("d",0,[2])]
-- Just ["d","c","b"]
findCycle :: Ord key => [(a, key, [key])] -> Maybe [a]
findCycle :: forall key a. Ord key => [(a, key, [key])] -> Maybe [a]
findCycle [(a, key, [key])]
es = (forall s. ST s (Maybe [a])) -> Maybe [a]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe [a])) -> Maybe [a])
-> (forall s. ST s (Maybe [a])) -> Maybe [a]
forall a b. (a -> b) -> a -> b
$ do
  STUArray s Vertex Bool
visited <- (Vertex, Vertex) -> Bool -> ST s (STUArray s Vertex Bool)
forall i. Ix i => (i, i) -> Bool -> ST s (STUArray s i Bool)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Vertex
0, Vertex
n Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
1) Bool
False
  (([Vertex] -> [a]) -> Maybe [Vertex] -> Maybe [a]
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([Vertex] -> [a]) -> Maybe [Vertex] -> Maybe [a])
-> ((Vertex -> a) -> [Vertex] -> [a])
-> (Vertex -> a)
-> Maybe [Vertex]
-> Maybe [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vertex -> a) -> [Vertex] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map) ((a, key, [key]) -> a
forall {a} {b} {c}. (a, b, c) -> a
fst3 ((a, key, [key]) -> a)
-> (Vertex -> (a, key, [key])) -> Vertex -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vertex -> (a, key, [key])
v2l) (Maybe [Vertex] -> Maybe [a])
-> ST s (Maybe [Vertex]) -> ST s (Maybe [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STUArray s Vertex Bool
-> DFSPath -> [Vertex] -> ST s (Maybe [Vertex])
forall s.
STUArray s Vertex Bool
-> DFSPath -> [Vertex] -> ST s (Maybe [Vertex])
dfsL STUArray s Vertex Bool
visited DFSPath
emptyDFSPath [Vertex
0 .. Vertex
n Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
1]
 where
  n :: Vertex
n = [(a, key, [key])] -> Vertex
forall a. [a] -> Vertex
forall (t :: * -> *) a. Foldable t => t a -> Vertex
length [(a, key, [key])]
es
  (Graph
g, Vertex -> (a, key, [key])
v2l, key -> Maybe Vertex
_) = [(a, key, [key])]
-> (Graph, Vertex -> (a, key, [key]), key -> Maybe Vertex)
forall key node.
Ord key =>
[(node, key, [key])]
-> (Graph, Vertex -> (node, key, [key]), key -> Maybe Vertex)
graphFromEdges [(a, key, [key])]
es
  fst3 :: (a, b, c) -> a
fst3 (a
a, b
_, c
_) = a
a

  dfsL :: STUArray s Vertex Bool -> DFSPath -> [Vertex] -> ST s (Maybe [Vertex])
  dfsL :: forall s.
STUArray s Vertex Bool
-> DFSPath -> [Vertex] -> ST s (Maybe [Vertex])
dfsL STUArray s Vertex Bool
_ DFSPath
_ [] = Maybe [Vertex] -> ST s (Maybe [Vertex])
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [Vertex]
forall a. Maybe a
Nothing
  dfsL STUArray s Vertex Bool
visited DFSPath
path (Vertex
v : [Vertex]
vs) = do
    Maybe [Vertex]
found <- STUArray s Vertex Bool
-> DFSPath -> Vertex -> ST s (Maybe [Vertex])
forall s.
STUArray s Vertex Bool
-> DFSPath -> Vertex -> ST s (Maybe [Vertex])
dfs STUArray s Vertex Bool
visited DFSPath
path Vertex
v
    case Maybe [Vertex]
found of
      Maybe [Vertex]
Nothing -> STUArray s Vertex Bool
-> DFSPath -> [Vertex] -> ST s (Maybe [Vertex])
forall s.
STUArray s Vertex Bool
-> DFSPath -> [Vertex] -> ST s (Maybe [Vertex])
dfsL STUArray s Vertex Bool
visited DFSPath
path [Vertex]
vs
      Just [Vertex]
cyc -> Maybe [Vertex] -> ST s (Maybe [Vertex])
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Vertex] -> Maybe [Vertex]
forall a. a -> Maybe a
Just [Vertex]
cyc)

  dfs :: STUArray s Vertex Bool -> DFSPath -> Vertex -> ST s (Maybe [Vertex])
  dfs :: forall s.
STUArray s Vertex Bool
-> DFSPath -> Vertex -> ST s (Maybe [Vertex])
dfs STUArray s Vertex Bool
visited p :: DFSPath
p@(DFSPath IntSet
pathMembers [Vertex]
path) Vertex
v
    | Vertex
v Vertex -> IntSet -> Bool
`IS.member` IntSet
pathMembers = Maybe [Vertex] -> ST s (Maybe [Vertex])
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe [Vertex] -> ST s (Maybe [Vertex]))
-> ([Vertex] -> Maybe [Vertex])
-> [Vertex]
-> ST s (Maybe [Vertex])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Vertex] -> Maybe [Vertex]
forall a. a -> Maybe a
Just ([Vertex] -> Maybe [Vertex])
-> ([Vertex] -> [Vertex]) -> [Vertex] -> Maybe [Vertex]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vertex
v Vertex -> [Vertex] -> [Vertex]
forall a. a -> [a] -> [a]
:) ([Vertex] -> [Vertex])
-> ([Vertex] -> [Vertex]) -> [Vertex] -> [Vertex]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Vertex] -> [Vertex]
forall a. [a] -> [a]
reverse ([Vertex] -> ST s (Maybe [Vertex]))
-> [Vertex] -> ST s (Maybe [Vertex])
forall a b. (a -> b) -> a -> b
$ (Vertex -> Bool) -> [Vertex] -> [Vertex]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
/= Vertex
v) [Vertex]
path
    | Bool
otherwise = do
        Bool
vis <- STUArray s Vertex Bool -> Vertex -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Vertex Bool
visited Vertex
v
        case Bool
vis of
          Bool
True -> Maybe [Vertex] -> ST s (Maybe [Vertex])
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [Vertex]
forall a. Maybe a
Nothing
          Bool
False -> STUArray s Vertex Bool
-> DFSPath -> [Vertex] -> ST s (Maybe [Vertex])
forall s.
STUArray s Vertex Bool
-> DFSPath -> [Vertex] -> ST s (Maybe [Vertex])
dfsL STUArray s Vertex Bool
visited (DFSPath -> Vertex -> DFSPath
appendPath DFSPath
p Vertex
v) (Graph
g Graph -> Vertex -> [Vertex]
forall i e. Ix i => Array i e -> i -> e
! Vertex
v)

failOnCyclicGraph ::
  Ord key =>
  Text ->
  (a -> Text) ->
  [(a, key, [key])] ->
  Either Text ()
failOnCyclicGraph :: forall key a.
Ord key =>
Text -> (a -> Text) -> [(a, key, [key])] -> Either Text ()
failOnCyclicGraph Text
graphType a -> Text
keyFunction [(a, key, [key])]
gEdges =
  Maybe [a] -> ([a] -> Either Text Any) -> Either Text ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([(a, key, [key])] -> Maybe [a]
forall key a. Ord key => [(a, key, [key])] -> Maybe [a]
findCycle [(a, key, [key])]
gEdges) (([a] -> Either Text Any) -> Either Text ())
-> ([a] -> Either Text Any) -> Either Text ()
forall a b. (a -> b) -> a -> b
$ \[a]
cyc ->
    Text -> Either Text Any
forall a b. a -> Either a b
Left (Text -> Either Text Any) -> Text -> Either Text Any
forall a b. (a -> b) -> a -> b
$
      [Text] -> Text
T.unwords
        [ Text
graphType
        , Text
"graph contains a cycle:"
        , Text -> Text
brackets (Text -> Text) -> ([a] -> Text) -> [a] -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Text] -> Text
T.intercalate Text
" -> " ([Text] -> Text) -> ([a] -> [Text]) -> [a] -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Text) -> [a] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Text
keyFunction ([a] -> Text) -> [a] -> Text
forall a b. (a -> b) -> a -> b
$ [a]
cyc
        ]