packages feed

algebraic-edge-graphs-0.1.1: src/EdgeGraph/AdjacencyMap/Internal.hs

-----------------------------------------------------------------------------
-- |
-- Module     : EdgeGraph.AdjacencyMap.Internal
-- Copyright  : (c) Jack Liell-Cock 2025-2026
-- License    : MIT (see the file LICENSE)
-- Maintainer : jackliellcock@gmail.com
-- Stability  : unstable
--
-- This module exposes the implementation of edge-indexed adjacency maps.
-- An t'AdjacencyMap' stores, for each edge, four sets describing its
-- neighbourhood in the flow representation. This is equivalent to 'EdgeGraph.Incidence.Internal.Incidence'
-- but indexed by edge for O(log n) lookups.
--
-- The API is unstable and unsafe. Where possible use the non-internal module
-- "EdgeGraph.AdjacencyMap" instead.
--
-----------------------------------------------------------------------------
module EdgeGraph.AdjacencyMap.Internal (
  -- * Data structure
  Adjacency (..), AdjacencyMap (..), consistent,

  -- * Conversion
  toIncidence, fromIncidence,

  -- * Basic graph construction primitives
  empty, edge, overlay, into, pits, tips, edges, fromNodeList, fromIncidenceList,

  -- * Graph properties
  nodeList, nodeSet, edgeSet, edgeList, adjacencyList,
  edgeCount, nodeCount, isEmpty, hasEdge,

  -- * Graph transformation
  removeEdge, detachPit, detachTip, gmap, induce
) where

import Data.Map.Strict (Map)
import Data.Set (Set)

import qualified Data.Map.Strict              as Map
import qualified Data.Set                     as Set
import qualified EdgeGraph.Class              as C
import qualified EdgeGraph.Incidence.Internal as I

-- | Neighbourhood information for a single edge in the graph.
--
-- For an edge @a@, this records its relationships to other edges through
-- the two nodes it connects (its source and sink in the flow representation):
--
-- * 'forks': all edges departing from the same source node as @a@ (including @a@)
-- * 'joins': all edges arriving at the same sink node as @a@ (including @a@)
-- * 'preds': edges arriving at @a@'s source node (predecessors)
-- * 'succs': edges departing from @a@'s sink node (successors)
data Adjacency a = Adjacency
  { forks :: Set a
  , joins :: Set a
  , preds :: Set a
  , succs :: Set a
  } deriving (Eq, Ord, Show)

-- | The t'AdjacencyMap' data type represents an edge-indexed graph as a map
-- from each edge to its t'Adjacency' information. This is an equivalent
-- representation to 'EdgeGraph.Incidence.Internal.Incidence' (the canonical flow representation for
-- algebraic edge graphs), but indexed by edge for efficient lookups.
--
-- The 'Eq' instance is derived from the underlying 'Map' and is correct
-- because the representation is canonical (one-to-one with 'EdgeGraph.Incidence.Internal.Incidence').
newtype AdjacencyMap a = AdjacencyMap
  { adjacencyMap :: Map a (Adjacency a)
  } deriving Eq

instance (Ord a, Show a) => Show (AdjacencyMap a) where
  show = show . toIncidence

instance Ord a => C.EdgeGraph (AdjacencyMap a) where
  type Edge (AdjacencyMap a) = a
  empty   = empty
  edge    = edge
  overlay = overlay
  into    = into
  pits    = pits
  tips    = tips

-- | Convert an t'AdjacencyMap' to its equivalent 'EdgeGraph.Incidence.Internal.Incidence' representation.
--
-- For each edge, reconstruct its source node (pitNode) and sink node (tipNode),
-- then collect all unique nodes into a set.
--
-- @
-- 'toIncidence' 'EdgeGraph.AdjacencyMap.Internal.empty'    == 'EdgeGraph.Incidence.Internal.empty'
-- 'toIncidence' ('EdgeGraph.AdjacencyMap.Internal.edge' x) == 'EdgeGraph.Incidence.Internal.edge' x
-- @
toIncidence :: Ord a => AdjacencyMap a -> I.Incidence a
toIncidence (AdjacencyMap m)
  | Map.null m = I.Incidence Set.empty
  | otherwise  = I.Incidence $ Set.fromList allNodes
  where
    allNodes = concatMap nodesPair (Map.elems m)
    nodesPair adj =
      [ I.Node (preds adj) (forks adj)   -- source node
      , I.Node (joins adj) (succs adj)   -- sink node
      ]

