{-|
Module      : Data.Graph.AdjacencyList
Description : Core graph types and constructors
Copyright   : Thodoris Papakonstantinou, 2017-2026
License     : LGPL-3
Maintainer  : dev@tpapak.com
Stability   : experimental
Portability : POSIX

Core types and constructors for directed graphs using adjacency list
representation.

A 'Graph' stores its vertex set, an 'EdgeMap' for edge-attribute lookup,
and a closure-based 'Neighbors' function for O(log V) neighbor access.
Undirected graphs are represented by including both directions of each edge.
 -}

{-# LANGUAGE DeriveGeneric #-}  

module Data.Graph.AdjacencyList
    ( Vertex (..)
    , Edge (..)
    , Neighbors (..)
    , EdgeMap (..)
    -- * Graph definition
    , Graph (..)
    , fromTuple
    , toTuple
    -- * createGraph: Graph constructor
    , createGraph
    -- * graph from list of Edges
    , graphFromEdges
    , edges
    , reverseEdge
    , reverseEdges
    , reverseGraph
    -- * filterVertices
    , filterVertices
    -- * filterEdges
    , filterEdges
    -- * makeUndirected
    , makeUndirected
    , adjacentEdges
    , edgesFromNeighbors
    , adjacencyMap
    , edgeExists
    , edgeIndex
    , from
    , to
    , numVertices
    , numEdges
    , removeReverseEdges
    , completeGraph
    ) where

import Data.List
import Data.List.Unique
import Data.Maybe
import qualified Data.Map.Lazy as M
import qualified Data.IntMap.Lazy as IM
import qualified Data.Set as Set
import qualified GHC.Generics as Gen
import qualified Data.Binary as Bin

-- | A vertex identifier (non-negative integer).
type Vertex = Int

-- | A directed edge from one vertex to another.
data Edge = Edge Vertex Vertex 
  deriving (Eq Edge
Eq Edge =>
(Edge -> Edge -> Ordering)
-> (Edge -> Edge -> Bool)
-> (Edge -> Edge -> Bool)
-> (Edge -> Edge -> Bool)
-> (Edge -> Edge -> Bool)
-> (Edge -> Edge -> Edge)
-> (Edge -> Edge -> Edge)
-> Ord Edge
Edge -> Edge -> Bool
Edge -> Edge -> Ordering
Edge -> Edge -> Edge
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Edge -> Edge -> Ordering
compare :: Edge -> Edge -> Ordering
$c< :: Edge -> Edge -> Bool
< :: Edge -> Edge -> Bool
$c<= :: Edge -> Edge -> Bool
<= :: Edge -> Edge -> Bool
$c> :: Edge -> Edge -> Bool
> :: Edge -> Edge -> Bool
$c>= :: Edge -> Edge -> Bool
>= :: Edge -> Edge -> Bool
$cmax :: Edge -> Edge -> Edge
max :: Edge -> Edge -> Edge
$cmin :: Edge -> Edge -> Edge
min :: Edge -> Edge -> Edge
Ord, (forall x. Edge -> Rep Edge x)
-> (forall x. Rep Edge x -> Edge) -> Generic Edge
forall x. Rep Edge x -> Edge
forall x. Edge -> Rep Edge x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Edge -> Rep Edge x
from :: forall x. Edge -> Rep Edge x
$cto :: forall x. Rep Edge x -> Edge
to :: forall x. Rep Edge x -> Edge
Gen.Generic)
instance Bin.Binary Edge

instance Show Edge where
 show :: Edge -> String
show (Edge Int
s Int
t) = String
"[" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
s String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"->" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
t String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"]"

instance Eq Edge where
  Edge
a == :: Edge -> Edge -> Bool
== Edge
b = Edge -> Int
from Edge
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Edge -> Int
from Edge
b Bool -> Bool -> Bool
&& Edge -> Int
to Edge
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Edge -> Int
to Edge
b

-- | Map from edges to their sequential index (1-based).
-- Used for edge-attribute lookup.
type EdgeMap = M.Map Edge Int

-- | Takes vertex and outputs neighboring vertices.
-- The Neighbors type is the function from a vertex to its neighbors
type Neighbors = (Vertex -> [Vertex])

-- | Graph definition of directed Graphs 
-- undirected graphs should include reverse edges.
data Graph = 
  Graph { Graph -> [Int]
vertices :: [Vertex] -- ^ The domain of the `neighbors` function. 
        -- It is usefull for finite graphs.
        , Graph -> EdgeMap
edgeMap :: EdgeMap -- ^ The edge map is necessary 
        -- for appointing edge attributes
        , Graph -> Neighbors
neighbors :: Neighbors -- ^ The `Adjacency List`
        }

-- | Check whether an edge exists in the graph.
edgeExists :: Graph -> Edge -> Bool
edgeExists :: Graph -> Edge -> Bool
edgeExists Graph
g Edge
e = Edge -> EdgeMap -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member Edge
e (Graph -> EdgeMap
edgeMap Graph
g)

-- | Gives the position of the edge to the edges list
edgeIndex :: Graph -> Edge -> Maybe Int
edgeIndex :: Graph -> Edge -> Maybe Int
edgeIndex Graph
g Edge
e = Edge -> EdgeMap -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Edge
e (EdgeMap -> Maybe Int) -> EdgeMap -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Graph -> EdgeMap
edgeMap Graph
g

-- | All edges of the graph, in 'EdgeMap' key order.
edges :: Graph -> [Edge]
edges :: Graph -> [Edge]
edges Graph
g = 
  ((Edge, Int) -> Edge) -> [(Edge, Int)] -> [Edge]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Edge, Int) -> Edge
