{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Graph.Base
-- Copyright   :  (c) Masahiro Sakai 2020
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable
--
-----------------------------------------------------------------------------
module ToySolver.Graph.Base
  ( EdgeLabeledGraph
  , Graph
  , graphToUnorderedEdges
  , graphFromUnorderedEdges
  , graphFromUnorderedEdgesWith
  , isIndependentSet
  ) where

import Control.Monad
import Data.Array.IArray
import Data.Array.ST
import qualified Data.IntMap.Lazy as IntMap
import Data.IntMap.Lazy (IntMap)
import qualified Data.IntSet as IntSet
import Data.IntSet (IntSet)

type EdgeLabeledGraph a = Array Int (IntMap a)

type Graph = EdgeLabeledGraph ()

graphToUnorderedEdges :: EdgeLabeledGraph a -> [(Int, Int, a)]
graphToUnorderedEdges :: forall a. EdgeLabeledGraph a -> [(Int, Int, a)]
graphToUnorderedEdges EdgeLabeledGraph a
g = do
  (Int
node1, IntMap a
nodes) <- forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs EdgeLabeledGraph a
g
  (Int
node2, a
a) <- forall a. IntMap a -> [(Int, a)]
IntMap.toList forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. Int -> IntMap a -> (IntMap a, IntMap a)
IntMap.split Int
node1 IntMap a
nodes
  forall (m :: * -> *) a. Monad m => a -> m a
return (Int
node1, Int
node2, a
a)

graphFromUnorderedEdges :: Int -> [(Int, Int, a)] -> EdgeLabeledGraph a
graphFromUnorderedEdges :: forall a. Int -> [(Int, Int, a)] -> EdgeLabeledGraph a
graphFromUnorderedEdges = forall a.
(a -> a -> a) -> Int -> [(Int, Int, a)] -> EdgeLabeledGraph a
graphFromUnorderedEdgesWith forall a b. a -> b -> a
const

graphFromUnorderedEdgesWith :: (a -> a -> a) -> Int -> [(Int, Int, a)] -> EdgeLabeledGraph a
graphFromUnorderedEdgesWith :: forall a.
(a -> a -> a) -> Int -> [(Int, Int, a)] -> EdgeLabeledGraph a
graphFromUnorderedEdgesWith a -> a -> a
f Int
n [(Int, Int, a)]
es = forall i e. (forall s. ST s (STArray s i e)) -> Array i e
runSTArray forall a b. (a -> b) -> a -> b
$ do
  STArray s Int (IntMap a)
a <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
0, Int
nforall a. Num a => a -> a -> a
-Int
1) forall a. IntMap a
IntMap.empty
  let ins :: Int -> Int -> a -> m ()
ins Int
i Int
x a
l = do
        IntMap a
m <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Int (IntMap a)
a Int
i
        forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STArray s Int (IntMap a)
a Int
i forall a b. (a -> b) -> a -> b
$! forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IntMap.insertWith a -> a -> a
f Int
x a
l IntMap a
m
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Int, Int, a)]
es forall a b. (a -> b) -> a -> b
$ \(Int
node1, Int
node2, a
a) -> do
    forall {m :: * -> *}.
MArray (STArray s) (IntMap a) m =>
Int -> Int -> a -> m ()
ins Int
node1 Int
node2 a
a
    forall {m :: * -> *}.
MArray (STArray s) (IntMap a) m =>
Int -> Int -> a -> m ()
ins Int
node2 Int
node1 a
a
  forall (m :: * -> *) a. Monad m => a -> m a
return STArray s Int (IntMap a)
a

isIndependentSet :: EdgeLabeledGraph a -> IntSet -> Bool
isIndependentSet :: forall a. EdgeLabeledGraph a -> IntSet -> Bool
isIndependentSet EdgeLabeledGraph a
g IntSet
s = forall (t :: * -> *) a. Foldable t => t a -> Bool
null forall a b. (a -> b) -> a -> b
$ do
  (Int
node1, Int
node2, a
_) <- forall a. EdgeLabeledGraph a -> [(Int, Int, a)]
graphToUnorderedEdges EdgeLabeledGraph a
g
  forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ Int
node1 Int -> IntSet -> Bool
`IntSet.member` IntSet
s
  forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ Int
node2 Int -> IntSet -> Bool
`IntSet.member` IntSet
s
  forall (m :: * -> *) a. Monad m => a -> m a
return ()