-- | Convert an 'EdgeGraph.Incidence.Internal.Incidence' to an t'AdjacencyMap'.
--
-- Scans all nodes once to build edge-to-source and edge-to-sink maps,
-- then constructs the t'Adjacency' record for each edge.
--
-- @
-- 'fromIncidence' 'EdgeGraph.Incidence.Internal.empty'    == 'EdgeGraph.AdjacencyMap.Internal.empty'
-- 'fromIncidence' ('EdgeGraph.Incidence.Internal.edge' x) == 'EdgeGraph.AdjacencyMap.Internal.edge' x
-- @
fromIncidence :: Ord a => I.Incidence a -> AdjacencyMap a
fromIncidence (I.Incidence ns)
  | Set.null ns = AdjacencyMap Map.empty
  | otherwise   = AdjacencyMap $ Map.fromSet buildAdj allEdges
  where
    nl = Set.toList ns
    -- Build maps: edge -> the node where it's in pits (source),
    --             edge -> the node where it's in tips (sink)
    sourceMap = Map.fromList
      [ (a, n) | n <- nl, a <- Set.toList (I.nodePits n) ]
    sinkMap = Map.fromList
      [ (a, n) | n <- nl, a <- Set.toList (I.nodeTips n) ]
    allEdges = Map.keysSet sourceMap
    buildAdj a =
        let srcNode  = sourceMap Map.! a
            sinkNode = sinkMap   Map.! a
        in Adjacency
          { forks = I.nodePits srcNode
          , joins = I.nodeTips sinkNode
          , preds = I.nodeTips srcNode
          , succs = I.nodePits sinkNode
          }

-- | Check if an t'AdjacencyMap' is internally consistent.
--
-- Verifies the following invariants:
--
-- 1. Self-membership: @a ∈ forks(a)@ and @a ∈ joins(a)@
-- 2. Group equivalence: @b ∈ forks(a)@ implies @forks(b) = forks(a)@
-- 3. Group equivalence: @b ∈ joins(a)@ implies @joins(b) = joins(a)@
-- 4. Cross-consistency: @b ∈ preds(a)@ implies @succs(b) = forks(a)@ and @joins(b) = preds(a)@
-- 5. Cross-consistency: @b ∈ succs(a)@ implies @preds(b) = joins(a)@ and @forks(b) = succs(a)@
-- 6. All referenced edges exist in the map
--
-- @
-- consistent 'EdgeGraph.AdjacencyMap.Internal.empty'         == True
-- consistent ('EdgeGraph.AdjacencyMap.Internal.edge' x)      == True
-- consistent ('overlay' x y) == True
-- consistent ('into' x y)    == True
-- @
consistent :: Ord a => AdjacencyMap a -> Bool
consistent (AdjacencyMap m) =
  selfMembership && forkEquiv && joinEquiv &&
  predCross && succCross && allExist
  where
    keys = Map.keysSet m
    entries = Map.toList m
    selfMembership = all (\(a, adj) ->
      Set.member a (forks adj) && Set.member a (joins adj)) entries
    forkEquiv = all (\(_, adj) ->
      all (\b -> case Map.lookup b m of
        Just adjB -> forks adjB == forks adj
        Nothing   -> False
      ) (Set.toList $ forks adj)) entries
    joinEquiv = all (\(_, adj) ->
      all (\b -> case Map.lookup b m of
        Just adjB -> joins adjB == joins adj
        Nothing   -> False
      ) (Set.toList $ joins adj)) entries
    predCross = all (\(_, adj) ->
      all (\b -> case Map.lookup b m of
        Just adjB -> succs adjB == forks adj
                     && joins adjB == preds adj
        Nothing   -> False
      ) (Set.toList $ preds adj)) entries
    succCross = all (\(_, adj) ->
      all (\b -> case Map.lookup b m of
        Just adjB -> preds adjB == joins adj
                     && forks adjB == succs adj
        Nothing   -> False
      ) (Set.toList $ succs adj)) entries
    allExist = all (\(_, adj) ->
      Set.isSubsetOf (forks adj) keys &&
      Set.isSubsetOf (joins adj) keys &&
      Set.isSubsetOf (preds adj) keys &&
      Set.isSubsetOf (succs adj) keys) entries

