{-|
Module      : Data.Graph.AdjacencyList.DFS
Description : Depth-first search with topological sort and longest path
Copyright   : Thodoris Papakonstantinou, 2017-2026
License     : LGPL-3
Maintainer  : dev@tpapak.com
Stability   : experimental
Portability : POSIX

Depth-first search (DFS) on directed graphs.  Produces a topological ordering,
a visited-order list, and the set of discovered vertices.  Also provides
'longestPath' on DAGs and connectivity queries.
 -}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE BangPatterns #-}


module Data.Graph.AdjacencyList.DFS
  ( DFS (..)
  , dfs
  -- * Types
  , DAG
  , Distances
  -- * get longest path from a vertex to another
  , longestPath
  , postordering
  , areConnected
  , distances
  ) where

import Data.List
import Data.Maybe
import qualified Data.IntMap   as IM
import qualified Data.IntSet   as Set
import qualified Data.Sequence as Seq

import Data.Graph.AdjacencyList

-- | Result of a depth-first search from a single source vertex.
data DFS = DFS { DFS -> [Vertex]
topsort :: [Vertex]
                 -- ^ Vertices in reverse post-order (topological sort for DAGs).
               , DFS -> [Vertex]
visited :: [Vertex]
                 -- ^ Vertices in DFS visit order.
               , DFS -> IntSet
discovered   :: Set.IntSet
                 -- ^ Set of all discovered vertices.
               , DFS -> Vertex
called :: Int
                 -- ^ Number of DFS calls made.
               } deriving (DFS -> DFS -> Bool
(DFS -> DFS -> Bool) -> (DFS -> DFS -> Bool) -> Eq DFS
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DFS -> DFS -> Bool
== :: DFS -> DFS -> Bool
$c/= :: DFS -> DFS -> Bool
/= :: DFS -> DFS -> Bool
Eq, Vertex -> DFS -> ShowS
[DFS] -> ShowS
DFS -> String
(Vertex -> DFS -> ShowS)
-> (DFS -> String) -> ([DFS] -> ShowS) -> Show DFS
forall a.
(Vertex -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Vertex -> DFS -> ShowS
showsPrec :: Vertex -> DFS -> ShowS
$cshow :: DFS -> String
show :: DFS -> String
$cshowList :: [DFS] -> ShowS
showList :: [DFS] -> ShowS
Show)

initialDFS :: DFS
initialDFS :: DFS
initialDFS = DFS { topsort :: [Vertex]
topsort = []
                 , discovered :: IntSet
discovered   = IntSet
Set.empty
                 , visited :: [Vertex]
visited = []
                 , called :: Vertex
called = Vertex
0
                 }

-- | Depth first search
dfs :: Graph -> Vertex -> DFS
dfs :: Graph -> Vertex -> DFS
dfs Graph
g Vertex
s = 
  let vset :: IntSet
vset = [Vertex] -> IntSet
Set.fromList (Graph -> [Vertex]
vertices Graph
g)
  in if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Vertex -> IntSet -> Bool
Set.member Vertex
s IntSet
vset
     then DFS
initialDFS
     else
       let depthFirstSearch :: Vertex -> DFS -> DFS
           depthFirstSearch :: Vertex -> DFS -> DFS
depthFirstSearch Vertex
v DFS
ac
              | Vertex -> IntSet -> Bool
Set.member Vertex
v (DFS -> IntSet
discovered DFS
ac) = DFS
ac
              | Bool
otherwise =
              let -- Mark v as discovered BEFORE recursing (prevents revisits in cyclic graphs)
                  ac0 :: DFS
ac0 = DFS
ac { discovered = Set.insert v (discovered ac) }
                  ns :: [Vertex]
ns = Graph -> Neighbors
neighbors Graph
g Vertex
v
                  !ac' :: DFS
ac' = (DFS -> Vertex -> DFS) -> DFS -> [Vertex] -> DFS
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\DFS
ac'' Vertex
n -> if Bool -> Bool
not (Vertex -> IntSet -> Bool
Set.member Vertex
n (DFS -> IntSet
discovered DFS
ac''))
                                              then Vertex -> DFS -> DFS
depthFirstSearch Vertex
n DFS
ac''
                                              else DFS
ac''
                                ) DFS
