{-# LANGUAGE ScopedTypeVariables #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Algorithms.Graph.DFS
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--------------------------------------------------------------------------------
module Algorithms.Graph.DFS where

import           Control.Monad.ST            (ST, runST)
import           Data.Maybe
import           Data.PlanarGraph
import           Data.Tree
import qualified Data.Vector                 as V
import qualified Data.Vector.Generic         as GV
import qualified Data.Vector.Unboxed.Mutable as UMV
import qualified Data.IntSet as IntSet

-- | DFS on a planar graph.
--
-- Running time: \(O(n)\)
--
-- Note that since our planar graphs are always connected there is no need need
-- for dfs to take a list of start vertices.
dfs  :: forall s w v e f.
      PlanarGraph s w v e f -> VertexId s w -> Tree (VertexId s w)
dfs :: PlanarGraph s w v e f -> VertexId s w -> Tree (VertexId s w)
dfs PlanarGraph s w v e f
g = AdjacencyLists s w -> VertexId s w -> Tree (VertexId s w)
forall k (s :: k) (w :: World).
AdjacencyLists s w -> VertexId s w -> Tree (VertexId s w)
dfs' (PlanarGraph s w v e f -> AdjacencyLists s w
forall k (s :: k) (w :: World) v e f.
PlanarGraph s w v e f -> AdjacencyLists s w
adjacencyLists PlanarGraph s w v e f
g)

-- | Adjacency list representation of a graph: for each vertex we simply list
-- all connected neighbours.
type AdjacencyLists s w = V.Vector [VertexId s w]

-- | Transform into adjacencylist representation
adjacencyLists   :: PlanarGraph s w v e f -> AdjacencyLists s w
adjacencyLists :: PlanarGraph s w v e f -> AdjacencyLists s w
adjacencyLists PlanarGraph s w v e f
g = Vector (VertexId s w) -> [VertexId s w]
forall a. Vector a -> [a]
V.toList (Vector (VertexId s w) -> [VertexId s w])
-> (VertexId s w -> Vector (VertexId s w))
-> VertexId s w
-> [VertexId s w]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VertexId s w -> PlanarGraph s w v e f -> Vector (VertexId s w))
-> PlanarGraph s w v e f -> VertexId s w -> Vector (VertexId s w)
forall a b c. (a -> b -> c) -> b -> a -> c
flip VertexId s w -> PlanarGraph s w v e f -> Vector (VertexId s w)
forall k (s :: k) (w :: World) v e f.
VertexId s w -> PlanarGraph s w v e f -> Vector (VertexId s w)
neighboursOf PlanarGraph s w v e f
g (VertexId s w -> [VertexId s w])
-> Vector (VertexId s w) -> AdjacencyLists s w
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PlanarGraph s w v e f -> Vector (VertexId s w)
forall k (s :: k) (w :: World) v e f.
PlanarGraph s w v e f -> Vector (VertexId s w)
vertices' PlanarGraph s w v e f
g

-- | DFS, from a given vertex, on a graph in AdjacencyLists representation.
--
-- Running time: \(O(n)\)
dfs'          :: forall s w. AdjacencyLists s w -> VertexId s w -> Tree (VertexId s w)
dfs' :: AdjacencyLists s w -> VertexId s w -> Tree (VertexId s w)
dfs' AdjacencyLists s w
g VertexId s w
start = (forall s. ST s (Tree (VertexId s w))) -> Tree (VertexId s w)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Tree (VertexId s w))) -> Tree (VertexId s w))
-> (forall s. ST s (Tree (VertexId s w))) -> Tree (VertexId s w)
forall a b. (a -> b) -> a -> b
$ do
                 MVector s Bool
bv     <- Int -> Bool -> ST s (MVector (PrimState (ST s)) Bool)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UMV.replicate Int
n Bool
False -- bit vector of marks
                 -- start will be unvisited, thus the fromJust is safe
                 Maybe (Tree (VertexId s w)) -> Tree (VertexId s w)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Tree (VertexId s w)) -> Tree (VertexId s w))
-> ST s (Maybe (Tree (VertexId s w))) -> ST s (Tree (VertexId s w))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector s Bool
-> VertexId s w -> ST s (Maybe (Tree (VertexId s w)))
forall s'.
MVector s' Bool
-> VertexId s w -> ST s' (Maybe (Tree (VertexId s w)))
dfs'' MVector s Bool
bv VertexId s w
start
  where
    n :: Int
n = AdjacencyLists s w -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
GV.length AdjacencyLists s w
g

    neighs              :: VertexId s w -> [VertexId s w]
    neighs :: VertexId s w -> [VertexId s w]
neighs (VertexId Int
u) = AdjacencyLists s w
g AdjacencyLists s w -> Int -> [VertexId s w]
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
GV.! Int
u

    visit :: MVector (PrimState m) Bool -> VertexId s w -> m ()
visit   MVector (PrimState m) Bool
bv (VertexId Int
i) = MVector (PrimState m) Bool -> Int -> Bool -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UMV.write MVector (PrimState m) Bool
bv Int
i Bool
True
    visited :: MVector (PrimState m) a -> VertexId s w -> m a
visited MVector (PrimState m) a
bv (VertexId Int
i) = MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UMV.read  MVector (PrimState m) a
bv Int
i
    dfs''      :: UMV.MVector s' Bool -> VertexId s w
               -> ST s' (Maybe (Tree (VertexId s w)))
    dfs'' :: MVector s' Bool