forall a b. (a, b) -> a
fst ([(Edge, Int)] -> [Edge]) -> [(Edge, Int)] -> [Edge]
forall a b. (a -> b) -> a -> b
$ EdgeMap -> [(Edge, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (EdgeMap -> [(Edge, Int)]) -> EdgeMap -> [(Edge, Int)]
forall a b. (a -> b) -> a -> b
$ Graph -> EdgeMap
edgeMap Graph
g

edgeMapFromEdges :: [Edge] -> EdgeMap
edgeMapFromEdges :: [Edge] -> EdgeMap
edgeMapFromEdges [Edge]
es =
  [(Edge, Int)] -> EdgeMap
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Edge, Int)] -> EdgeMap) -> [(Edge, Int)] -> EdgeMap
forall a b. (a -> b) -> a -> b
$ [Edge] -> [Int] -> [(Edge, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Edge]
es [Int
1..]

-- | Source vertex of an edge.
from :: Edge -> Vertex
from :: Edge -> Int
from (Edge Int
s Int
t) = Int
s

-- | Target vertex of an edge.
to :: Edge -> Vertex
to :: Edge -> Int
to (Edge Int
s Int
t) = Int
t

-- | Construct an 'Edge' from a @(source, target)@ tuple.
fromTuple :: (Vertex, Vertex) -> Edge
fromTuple :: (Int, Int) -> Edge
fromTuple (Int
s,Int
t) = Int -> Int -> Edge
Edge Int
s Int
t

-- | Convert an 'Edge' to a @(source, target)@ tuple.
toTuple :: Edge -> (Vertex, Vertex)
toTuple :: Edge -> (Int, Int)
toTuple (Edge Int
s Int
t) = (Int
s,Int
t)

-- | Reverse the direction of an edge.
reverseEdge :: Edge -> Edge
reverseEdge :: Edge -> Edge
reverseEdge (Edge Int
s Int
t) = Int -> Int -> Edge
Edge Int
t Int
s

-- | All edges of the graph with reversed direction.
reverseEdges :: Graph -> [Edge]
reverseEdges :: Graph -> [Edge]
reverseEdges Graph
g = (Edge -> Edge) -> [Edge] -> [Edge]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Edge -> Edge
reverseEdge ([Edge] -> [Edge]) -> [Edge] -> [Edge]
forall a b. (a -> b) -> a -> b
$ Graph -> [Edge]
edges Graph
g

-- | Number of vertices in the graph.
numVertices :: Graph -> Int
numVertices :: Graph -> Int
numVertices Graph
g = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ Graph -> [Int]
vertices Graph
g

-- | Number of edges in the graph.
numEdges :: Graph -> Int
numEdges :: Graph -> Int
numEdges Graph
g = [Edge] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Edge] -> Int) -> [Edge] -> Int
forall a b. (a -> b) -> a -> b
$ Graph -> [Edge]
edges Graph
g