-- | Construct the /empty graph/.
-- Complexity: /O(1)/ time and memory.
--
-- @
-- 'isEmpty' empty == True
-- @
empty :: AdjacencyMap a
empty = AdjacencyMap Map.empty

-- | Construct the graph comprising /a single edge/.
-- Complexity: /O(1)/ time and memory.
--
-- @
-- 'isEmpty' ('EdgeGraph.AdjacencyMap.Internal.edge' x)   == False
-- 'hasEdge' x ('EdgeGraph.AdjacencyMap.Internal.edge' x) == True
-- 'edgeCount' ('EdgeGraph.AdjacencyMap.Internal.edge' x) == 1
-- 'nodeCount' ('EdgeGraph.AdjacencyMap.Internal.edge' x) == 2
-- @
edge :: Ord a => a -> AdjacencyMap a
edge a = AdjacencyMap $ Map.singleton a $ Adjacency
  { forks = Set.singleton a
  , joins = Set.singleton a
  , preds = Set.empty
  , succs = Set.empty
  }

-- | /Overlay/ two graphs. This computes the least upper bound of two flow
-- representations by merging nodes that share edges.
--
-- @
-- 'isEmpty' ('overlay' x y)       == 'isEmpty' x && 'isEmpty' y
-- 'overlay' 'EdgeGraph.AdjacencyMap.Internal.empty' x             == x
-- 'overlay' x     'EdgeGraph.AdjacencyMap.Internal.empty'         == x
-- 'overlay' x     y             == 'overlay' y x
-- 'overlay' x     ('overlay' y z) == 'overlay' ('overlay' x y) z
-- 'overlay' x     x             == x
-- @
overlay :: Ord a => AdjacencyMap a -> AdjacencyMap a -> AdjacencyMap a
overlay x y = fromIncidence $ I.overlay (toIncidence x) (toIncidence y)

-- | /Into/ two graphs. Connects the sink side of the left graph to the
-- source side of the right graph, creating a sequential composition.
--
-- @
-- 'isEmpty' ('into' x y) == 'isEmpty' x && 'isEmpty' y
-- 'into' 'EdgeGraph.AdjacencyMap.Internal.empty' x       == x
-- 'into' x 'EdgeGraph.AdjacencyMap.Internal.empty'       == x
-- @
into :: Ord a => AdjacencyMap a -> AdjacencyMap a -> AdjacencyMap a
into x y = fromIncidence $ I.into (toIncidence x) (toIncidence y)

-- | /Pits/ two graphs. Connects where outgoing edges (pits) overlap,
-- causing source-side merging.
--
-- @
-- 'isEmpty' ('pits' x y) == 'isEmpty' x && 'isEmpty' y
-- @
pits :: Ord a => AdjacencyMap a -> AdjacencyMap a -> AdjacencyMap a
pits x y = fromIncidence $ I.pits (toIncidence x) (toIncidence y)

-- | /Tips/ two graphs. Connects where incoming edges (tips) overlap,
-- causing sink-side merging.
--
-- @
-- 'isEmpty' ('tips' x y) == 'isEmpty' x && 'isEmpty' y
-- @
tips :: Ord a => AdjacencyMap a -> AdjacencyMap a -> AdjacencyMap a
tips x y = fromIncidence $ I.tips (toIncidence x) (toIncidence y)

-- | Construct a graph from a given list of edges by overlaying them.
--
-- @
-- edges []  == 'EdgeGraph.AdjacencyMap.Internal.empty'
-- edges [x] == 'EdgeGraph.AdjacencyMap.Internal.edge' x
-- @
edges :: Ord a => [a] -> AdjacencyMap a
edges = fromIncidence . I.edges

-- | Construct a graph from a list of nodes. The nodes are normalized
-- (merged where they share labels).
fromNodeList :: Ord a => [I.Node a] -> AdjacencyMap a
fromNodeList = fromIncidence . I.fromNodeList