-> VertexId s w -> ST s' (Maybe (Tree (VertexId s w)))
dfs'' MVector s' Bool
bv VertexId s w
u = MVector (PrimState (ST s')) Bool -> VertexId s w -> ST s' Bool
forall k (m :: * -> *) a (s :: k) (w :: World).
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> VertexId s w -> m a
visited MVector s' Bool
MVector (PrimState (ST s')) Bool
bv VertexId s w
u ST s' Bool
-> (Bool -> ST s' (Maybe (Tree (VertexId s w))))
-> ST s' (Maybe (Tree (VertexId s w)))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                   Bool
True  -> Maybe (Tree (VertexId s w)) -> ST s' (Maybe (Tree (VertexId s w)))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Tree (VertexId s w))
forall a. Maybe a
Nothing
                   Bool
False -> do
                              MVector (PrimState (ST s')) Bool -> VertexId s w -> ST s' ()
forall k (m :: * -> *) (s :: k) (w :: World).
PrimMonad m =>
MVector (PrimState m) Bool -> VertexId s w -> m ()
visit MVector s' Bool
MVector (PrimState (ST s')) Bool
bv VertexId s w
u
                              Tree (VertexId s w) -> Maybe (Tree (VertexId s w))
forall a. a -> Maybe a
Just (Tree (VertexId s w) -> Maybe (Tree (VertexId s w)))
-> ([Maybe (Tree (VertexId s w))] -> Tree (VertexId s w))
-> [Maybe (Tree (VertexId s w))]
-> Maybe (Tree (VertexId s w))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VertexId s w -> Forest (VertexId s w) -> Tree (VertexId s w)
forall a. a -> Forest a -> Tree a
Node VertexId s w
u (Forest (VertexId s w) -> Tree (VertexId s w))
-> ([Maybe (Tree (VertexId s w))] -> Forest (VertexId s w))
-> [Maybe (Tree (VertexId s w))]
-> Tree (VertexId s w)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (Tree (VertexId s w))] -> Forest (VertexId s w)
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Tree (VertexId s w))] -> Maybe (Tree (VertexId s w)))
-> ST s' [Maybe (Tree (VertexId s w))]
-> ST s' (Maybe (Tree (VertexId s w)))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VertexId s w -> ST s' (Maybe (Tree (VertexId s w))))
-> [VertexId s w] -> ST s' [Maybe (Tree (VertexId s w))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (MVector s' Bool
-> VertexId s w -> ST s' (Maybe (Tree (VertexId s w)))
forall s'.
MVector s' Bool
-> VertexId s w -> ST s' (Maybe (Tree (VertexId s w)))
dfs'' MVector s' Bool
bv) (VertexId s w -> [VertexId s w]
neighs VertexId s w
u)

-- | DFS, from a given vertex, on a graph in AdjacencyLists representation. Cycles are not removed.
-- If your graph may contain cycles, see 'dfsFilterCycles'.
--
-- Running time: \(O(k)\), where \(k\) is the number of nodes consumed.
dfsSensitive :: forall s w. (VertexId s w -> [VertexId s w]) -> VertexId s w -> Tree (VertexId s w)
dfsSensitive :: (VertexId s w -> [VertexId s w])
-> VertexId s w -> Tree (VertexId s w)
dfsSensitive VertexId s w -> [VertexId s w]
neighs VertexId s w
u = VertexId s w -> Forest (VertexId s w) -> Tree (VertexId s w)
forall a. a -> Forest a -> Tree a
Node VertexId s w
u (Forest (VertexId s w) -> Tree (VertexId s w))
-> Forest (VertexId s w) -> Tree (VertexId s w)
forall a b. (a -> b) -> a -> b
$ (VertexId s w -> Tree (VertexId s w))
-> [VertexId s w] -> Forest (VertexId s w)
forall a b. (a -> b) -> [a] -> [b]
map ((VertexId s w -> [VertexId s w])
-> VertexId s w -> Tree (VertexId s w)
forall k (s :: k) (w :: World).
(VertexId s w -> [VertexId s w])
-> VertexId s w -> Tree (VertexId s w)
dfsSensitive VertexId s w -> [VertexId s w]
neighs) (VertexId s w -> [VertexId s w]
neighs VertexId s w
u)

-- | Remove infinite cycles from a DFS search tree.
dfsFilterCycles :: Tree (VertexId s w) -> Tree (VertexId s w)
dfsFilterCycles :: Tree (VertexId s w) -> Tree (VertexId s w)
dfsFilterCycles = IntSet -> Tree (VertexId s w) -> Tree (VertexId s w)
forall k (s :: k) (w :: World).
IntSet -> Tree (VertexId s w) -> Tree (VertexId s w)
worker IntSet
IntSet.empty
  where
    worker :: IntSet -> Tree (VertexId s w) -> Tree (VertexId s w)
worker IntSet
seen (Node VertexId s w
root Forest (VertexId s w)
forest) = VertexId s w -> Forest (VertexId s w) -> Tree (VertexId s w)
forall a. a -> Forest a -> Tree a
Node VertexId s w
root
      [ VertexId s w -> Forest (VertexId s w) -> Tree (VertexId s w)
forall a. a -> Forest a -> Tree a
Node (Int -> VertexId s w
forall k (s :: k) (w :: World). Int -> VertexId s w
VertexId Int
v) ((Tree (VertexId s w) -> Tree (VertexId s w))
-> Forest (VertexId s w) -> Forest (VertexId s w)
forall a b. (a -> b) -> [a] -> [b]
map (IntSet -> Tree (VertexId s w) -> Tree (VertexId s w)
worker (Int -> IntSet -> IntSet
IntSet.insert Int
v IntSet
seen)) Forest (VertexId s w)
sub)
      | Node (VertexId Int
v) Forest (VertexId s w)
sub <- Forest (VertexId s w)
forest
      , Int
v Int -> IntSet -> Bool
`IntSet.notMember` IntSet
seen
      ]