module LOAG.Graphs where

import Control.Monad (forM, forM_)
import Control.Monad.ST
import Control.Monad.State
import CommonTypes
import Data.STRef
import Data.Maybe (catMaybes, isNothing, fromJust)
import Data.Tuple (swap)
import qualified Data.Array as A
import           Data.Array.IArray (amap)
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import qualified Data.Map as M
import qualified Data.Set as S
import Data.Array.MArray (mapArray)
import Data.Array.ST

type Schedule   = (A.Array Vertex (Maybe Int), A.Array Int [Vertex])
type Vertex     = Int
type Cycle      = IS.IntSet
type Vertices   = IS.IntSet
type Edge       = (Vertex, Vertex)
type Edges      = S.Set Edge
-- Maps that are suitable for Graphs (from 1 node to a set of nodes)

type Graph s    = (DirGraphRef s, DirGraphRef s)
-- | Frozen version of a graph

type FrGraph    = (DirGraph, DirGraph)
type DirGraph   = A.Array Vertex Vertices
type DirGraphRef s = STArray s Vertex Vertices

-- |----------------------------------------------------------------------

-- | Functions for changing the state within AOAG

-- |  possibly catching errors from creating cycles


addEDs :: Graph s -> [Edge] -> (ST s) (Maybe (Edge, Cycle))
addEDs :: Graph s -> [Edge] -> ST s (Maybe (Edge, Cycle))
addEDs Graph s
_ [] = Maybe (Edge, Cycle) -> ST s (Maybe (Edge, Cycle))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Edge, Cycle)
forall a. Maybe a
Nothing
addEDs Graph s
edp (Edge
e:[Edge]
es) = do
    Either Cycle [Edge]
res <- Edge
e Edge -> Graph s -> ST s (Either Cycle [Edge])
forall s. Edge -> Graph s -> ST s (Either Cycle [Edge])
`inserT` Graph s
edp
    case Either Cycle [Edge]
res of
        Right [Edge]
_ -> Graph s -> [Edge] -> ST s (Maybe (Edge, Cycle))
forall s. Graph s -> [Edge] -> ST s (Maybe (Edge, Cycle))
addEDs Graph s
edp [Edge]
es
        Left Cycle
c  -> Maybe (Edge, Cycle) -> ST s (Maybe (Edge, Cycle))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Edge, Cycle) -> ST s (Maybe (Edge, Cycle)))
-> Maybe (Edge, Cycle) -> ST s (Maybe (Edge, Cycle))
forall a b. (a -> b) -> a -> b
$ (Edge, Cycle) -> Maybe (Edge, Cycle)
forall a. a -> Maybe a
Just (Edge
e,Cycle
c)
       
-- | Draws an edge from one node to another, by adding the latter to the

--    node set of the first

insErt :: Edge -> Graph s -> (ST s) ()
insErt :: Edge -> Graph s -> ST s ()
insErt (Vertex
f, Vertex
t) g :: Graph s
g@(DirGraphRef s
ft,DirGraphRef s
tf) = do 
    Cycle
ts <- DirGraphRef s -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
ft Vertex
f
    Cycle
fs <- DirGraphRef s -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
tf Vertex
t
    DirGraphRef s -> Vertex -> Cycle -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray DirGraphRef s
ft Vertex
f (Vertex
t Vertex -> Cycle -> Cycle
`IS.insert` Cycle
ts)
    DirGraphRef s -> Vertex -> Cycle -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray DirGraphRef s
tf Vertex
t (Vertex
f Vertex -> Cycle -> Cycle
`IS.insert` Cycle
fs)

removE :: Edge -> Graph s -> (ST s) ()
removE :: Edge -> Graph s -> ST s ()
removE e :: Edge
e@(Vertex
f,Vertex
t) g :: Graph s
g@(DirGraphRef s
ft,DirGraphRef s
tf) = do 
    Cycle
ts <- DirGraphRef s -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
ft Vertex
f
    Cycle
fs <- DirGraphRef s -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
tf Vertex
t
    DirGraphRef s -> Vertex -> Cycle -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray DirGraphRef s
ft Vertex
f (Vertex
t Vertex -> Cycle -> Cycle
`IS.delete` Cycle
ts)
    DirGraphRef s -> Vertex -> Cycle -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray DirGraphRef s
tf Vertex
t (Vertex
f Vertex -> Cycle -> Cycle
`IS.delete` Cycle
fs)


