{- HLINT ignore "Use second" -}
module Mischief.ECS.Graph where

import Control.Monad
import Data.Foldable
import Data.IORef
import Data.List qualified as List
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe
import Data.Set (Set)
import Data.Set qualified as Set
import Mischief.ECS.Vec (IOVec)
import Mischief.ECS.Vec qualified as Vec

data Graph a = Graph {forall a. Graph a -> IORef (Map a Int)
nodes :: IORef (Map a Int), forall a. Graph a -> IORef Int
counter :: IORef Int, forall a. Graph a -> IOVec (Maybe a, Set Int)
edges :: IOVec (Maybe a, Set Int)}

newGraph :: forall a. IO (Graph a)
newGraph :: forall a. IO (Graph a)
newGraph = do
  nodes <- Map a Int -> IO (IORef (Map a Int))
forall a. a -> IO (IORef a)
newIORef Map a Int
forall k a. Map k a
Map.empty
  counter <- newIORef 0
  edges <- Vec.new 16

  return Graph {nodes, counter, edges}

addNode :: (Ord a) => a -> Graph a -> IO Int
addNode :: forall a. Ord a => a -> Graph a -> IO Int
addNode a
node Graph {IORef (Map a Int)
nodes :: forall a. Graph a -> IORef (Map a Int)
nodes :: IORef (Map a Int)
nodes, IORef Int
counter :: forall a. Graph a -> IORef Int
counter :: IORef Int
counter, IOVec (Maybe a, Set Int)
edges :: forall a. Graph a -> IOVec (Maybe a, Set Int)
edges :: IOVec (Maybe a, Set Int)
edges} = do
  maybeNode <- a -> Graph a -> IO (Maybe Int)
forall a. Ord a => a -> Graph a -> IO (Maybe Int)
tryGetNode a
node Graph {IORef (Map a Int)
nodes :: IORef (Map a Int)
nodes :: IORef (Map a Int)
nodes, IORef Int
counter :: IORef Int
counter :: IORef Int
counter, IOVec (Maybe a, Set Int)
edges :: IOVec (Maybe a, Set Int)
edges :: IOVec (Maybe a, Set Int)
edges}
  case maybeNode of
    Just Int
x -> do
      a -> Int -> Graph a -> IO ()
forall a. a -> Int -> Graph a -> IO ()
setNode a
node Int
x Graph {IORef (Map a Int)
nodes :: IORef (Map a Int)
nodes :: IORef (Map a Int)
nodes, IORef Int
counter :: IORef Int
counter :: IORef Int
counter, IOVec (Maybe a, Set Int)
edges :: IOVec (Maybe a, Set Int)
edges :: IOVec (Maybe a, Set Int)
edges}
      Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
x
    Maybe Int
Nothing -> do
      index <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
counter
      modifyIORef' counter (+ 1)
      Vec.pushBack edges (Just node, Set.empty)
      modifyIORef' nodes (Map.insert node index)
      return index

setNode :: a -> Int -> Graph a -> IO ()
setNode :: forall a. a -> Int -> Graph a -> IO ()
setNode a
node Int
index Graph {IOVec (Maybe a, Set Int)
edges :: forall a. Graph a -> IOVec (Maybe a, Set Int)
edges :: IOVec (Maybe a, Set Int)
edges} = do
  Vec (PrimState IO) (Maybe a, Set Int)
-> Int -> ((Maybe a, Set Int) -> (Maybe a, Set Int)) -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> (a -> a) -> m ()
Vec.modify_ IOVec (Maybe a, Set Int)
Vec (PrimState IO) (Maybe a, Set Int)
edges Int
index (\(Maybe a
_, Set Int
x) -> (a -> Maybe a
forall a. a -> Maybe a
Just a
node, Set Int
x))

tryGetNode :: (Ord a) => a -> Graph a -> IO (Maybe Int)
tryGetNode :: forall a. Ord a => a -> Graph a -> IO (Maybe Int)
tryGetNode a
node Graph a
graph = do
  nodes <- IORef (Map a Int) -> IO (Map a Int)
forall a. IORef a -> IO a
readIORef Graph a
graph.nodes
  return $ Map.lookup node nodes

-- data GetOrAddResult = AddedNode Int | GotNode Int

getOrAddNode :: (Ord a) => a -> Graph a -> IO Int
getOrAddNode :: forall a. Ord a => a -> Graph a -> IO Int
getOrAddNode a
node Graph a
graph = do
  nodes' <- IORef (Map a Int) -> IO (Map a Int)
forall a. IORef a -> IO a
readIORef Graph a
graph.nodes
  case Map.lookup node nodes' of
    Just Int
x -> Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
x
    Maybe Int