instance Eq Graph where
  == :: Graph -> Graph -> Bool
(==) Graph
g1 Graph
g2 = ([Int] -> [Int]
forall a. Ord a => [a] -> [a]
sort (Graph -> [Int]
vertices Graph
g1) [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> [Int]
forall a. Ord a => [a] -> [a]
sort (Graph -> [Int]
vertices Graph
g2))
               Bool -> Bool -> Bool
&& ([Edge] -> [Edge]
forall a. Ord a => [a] -> [a]
sort (Graph -> [Edge]
edges Graph
g1) [Edge] -> [Edge] -> Bool
forall a. Eq a => a -> a -> Bool
== [Edge] -> [Edge]
forall a. Ord a => [a] -> [a]
sort (Graph -> [Edge]
edges Graph
g2))

instance Show Graph where
  show :: Graph -> String
show Graph
g = String
"vertices: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [Int] -> String
forall a. Show a => a -> String
show (Graph -> [Int]
vertices Graph
g) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++
            String
"edges: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [Edge] -> String
forall a. Show a => a -> String
show (Graph -> [Edge]
edges Graph
g) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\n"

-- | Graph constructor given a neighbors function
createGraph :: [Vertex] -> Neighbors -> Graph
createGraph :: [Int] -> Neighbors -> Graph
createGraph [Int]
vs Neighbors
neis =
  let emap :: EdgeMap
emap = [Edge] -> EdgeMap
edgeMapFromEdges ([Edge] -> EdgeMap) -> [Edge] -> EdgeMap
forall a b. (a -> b) -> a -> b
$ Neighbors -> [Int] -> [Edge]
edgesFromNeighbors Neighbors
neis [Int]
vs
   in Graph { vertices :: [Int]
vertices = [Int]
vs
            , neighbors :: Neighbors
neighbors = Neighbors
neis
            , edgeMap :: EdgeMap
edgeMap = EdgeMap
emap
            }

-- | Graph constructor given a list of edges.
--
-- Builds the adjacency map in a single O(E) pass using 'IM.fromListWith',
-- then wraps it in a closure for O(log V) neighbor lookup.
graphFromEdges :: [Edge] -> Graph
graphFromEdges :: [Edge] -> Graph
graphFromEdges [Edge]
es = 
  let vs :: [Int]
vs = Set Int -> [Int]
forall a. Set a -> [a]
Set.toList (Set Int -> [Int]) -> Set Int -> [Int]
forall a b. (a -> b) -> a -> b
$ (Set Int -> Edge -> Set Int) -> Set Int -> [Edge] -> Set Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Set Int
ac (Edge Int
u Int
v) ->
             Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
Set.insert Int
u (Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
Set.insert Int
v Set Int
ac)) Set Int
forall a. Set a
Set.empty [Edge]
es
      esmap :: EdgeMap
esmap = [Edge] -> EdgeMap
edgeMapFromEdges [Edge]
es
      -- Build adjacency map in one pass: O(E log V) via fromListWith
      neimap :: IntMap [Int]