-- | Revert an edge in the graph

revErt :: Edge -> Graph s -> (ST s) ()
revErt :: Edge -> Graph s -> ST s ()
revErt Edge
e Graph s
g = do
    Bool
present <- Edge -> Graph s -> ST s Bool
forall s. Edge -> Graph s -> ST s Bool
member Edge
e Graph s
g
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
present (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ Edge -> Graph s -> ST s ()
forall s. Edge -> Graph s -> ST s ()
removE Edge
e Graph s
g ST s () -> ST s () -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Edge -> Graph s -> ST s ()
forall s. Edge -> Graph s -> ST s ()
insErt (Edge -> Edge
forall a b. (a, b) -> (b, a)
swap Edge
e) Graph s
g

-- | Assuming the given graph is already transitively closed, and

-- |    not cyclic, insert an 

-- |    edge such that the graph maintains transitively closed.

-- |    returns the cycle if this results in a cycle or returns a pair

-- |    (graph, edges) if not. Where graph is the new Graph and 

-- |    edges represent the edges that were required for transitively

-- |    closing the graph.

inserT :: Edge -> Graph s -> (ST s) (Either Cycle [Edge])
inserT :: Edge -> Graph s -> ST s (Either Cycle [Edge])
inserT e :: Edge
e@(Vertex
f, Vertex
t) g :: Graph s
g@(DirGraphRef s
gft,DirGraphRef s
gtf)
    | Vertex
f Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
t     = Either Cycle [Edge] -> ST s (Either Cycle [Edge])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Cycle [Edge] -> ST s (Either Cycle [Edge]))
-> Either Cycle [Edge] -> ST s (Either Cycle [Edge])
forall a b. (a -> b) -> a -> b
$ Cycle -> Either Cycle [Edge]
forall a b. a -> Either a b
Left (Cycle -> Either Cycle [Edge]) -> Cycle -> Either Cycle [Edge]
forall a b. (a -> b) -> a -> b
$ Vertex -> Cycle
IS.singleton Vertex
f
    | Bool
otherwise  = do
        Bool
present <- Edge -> Graph s -> ST s Bool
forall s. Edge -> Graph s -> ST s Bool
member Edge
e Graph s
g
        if Bool
present 
         then (Either Cycle [Edge] -> ST s (Either Cycle [Edge])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Cycle [Edge] -> ST s (Either Cycle [Edge]))
-> Either Cycle [Edge] -> ST s (Either Cycle [Edge])
forall a b. (a -> b) -> a -> b
$ [Edge] -> Either Cycle [Edge]
forall a b. b -> Either a b
Right [])
         else do
          Cycle
rs <- DirGraphRef s -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gtf Vertex
f
          Cycle
us <- DirGraphRef s -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gft Vertex
t
          Cycle
pointsToF <- DirGraphRef s -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gtf Vertex
f
          Cycle
pointsToT <- DirGraphRef s -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gtf Vertex
t
          Cycle
tPointsTo <- DirGraphRef s -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gft Vertex
t
          let new2t :: Cycle
new2t = Cycle
pointsToF Cycle -> Cycle -> Cycle
IS.\\ Cycle
pointsToT
          -- extras from f connects all new nodes pointing to f with t

          let extraF :: [Edge]
extraF = ([Edge] -> Vertex -> [Edge]) -> [Edge] -> Cycle -> [Edge]
forall a. (a -> Vertex -> a) -> a -> Cycle -> a
IS.foldl' (\[Edge]
acc Vertex
tf -> (Vertex
tf,Vertex
t) Edge -> [Edge] -> [Edge]
forall a. a -> [a] -> [a]
: [Edge]
acc) [] Cycle
new2t
          -- extras of t connects all nodes that will be pointing to t

          -- in the new graph, with all the nodes t points to in the

          -- current graph

          STRef s [Edge]