Nothing -> do
      index <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef Graph a
graph.counter
      modifyIORef' graph.counter (+ 1)
      Vec.pushBack graph.edges (Nothing, Set.empty)
      modifyIORef' graph.nodes (Map.insert node index)
      return index

addEdge :: (Ord a) => (a, a) -> Graph a -> IO ()
addEdge :: forall a. Ord a => (a, a) -> Graph a -> IO ()
addEdge (a
a, a
b) Graph a
graph = do
  a' <- a -> Graph a -> IO Int
forall a. Ord a => a -> Graph a -> IO Int
getOrAddNode a
a Graph a
graph
  b' <- getOrAddNode b graph
  Vec.modify_ graph.edges b' (\(Maybe a
x, Set Int
l) -> (Maybe a
x, Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
Set.insert Int
a' Set Int
l))

takeRemoveableNodes :: IOVec (Maybe a, Set Int) -> IO [a]
takeRemoveableNodes :: forall a. IOVec (Maybe a, Set Int) -> IO [a]
takeRemoveableNodes IOVec (Maybe a, Set Int)
edges = do
  edgeList <- Vec (PrimState IO) (Maybe a, Set Int) -> IO [(Maybe a, Set Int)]
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m [a]
Vec.toList IOVec (Maybe a, Set Int)
Vec (PrimState IO) (Maybe a, Set Int)
edges
  len <- Vec.length edges
  let nodes = (((Maybe a, Set Int), Int) -> [(a, Int)])
-> [((Maybe a, Set Int), Int)] -> [(a, Int)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Maybe a, Set Int), Int) -> [(a, Int)]
forall {a} {b} {b}. ((Maybe a, b), b) -> [(a, b)]
unwrap ((((Maybe a, Set Int), Int) -> Bool)
-> [((Maybe a, Set Int), Int)] -> [((Maybe a, Set Int), Int)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\((Maybe a
a, Set Int
x), Int
_) -> Set Int -> Bool
forall a. Set a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Set Int
x Bool -> Bool -> Bool
&& Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
a) ([(Maybe a, Set Int)] -> [Int] -> [((Maybe a, Set Int), Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Maybe a, Set Int)]
edgeList [Int
0 :: Int ..]))
  for_ nodes $ \(a
_, Int
i) -> Vec (PrimState IO) (Maybe a, Set Int)
-> Int -> ((Maybe a, Set Int) -> (Maybe a, Set Int)) -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> (a -> a) -> m ()
Vec.modify_ IOVec (Maybe a, Set Int)
Vec (PrimState IO) (Maybe a, Set Int)
edges Int
i (\(Maybe a
_, Set Int
l) -> (Maybe a
forall a. Maybe a
Nothing, Set Int
l))
  for_ nodes $ \(a
_, Int
i) -> [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> Vec (PrimState IO) (Maybe a, Set Int)
-> Int -> ((Maybe a, Set Int) -> (Maybe a, Set Int)) -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> (a -> a) -> m ()
Vec.modify_ IOVec (Maybe a, Set Int)
Vec (PrimState IO) (Maybe a, Set Int)
edges Int
j (\(Maybe a
a, Set Int
l) -> (Maybe a
a, Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
Set.delete Int
i Set Int
l))
  return (map fst nodes)
  where
    unwrap :: ((Maybe a, b), b) -> [(a, b)]
unwrap ((Just a
a, b
_), b
i) = [(a
a, b
i)]
    unwrap ((Maybe a, b), b)
_ = []

getNodes :: Graph a -> IO [[a]]
getNodes :: forall a. Graph a -> IO [[a]]
getNodes Graph {IOVec (Maybe a, Set Int)
edges :: forall a. Graph a -> IOVec (Maybe a, Set Int)
edges :: IOVec (Maybe a, Set Int)
edges} = do
  edges' <- Vec (PrimState IO) (Maybe a, Set Int)
-> IO (Vec (PrimState IO) (Maybe a, Set Int))
forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> m (Vec (PrimState m) a)
Vec.clone IOVec (Maybe a, Set Int)
Vec (PrimState IO) (Maybe a, Set Int)
edges
  res <- newIORef []
  step res edges'
  readIORef res
  where
    step :: IORef [[a]] -> IOVec (Maybe a, Set Int) -> IO ()
step IORef [[a]]
res IOVec (Maybe a, Set Int)
edges' = do
      nodes <- IOVec (Maybe a, Set Int) -> IO [a]
forall a. IOVec (Maybe a, Set Int) -> IO [a]
takeRemoveableNodes IOVec (Maybe a, Set Int)
edges'
      unless (null nodes) $ do
        modifyIORef' res (++ [nodes])
        step res edges'