-- | Construct a graph from a list of (tips, pits) pairs, where each pair
-- represents a node with its incoming edges (tips) and outgoing edge
-- labels (pits). The resulting graph is normalized (nodes sharing labels
-- are merged).
--
-- @
-- fromIncidenceList []                  == 'EdgeGraph.AdjacencyMap.Internal.empty'
-- fromIncidenceList [([],[x]),([x],[])] == 'EdgeGraph.AdjacencyMap.Internal.edge' x
-- @
fromIncidenceList :: Ord a => [([a], [a])] -> AdjacencyMap a
fromIncidenceList = fromIncidence . I.fromIncidenceList

-- | The sorted list of nodes of a graph.
nodeList :: Ord a => AdjacencyMap a -> [I.Node a]
nodeList = I.nodeList . toIncidence

-- | The set of nodes of a graph.
nodeSet :: Ord a => AdjacencyMap a -> Set (I.Node a)
nodeSet = I.nodeSet . toIncidence

-- | The number of nodes in a graph.
nodeCount :: Ord a => AdjacencyMap a -> Int
nodeCount = I.nodeCount . toIncidence

-- | Check if a graph is empty.
-- Complexity: /O(1)/ time.
isEmpty :: AdjacencyMap a -> Bool
isEmpty (AdjacencyMap m) = Map.null m

-- | The set of all distinct edges.
edgeSet :: AdjacencyMap a -> Set a
edgeSet (AdjacencyMap m) = Map.keysSet m

-- | The sorted list of all distinct edges.
edgeList :: AdjacencyMap a -> [a]
edgeList (AdjacencyMap m) = Map.keys m

-- | The sorted /adjacency list/ of a graph. Each entry is an edge
-- paired with its t'Adjacency' record.
-- Complexity: /O(n)/ time and memory.
--
-- @
-- adjacencyList 'EdgeGraph.AdjacencyMap.Internal.empty'    == []
-- adjacencyList ('EdgeGraph.AdjacencyMap.Internal.edge' x) == [(x, Adjacency ('Data.Set.singleton' x) ('Data.Set.singleton' x) 'Data.Set.empty' 'Data.Set.empty')]
-- @
adjacencyList :: AdjacencyMap a -> [(a, Adjacency a)]
adjacencyList (AdjacencyMap m) = Map.toAscList m

-- | The number of distinct edges.
edgeCount :: AdjacencyMap a -> Int
edgeCount (AdjacencyMap m) = Map.size m

-- | Check if a graph contains a given edge.
-- Complexity: /O(log n)/ time.
hasEdge :: Ord a => a -> AdjacencyMap a -> Bool
hasEdge a (AdjacencyMap m) = Map.member a m

-- | Remove an edge from the graph, updating all neighbour references.
-- Complexity: /O((|forks| + |joins| + |preds| + |succs|) * log n)/ time.
--
-- @
-- removeEdge x ('EdgeGraph.AdjacencyMap.Internal.edge' x) == 'EdgeGraph.AdjacencyMap.Internal.empty'
-- @
removeEdge :: Ord a => a -> AdjacencyMap a -> AdjacencyMap a
removeEdge x (AdjacencyMap m) = case Map.lookup x m of
    Nothing  -> AdjacencyMap m
    Just adj ->
      let newForks = Set.delete x (forks adj)
          newJoins = Set.delete x (joins adj)
          m1 = Map.delete x m
          m2 = Set.foldl' (\acc b ->
            Map.adjust (\r -> r { forks = newForks }) b acc)
            m1 newForks
          m3 = Set.foldl' (\acc b ->
            Map.adjust (\r -> r { joins = newJoins }) b acc)
            m2 newJoins
          m4 = Set.foldl' (\acc b ->
            Map.adjust (\r -> r { succs = newForks }) b acc)
            m3 (preds adj)
          m5 = Set.foldl' (\acc b ->
            Map.adjust (\r -> r { preds = newJoins }) b acc)
            m4 (succs adj)
      in AdjacencyMap m5