neimap = ([Int] -> [Int] -> [Int]) -> [(Int, [Int])] -> IntMap [Int]
forall a. (a -> a -> a) -> [(Int, a)] -> IntMap a
IM.fromListWith [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
(++)
                  ([(Int, [Int])] -> IntMap [Int]) -> [(Int, [Int])] -> IntMap [Int]
forall a b. (a -> b) -> a -> b
$ (Edge -> (Int, [Int])) -> [Edge] -> [(Int, [Int])]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Edge Int
u Int
v) -> (Int
u, [Int
v])) [Edge]
es
      neis :: Neighbors
neis Int
v = case Int -> IntMap [Int] -> Maybe [Int]
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
v IntMap [Int]
neimap of
                 Maybe [Int]
Nothing -> []
                 Just [Int]
ns -> [Int]
ns
   in Graph { vertices :: [Int]
vertices = [Int]
vs
            , edgeMap :: EdgeMap
edgeMap = EdgeMap
esmap
            , neighbors :: Neighbors
neighbors = Neighbors
neis
            }

-- | Enumerate all edges implied by a 'Neighbors' function over a vertex set.
edgesFromNeighbors :: Neighbors -> [Vertex] -> [Edge]
edgesFromNeighbors :: Neighbors -> [Int] -> [Edge]
edgesFromNeighbors Neighbors
neis [Int]
vs = 
  let allneis :: [(Int, [Int])]
allneis = (Int -> (Int, [Int])) -> [Int] -> [(Int, [Int])]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
v -> (Int
v,Neighbors
neis Int
v)) [Int]
vs
   in ((Int, [Int]) -> [Edge] -> [Edge])
-> [Edge] -> [(Int, [Int])] -> [Edge]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(Int
v,[Int]
nv) [Edge]
ac -> 
             ((Int -> Edge) -> [Int] -> [Edge]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
n -> Int -> Int -> Edge
Edge Int
v Int
n) [Int]
nv) [Edge] -> [Edge] -> [Edge]
forall a. [a] -> [a] -> [a]
++ [Edge]
ac
             ) [] [(Int, [Int])]
allneis

-- | All outgoing edges from a vertex.
adjacentEdges :: Graph -> Vertex -> [Edge]
adjacentEdges :: Graph -> Int -> [Edge]
adjacentEdges Graph
g Int
v = (Int -> Edge) -> [Int] -> [Edge]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
n -> Int -> Int -> Edge
Edge Int
v Int
n) ([Int] -> [Edge]) -> [Int] -> [Edge]
forall a b. (a -> b) -> a -> b
$ Graph -> Neighbors
neighbors Graph
g Int
v

-- | Build an explicit adjacency map from the graph's 'Neighbors' closure.
adjacencyMap :: Graph -> IM.IntMap [Vertex]
adjacencyMap :: Graph -> IntMap [Int]
adjacencyMap Graph
g = [(Int, [Int])] -> IntMap [Int]
forall a. [(Int, a)] -> IntMap a
IM.fromList ([(Int, [Int])] -> IntMap [Int]) -> [(Int, [Int])] -> IntMap [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> (Int, [Int])) -> [Int] -> [(Int, [Int])]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
v -> (Int
v, (Graph -> Neighbors
neighbors Graph
g Int
v))) [Int]
vs
                 where vs :: [Int]
vs = Graph -> [Int]
vertices Graph
g

-- | Reverse all edges in the graph.
reverseGraph :: Graph -> Graph
reverseGraph :: Graph -> Graph
reverseGraph Graph
g =
  [Edge] -> Graph
graphFromEdges ([Edge] -> Graph) -> [Edge] -> Graph
forall a b. (a -> b) -> a -> b
$ Graph -> [Edge]
reverseEdges Graph
g

-- | Get the subgraph of a graph by including vertices satisfying given predicate.
filterVertices :: (Vertex -> Bool) -- ^ filter predicate
               -> Graph
               -> Graph
filterVertices :: (Int -> Bool) -> Graph -> Graph
filterVertices Int -> Bool
f Graph
g =
  let oldvs :: [Int]