all2tPointsTo <- [Edge] -> ST s (STRef s [Edge])
forall a s. a -> ST s (STRef s a)
newSTRef []
          [Vertex] -> (Vertex -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Cycle -> [Vertex]
IS.toList Cycle
tPointsTo) ((Vertex -> ST s ()) -> ST s ()) -> (Vertex -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Vertex
ft -> do
            [Edge]
current  <- STRef s [Edge] -> ST s [Edge]
forall s a. STRef s a -> ST s a
readSTRef STRef s [Edge]
all2tPointsTo
            Cycle
existing <- DirGraphRef s -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gtf Vertex
ft
            let new4ft :: [Edge]
new4ft = (Vertex -> Edge) -> [Vertex] -> [Edge]
forall a b. (a -> b) -> [a] -> [b]
map ((Vertex -> Vertex -> Edge) -> Vertex -> Vertex -> Edge
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) Vertex
ft) ([Vertex] -> [Edge]) -> [Vertex] -> [Edge]
forall a b. (a -> b) -> a -> b
$ Cycle -> [Vertex]
IS.toList (Cycle -> [Vertex]) -> Cycle -> [Vertex]
forall a b. (a -> b) -> a -> b
$ 
                            -- removing existing here matters a lot

                            (Vertex
f Vertex -> Cycle -> Cycle
`IS.insert` Cycle
pointsToF) Cycle -> Cycle -> Cycle
IS.\\ Cycle
existing
            STRef s [Edge] -> [Edge] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [Edge]
all2tPointsTo ([Edge] -> ST s ()) -> [Edge] -> ST s ()
forall a b. (a -> b) -> a -> b
$ [Edge]
current [Edge] -> [Edge] -> [Edge]
forall a. [a] -> [a] -> [a]
++ [Edge]
new4ft
                  
          [Edge]
extraT <- STRef s [Edge] -> ST s [Edge]
forall s a. STRef s a -> ST s a
readSTRef STRef s [Edge]
all2tPointsTo
        -- the extras consists of extras from f and extras from t

        -- both these extra sets dont contain edges if they are already 

        -- present in the old graph

          let extra :: [Edge]
extra  = [Edge]
extraF [Edge] -> [Edge] -> [Edge]
forall a. [a] -> [a] -> [a]
++ [Edge]
extraT
          (Edge -> ST s ()) -> [Edge] -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Edge -> Graph s -> ST s ()
forall s. Edge -> Graph s -> ST s ()
`insErt` Graph s
g) (Edge
e Edge -> [Edge] -> [Edge]
forall a. a -> [a] -> [a]
: [Edge]
extra) 
        -- the new graph contains a cycle if there is a self-edge

        -- this cycle will contain both f and t

          Bool
cyclic <- Edge -> Graph s -> ST s Bool
forall s. Edge -> Graph s -> ST s Bool
member (Vertex
f,Vertex
f) Graph s
g
          if Bool
cyclic
           then do
            Cycle
cycle <- DirGraphRef s -> ST s Cycle
forall s. STArray s Vertex Cycle -> ST s Cycle
getCycle DirGraphRef s
gft
            Either Cycle [Edge] -> ST s (Either Cycle [Edge])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Cycle [Edge] -> ST s (Either Cycle [Edge]))
-> Either Cycle [Edge] -> ST s (Either Cycle [Edge])
forall a b. (a -> b) -> a -> b
$ Cycle -> Either Cycle [Edge]
forall a b. a -> Either a b
Left Cycle
cycle
           else Either Cycle [Edge] -> ST s (Either Cycle [Edge])
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Cycle [Edge] -> ST s (Either Cycle [Edge]))
-> Either Cycle [Edge] -> ST s (Either Cycle [Edge])
forall a b. (a -> b) -> a -> b
$ [Edge] -> Either Cycle [Edge]
forall a b. b -> Either a b
Right [Edge]
extra
        
       where
        -- given that there is a cycle,all elements of this cycle are being

        -- pointed at by f. However, not all elements that f points to are 

        -- part of the cycle. Only those that point back to f.

        getCycle :: STArray s Vertex Vertices -> (ST s) Cycle
        getCycle :: STArray s Vertex Cycle -> ST s Cycle
getCycle STArray s Vertex Cycle
gft = do
            Cycle
ts <- STArray s Vertex Cycle -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Vertex Cycle
gft Vertex
f
            [Maybe Vertex]
mnodes <- [Vertex] -> (Vertex -> ST s (Maybe Vertex)) -> ST s [Maybe Vertex]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Cycle -> [Vertex]
IS.toList Cycle
ts) ((Vertex -> ST s (Maybe Vertex)) -> ST s [Maybe Vertex])
-> (Vertex -> ST s (Maybe Vertex)) -> ST s [Maybe Vertex]
forall a b. (a -> b) -> a -> b
$ \Vertex
t' -> do
                Cycle
fs' <- STArray s Vertex Cycle -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Vertex Cycle
gft Vertex
t'
                if Vertex