-- | Detach an edge from its source node. The edge gets a fresh source
-- while any other edges sharing the original source node remain together.
-- Complexity: /O((|forks| + |preds|) * log n)/ time.
--
-- @
-- detachPit x ('EdgeGraph.AdjacencyMap.Internal.edge' x)                 == 'EdgeGraph.AdjacencyMap.Internal.edge' x
-- detachPit 2 ('into' ('EdgeGraph.AdjacencyMap.Internal.edge' 1) ('EdgeGraph.AdjacencyMap.Internal.edge' 2)) == 'edges' [1, 2]
-- detachPit 1 ('pits' ('EdgeGraph.AdjacencyMap.Internal.edge' 1) ('EdgeGraph.AdjacencyMap.Internal.edge' 2)) == 'edges' [1, 2]
-- @
detachPit :: Ord a => a -> AdjacencyMap a -> AdjacencyMap a
detachPit a (AdjacencyMap m) = case Map.lookup a m of
  Nothing  -> AdjacencyMap m
  Just adj ->
    let oldForks = forks adj
        oldPreds = preds adj
        m1 = Map.adjust (\r -> r { forks = Set.singleton a
                                 , preds = Set.empty }) a m
        m2 = Set.foldl' (\acc b ->
          Map.adjust (\r -> r { forks = Set.delete a (forks r) }) b acc)
          m1 (Set.delete a oldForks)
        m3 = Set.foldl' (\acc b ->
          Map.adjust (\r -> r { succs = Set.delete a (succs r) }) b acc)
          m2 oldPreds
    in AdjacencyMap m3

-- | Detach an edge from its sink node. The edge gets a fresh sink
-- while any other edges sharing the original sink node remain together.
-- Complexity: /O((|joins| + |succs|) * log n)/ time.
--
-- @
-- detachTip x ('EdgeGraph.AdjacencyMap.Internal.edge' x)                 == 'EdgeGraph.AdjacencyMap.Internal.edge' x
-- detachTip 1 ('into' ('EdgeGraph.AdjacencyMap.Internal.edge' 1) ('EdgeGraph.AdjacencyMap.Internal.edge' 2)) == 'edges' [1, 2]
-- detachTip 1 ('tips' ('EdgeGraph.AdjacencyMap.Internal.edge' 1) ('EdgeGraph.AdjacencyMap.Internal.edge' 2)) == 'edges' [1, 2]
-- @
detachTip :: Ord a => a -> AdjacencyMap a -> AdjacencyMap a
detachTip a (AdjacencyMap m) = case Map.lookup a m of
  Nothing  -> AdjacencyMap m
  Just adj ->
    let oldJoins = joins adj
        oldSuccs = succs adj
        m1 = Map.adjust (\r -> r { joins = Set.singleton a
                                 , succs = Set.empty }) a m
        m2 = Set.foldl' (\acc b ->
          Map.adjust (\r -> r { joins = Set.delete a (joins r) }) b acc)
          m1 (Set.delete a oldJoins)
        m3 = Set.foldl' (\acc b ->
          Map.adjust (\r -> r { preds = Set.delete a (preds r) }) b acc)
          m2 oldSuccs
    in AdjacencyMap m3

-- | Transform a graph by applying a function to each edge.
--
-- @
-- gmap f 'EdgeGraph.AdjacencyMap.Internal.empty'    == 'EdgeGraph.AdjacencyMap.Internal.empty'
-- gmap f ('EdgeGraph.AdjacencyMap.Internal.edge' x) == 'EdgeGraph.AdjacencyMap.Internal.edge' (f x)
-- gmap id         == id
-- gmap f . gmap g == gmap (f . g)
-- @
gmap :: (Ord a, Ord b) => (a -> b) -> AdjacencyMap a -> AdjacencyMap b
gmap f = fromIncidence . I.gmap f . toIncidence

-- | Construct the /induced subgraph/ of a given graph by removing edges
-- that do not satisfy a given predicate. Keeps only edges where @p@ returns
-- 'True', and updates all neighbour sets accordingly.
-- Complexity: /O(n * m)/ time.
--
-- @
-- induce (const True)  x == x
-- induce (const False) x == 'EdgeGraph.AdjacencyMap.Internal.empty'
-- @
induce :: Ord a => (a -> Bool) -> AdjacencyMap a -> AdjacencyMap a
induce p (AdjacencyMap m) = AdjacencyMap $ Map.map updateAdj kept
  where
    kept     = Map.filterWithKey (\k _ -> p k) m
    keptKeys = Map.keysSet kept
    updateAdj adj = Adjacency
      { forks = Set.intersection (forks adj) keptKeys
      , joins = Set.intersection (joins adj) keptKeys
      , preds = Set.intersection (preds adj) keptKeys
      , succs = Set.intersection (succs adj) keptKeys
      }