ac0 [Vertex]
ns
                  res :: DFS
res = DFS
ac' { topsort = v : topsort ac'
                            -- Prepend to visited (reversed at end)
                            , visited = v : visited ac'
                            , called = called ac' + 1
                            }
               in DFS
res
           result :: DFS
result = Vertex -> DFS -> DFS
depthFirstSearch Vertex
s DFS
initialDFS
        in DFS
result { visited = reverse (visited result) }

-- | Post-order traversal (reverse of 'topsort').
postordering :: DFS -> [Vertex]
postordering :: DFS -> [Vertex]
postordering = [Vertex] -> [Vertex]
forall a. [a] -> [a]
reverse ([Vertex] -> [Vertex]) -> (DFS -> [Vertex]) -> DFS -> [Vertex]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DFS -> [Vertex]
topsort

-- | :)
type DAG = Graph

-- | Ginen a DAG and a vertex you get the distances
distances' :: DAG  -> Vertex -> IM.IntMap Vertex
distances' :: Graph -> Vertex -> IntMap Vertex
distances' Graph
g Vertex
s =
  let topsorted :: [Vertex]
topsorted = DFS -> [Vertex]
topsort (DFS -> [Vertex]) -> DFS -> [Vertex]
forall a b. (a -> b) -> a -> b
$ Graph -> Vertex -> DFS
dfs Graph
g Vertex
s
      initdists :: IntMap Vertex
initdists = (IntMap Vertex -> Vertex -> IntMap Vertex)
-> IntMap Vertex -> [Vertex] -> IntMap Vertex
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap Vertex
ac Vertex
v -> Vertex -> Vertex -> IntMap Vertex -> IntMap Vertex
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
v Vertex
0 IntMap Vertex
ac) IntMap Vertex
forall a. IntMap a
IM.empty ([Vertex] -> IntMap Vertex) -> [Vertex] -> IntMap Vertex
forall a b. (a -> b) -> a -> b
$ Graph -> [Vertex]
vertices Graph
g
   in (IntMap Vertex -> Vertex -> IntMap Vertex)
-> IntMap Vertex -> [Vertex] -> IntMap Vertex
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap Vertex
ac Vertex
v -> 
        let neis :: [Vertex]
neis = Graph -> Neighbors
neighbors Graph
g Vertex
v
            distv :: Vertex
distv = case Vertex -> IntMap Vertex -> Maybe Vertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v IntMap Vertex
ac of
                      Maybe Vertex
Nothing -> Vertex
0
                      Just Vertex
d -> Vertex
d
         in (IntMap Vertex -> Vertex -> IntMap Vertex)
-> IntMap Vertex -> [Vertex] -> IntMap Vertex
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap Vertex
dists' Vertex
nei -> 
           let neidist :: Vertex
neidist = case Vertex -> IntMap Vertex -> Maybe Vertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
nei IntMap Vertex
dists' of
                           Maybe Vertex
Nothing -> Vertex
0
                           Just Vertex
nd -> Vertex
nd
               newdist :: Vertex
newdist = Vertex -> Vertex -> Vertex
forall a. Ord a => a -> a -> a
max Vertex
neidist (Vertex
distvVertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+Vertex
1)
            in Vertex -> Vertex -> IntMap Vertex -> IntMap Vertex
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
nei Vertex
newdist IntMap Vertex
dists'
                  ) IntMap Vertex
ac [Vertex]
neis
      ) IntMap Vertex
initdists [Vertex]
topsorted

-- | Map from vertex to its distance (number of edges) from the source in a 'DAG'.
type Distances = IM.IntMap Vertex

-- | Ginen a DAG and a vertex you get the distances
distances :: DAG  -> DFS -> Vertex -> Distances
distances :: Graph -> DFS -> Vertex -> IntMap Vertex
distances Graph
g DFS
dfs' Vertex
s =
  let topsorted :: [Vertex]