f Vertex -> Cycle -> Bool
`IS.member` Cycle
fs'
                 then Maybe Vertex -> ST s (Maybe Vertex)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Vertex -> ST s (Maybe Vertex))
-> Maybe Vertex -> ST s (Maybe Vertex)
forall a b. (a -> b) -> a -> b
$ Vertex -> Maybe Vertex
forall a. a -> Maybe a
Just Vertex
t'
                 else Maybe Vertex -> ST s (Maybe Vertex)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Vertex -> ST s (Maybe Vertex))
-> Maybe Vertex -> ST s (Maybe Vertex)
forall a b. (a -> b) -> a -> b
$ Maybe Vertex
forall a. Maybe a
Nothing
            Cycle -> ST s Cycle
forall (m :: * -> *) a. Monad m => a -> m a
return (Cycle -> ST s Cycle) -> Cycle -> ST s Cycle
forall a b. (a -> b) -> a -> b
$ [Vertex] -> Cycle
IS.fromList ([Vertex] -> Cycle) -> [Vertex] -> Cycle
forall a b. (a -> b) -> a -> b
$ [Maybe Vertex] -> [Vertex]
forall a. [Maybe a] -> [a]
catMaybes [Maybe Vertex]
mnodes

-- | Check if a certain edge is part of a graph which means that,

-- |  the receiving node must be in the node set of the sending

member :: Edge -> Graph s -> (ST s) Bool
member :: Edge -> Graph s -> ST s Bool
member (Vertex
f, Vertex
t) (DirGraphRef s
ft, DirGraphRef s
tf) = do
    Cycle
ts <- DirGraphRef s -> Vertex -> ST s Cycle
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
ft Vertex
f
    Bool -> ST s Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> ST s Bool) -> Bool -> ST s Bool
forall a b. (a -> b) -> a -> b
$ Vertex -> Cycle -> Bool
IS.member Vertex
t Cycle
ts

-- | Check whether an edge is part of a frozen graph

fr_member :: FrGraph -> Edge -> Bool
fr_member :: FrGraph -> Edge -> Bool
fr_member (DirGraph
ft, DirGraph
tf) (Vertex
f, Vertex
t) = Vertex -> Cycle -> Bool
IS.member Vertex
t (DirGraph
ft DirGraph -> Vertex -> Cycle
forall i e. Ix i => Array i e -> i -> e
A.! Vertex
f)

-- | Flatten a graph, meaning that we transform this graph to 

-- |  a set of Edges by combining a sending node with all the

-- |  receiving nodes in its node set

flatten :: Graph s -> (ST s) Edges 
flatten :: Graph s -> ST s Edges
flatten (DirGraphRef s
gft, DirGraphRef s
_) = do
    [(Vertex, Cycle)]
list <- DirGraphRef s -> ST s [(Vertex, Cycle)]
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m [(i, e)]
getAssocs DirGraphRef s
gft
    Edges -> ST s Edges
forall (m :: * -> *) a. Monad m => a -> m a
return (Edges -> ST s Edges) -> Edges -> ST s Edges
forall a b. (a -> b) -> a -> b
$ [Edge] -> Edges
forall a. Ord a => [a] -> Set a
S.fromList ([Edge] -> Edges) -> [Edge] -> Edges
forall a b. (a -> b) -> a -> b
$ ((Vertex, Cycle) -> [Edge]) -> [(Vertex, Cycle)] -> [Edge]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap 
                (\(Vertex
f, Cycle
ts) -> (Vertex -> Edge) -> [Vertex] -> [Edge]
forall a b. (a -> b) -> [a] -> [b]
map ((,) Vertex
f) ([Vertex] -> [Edge]) -> [Vertex] -> [Edge]
forall a b. (a -> b) -> a -> b
$ Cycle -> [Vertex]
IS.toList Cycle
ts) [(Vertex, Cycle)]
list

freeze_graph :: Graph s -> (ST s) FrGraph
freeze_graph :: Graph s -> ST s FrGraph
freeze_graph (DirGraphRef s
mf, DirGraphRef s
mt) = do
    DirGraph
fr_f <- DirGraphRef s -> ST s DirGraph
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
freeze DirGraphRef s
mf
    DirGraph
fr_t <- DirGraphRef s -> ST s DirGraph
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
freeze DirGraphRef s
mt
    FrGraph -> ST s FrGraph
forall (m :: * -> *) a. Monad m => a -> m a
return (DirGraph
fr_f, DirGraph
fr_t)