oldvs = Graph -> [Int]
vertices Graph
g
      vs :: [Int]
vs = (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter Int -> Bool
f [Int]
oldvs 
      neis :: Neighbors
neis Int
v = 
        let ns :: [Int]
ns = Graph -> Neighbors
neighbors Graph
g Int
v
         in (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter Int -> Bool
f [Int]
ns
   in [Int] -> Neighbors -> Graph
createGraph [Int]
vs Neighbors
neis

-- | Get the subgraph of a graph by including edges satisfying given predicate.
filterEdges :: (Edge -> Bool) -> Graph -> Graph
filterEdges :: (Edge -> Bool) -> Graph -> Graph
filterEdges Edge -> Bool
f Graph
g =
  let vs :: [Int]
vs = Graph -> [Int]
vertices Graph
g
      neis :: Neighbors
neis Int
v = 
        let neis :: [Int]
neis = Graph -> Neighbors
neighbors Graph
g Int
v
         in (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (\Int
n -> Edge -> Bool
f (Int -> Int -> Edge
Edge Int
v Int
n)) [Int]
neis
   in [Int] -> Neighbors -> Graph
createGraph [Int]
vs Neighbors
neis

-- | Make a graph undirected by adding all missing reverse edges.
makeUndirected :: Graph -- ^ directed graph
               -> Graph -- ^ undirected graph
makeUndirected :: Graph -> Graph
makeUndirected Graph
g =
  let rg :: Graph
rg = Graph -> Graph
reverseGraph Graph
g
      vs :: [Int]
vs = Graph -> [Int]
vertices Graph
g
      newnei :: Neighbors
newnei Int
v = 
        let nei :: [Int]
nei = Graph -> Neighbors
neighbors Graph
g Int
v
            rnei :: [Int]
rnei = Graph -> Neighbors
neighbors Graph
rg Int
v
         in [Int] -> [Int]
forall a. Ord a => [a] -> [a]
sortUniq ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ [Int]
nei [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
rnei
   in [Int] -> Neighbors -> Graph
createGraph [Int]
vs Neighbors
newnei

-- | Make a graph directed by removing randomly reverse edges
removeReverseEdges :: Graph -- ^ Graph with reverse edges
                   -> Graph -- ^ Directected graph
removeReverseEdges :: Graph -> Graph
removeReverseEdges Graph
g =
  let unes :: [Edge]
unes = [Edge] -> [Edge]
forall a. Ord a => [a] -> [a]
sort ([Edge] -> [Edge]) -> [Edge] -> [Edge]
forall a b. (a -> b) -> a -> b
$ Graph -> [Edge]
edges Graph
g
      dires :: [Edge]
dires = (Edge -> Bool) -> [Edge] -> [Edge]
forall a. (a -> Bool) -> [a] -> [a]
filter (\Edge
e -> Edge -> [Edge] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem (Edge -> Edge
reverseEdge Edge
e) 
                             ((Edge -> Bool) -> [Edge] -> [Edge]
forall a. (a -> Bool) -> [a] -> [a]
filter (\Edge
e' -> Edge
e' Edge -> Edge -> Bool
forall a. Ord a => a -> a -> Bool
> Edge
e) [Edge]
unes) 
                     ) [Edge]
unes
   in [Edge] -> Graph
graphFromEdges [Edge]
dires


-- | Complete undirected graph from number of vertices
completeGraph :: Int -> Graph
completeGraph :: Int -> Graph
completeGraph Int
n =
  let es :: [Edge]
es = [Edge
e | Edge
e <- Int -> Int -> Edge
Edge (Int -> Int -> Edge) -> [Int] -> [Int -> Edge]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
1..Int
n] [Int -> Edge] -> [Int] -> [Edge]
forall a b. [a -> b] -> [a] -> [b]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Int
1..Int
n], (\(Edge Int
s Int
t) -> Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
t ) Edge
e]
   in [Edge] -> Graph
graphFromEdges [Edge]
es