topsorted = DFS -> [Vertex]
topsort (DFS -> [Vertex]) -> DFS -> [Vertex]
forall a b. (a -> b) -> a -> b
$ DFS
dfs'
      !initdists :: IntMap Vertex
initdists = (IntMap Vertex -> Vertex -> IntMap Vertex)
-> IntMap Vertex -> [Vertex] -> IntMap Vertex
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap Vertex
ac Vertex
v -> Vertex -> Vertex -> IntMap Vertex -> IntMap Vertex
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
v Vertex
0 IntMap Vertex
ac) IntMap Vertex
forall a. IntMap a
IM.empty ([Vertex] -> IntMap Vertex) -> [Vertex] -> IntMap Vertex
forall a b. (a -> b) -> a -> b
$ Graph -> [Vertex]
vertices Graph
g
   in (IntMap Vertex -> Vertex -> IntMap Vertex)
-> IntMap Vertex -> [Vertex] -> IntMap Vertex
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap Vertex
ac Vertex
v -> 
        let neis :: [Vertex]
neis = Graph -> Neighbors
neighbors Graph
g Vertex
v
            distv :: Vertex
distv = case Vertex -> IntMap Vertex -> Maybe Vertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v IntMap Vertex
ac of
                      Maybe Vertex
Nothing -> Vertex
0
                      Just Vertex
d -> Vertex
d
         in (IntMap Vertex -> Vertex -> IntMap Vertex)
-> IntMap Vertex -> [Vertex] -> IntMap Vertex
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntMap Vertex
dists' Vertex
nei -> 
           let neidist :: Vertex
neidist = case Vertex -> IntMap Vertex -> Maybe Vertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
nei IntMap Vertex
dists' of
                           Maybe Vertex
Nothing -> Vertex
0
                           Just Vertex
nd -> Vertex
nd
               newdist :: Vertex
newdist = Vertex -> Vertex -> Vertex
forall a. Ord a => a -> a -> a
max Vertex
neidist (Vertex
distvVertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+Vertex
1)
            in Vertex -> Vertex -> IntMap Vertex -> IntMap Vertex
forall a. Vertex -> a -> IntMap a -> IntMap a
IM.insert Vertex
nei Vertex
newdist IntMap Vertex
dists'
                  ) IntMap Vertex
ac [Vertex]
neis
      ) IntMap Vertex
initdists [Vertex]
topsorted

type TopologicalSorting = [Vertex]
-- |checks if s is predecessor of t
dependsOn :: TopologicalSorting -> Vertex -> Vertex -> Bool
dependsOn :: [Vertex] -> Vertex -> Vertex -> Bool
dependsOn [Vertex]
topsorted Vertex
t Vertex
s = Vertex -> [Vertex] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem Vertex
t (([Vertex], [Vertex]) -> [Vertex]
forall a b. (a, b) -> b
snd ((Vertex -> Bool) -> [Vertex] -> ([Vertex], [Vertex])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
(==) Vertex
s) [Vertex]
topsorted))

-- | Check whether vertex @v@ is reachable from vertex @u@ according to the
-- given distance map (distance > 0 means reachable; @u@ is reachable from itself).
areConnected :: Distances -> Vertex -> Vertex -> Bool
areConnected :: IntMap Vertex -> Vertex -> Vertex -> Bool
areConnected IntMap Vertex
dists Vertex
u Vertex
v = (Maybe Vertex -> Vertex
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Vertex -> Vertex) -> Maybe Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$ Vertex -> IntMap Vertex -> Maybe Vertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
v IntMap Vertex
dists) Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
> Vertex
0 Bool -> Bool -> Bool
|| Vertex
v Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
u

-- |Longest path from tail to nose
longestPath :: Graph -> Vertex -> Vertex -> [Edge]
longestPath :: Graph -> Vertex -> Vertex -> [Edge]
longestPath Graph
g Vertex
s Vertex
t =
  let dfs' :: DFS
dfs' = Graph -> Vertex -> DFS
dfs Graph
g Vertex
s
      topsorted :: [Vertex]
topsorted = DFS -> [Vertex]
topsort DFS
dfs'
      dists :: IntMap Vertex
dists = Graph -> DFS -> Vertex -> IntMap Vertex
distances Graph
g DFS
dfs' Vertex
s
      revg :: Graph
revg = Graph -> Graph
reverseGraph Graph
g
      disconnected :: [Vertex]
disconnected = (Vertex -> Bool) -> [Vertex] -> [Vertex]
forall a. (a -> Bool) -> [a] -> [a]
filter (\Vertex
n -> Bool -> Bool
not (IntMap Vertex -> Vertex -> Vertex -> Bool
areConnected IntMap Vertex
dists Vertex
s Vertex
n)) ([Vertex] -> [Vertex]) -> [Vertex] -> [Vertex]
forall a b. (a -> b) -> a -> b
$ Graph -> [Vertex]
vertices Graph
g
   in if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Vertex] -> Vertex -> Vertex -> Bool
dependsOn [Vertex]
topsorted Vertex
t Vertex
s
     then []
     else 
       if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Vertex] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Vertex]
disconnected
          then
            let cleangraph :: Graph
cleangraph = (Vertex -> Bool) -> Graph -> Graph
filterVertices (\Vertex
v -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Vertex -> [Vertex] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem Vertex
v [Vertex]
disconnected) Graph
g
             in Graph -> Vertex -> Vertex -> [Edge]
longestPath Graph
cleangraph Vertex
s Vertex
t
          else
            let path' :: Vertex -> [Edge] -> [Edge]
                path' :: Vertex -> [Edge] -> [Edge]
path' Vertex
v [Edge]
p 
                  | Vertex
v Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
s = [Edge]
p
                  | Bool
otherwise = 
                         let parents :: [Vertex]
parents = Graph -> Neighbors
neighbors Graph
revg Vertex
v
                          in if [Vertex] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Vertex]
parents
                                then []
                                else
                                  if [Vertex]
parents [Vertex] -> [Vertex] -> Bool
forall a. Eq a => a -> a -> Bool
== [Vertex
s]
                                   then (Vertex -> Vertex -> Edge
Edge Vertex
s Vertex
v)Edge -> [Edge] -> [Edge]
forall a. a -> [a] -> [a]
:[Edge]
p
                                   else 
                                     let pred :: Vertex
                                         pred :: Vertex
pred = (Vertex, Vertex) -> Vertex
forall a b. (a, b) -> a
fst ((Vertex, Vertex) -> Vertex) -> (Vertex, Vertex) -> Vertex
forall a b. (a -> b) -> a -> b
$ ((Vertex, Vertex) -> Vertex -> (Vertex, Vertex))
-> (Vertex, Vertex) -> [Vertex] -> (Vertex, Vertex)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
                                           (\(Vertex
prevmax,Vertex
maxdist) Vertex
parent ->
                                             let currentDist :: (Vertex, Vertex)
currentDist =
                                                   case Vertex -> IntMap Vertex -> Maybe Vertex
forall a. Vertex -> IntMap a -> Maybe a
IM.lookup Vertex
parent IntMap Vertex
dists of
                                                     Maybe Vertex
Nothing -> (Vertex
0,Vertex
0)
                                                     Just Vertex
d -> (Vertex
parent,Vertex
d)
                                              in if Vertex
maxdist Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
< (Vertex, Vertex) -> Vertex
forall a b. (a, b) -> b
snd (Vertex, Vertex)
currentDist
                                                    then (Vertex, Vertex)
currentDist
                                                    else (Vertex
prevmax,Vertex
maxdist)
                                                    ) (Vertex
0,Vertex
0) [Vertex]
parents
                                      in  Vertex -> [Edge] -> [Edge]
path' Vertex
pred ([Edge] -> [Edge]) -> [Edge] -> [Edge]
forall a b. (a -> b) -> a -> b
$ (Vertex -> Vertex -> Edge
Edge Vertex
pred Vertex
v)Edge -> [Edge] -> [Edge]
forall a. a -> [a] -> [a]
: [Edge]
p
             in Vertex -> [Edge] -> [Edge]
path' Vertex
t []