{-# LANGUAGE TypeFamilies #-}

module Data.Graph.Dense where

import Control.Monad.ST
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

data DenseGraph a = DenseGraph
  { forall a. DenseGraph a -> Int
numVerticesDG :: !Int
  , forall a. DenseGraph a -> Vector a
adjacentDG :: !(U.Vector a)
  }

instance (U.Unbox a, Show a) => Show (DenseGraph a) where
  show :: DenseGraph a -> String
show = [[a]] -> String
forall a. Show a => a -> String
show ([[a]] -> String)
-> (DenseGraph a -> [[a]]) -> DenseGraph a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DenseGraph a -> [[a]]
forall a. Unbox a => DenseGraph a -> [[a]]
toListDG

indexDG :: Int -> Int -> Int -> Int
indexDG :: Int -> Int -> Int -> Int
indexDG Int
n Int
i Int
j = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j
{-# INLINE indexDG #-}

matDG :: (U.Unbox a) => DenseGraph a -> Int -> Int -> a
matDG :: forall a. Unbox a => DenseGraph a -> Int -> Int -> a
matDG DenseGraph a
gr Int
i Int
j = Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex (DenseGraph a -> Vector a
forall a. DenseGraph a -> Vector a
adjacentDG DenseGraph a
gr) (Int -> Int -> Int -> Int
indexDG (DenseGraph a -> Int
forall a. DenseGraph a -> Int
numVerticesDG DenseGraph a
gr) Int
i Int
j)
{-# INLINE matDG #-}

adjDG :: (U.Unbox a) => DenseGraph a -> Int -> U.Vector a
adjDG :: forall a. Unbox a => DenseGraph a -> Int -> Vector a
adjDG DenseGraph a
gr Int
i =
  Int -> Int -> Vector a -> Vector a
forall a. Unbox a => Int -> Int -> Vector a -> Vector a
U.unsafeSlice
    (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* DenseGraph a -> Int
forall a. DenseGraph a -> Int
numVerticesDG DenseGraph a
gr)
    (DenseGraph a -> Int
forall a. DenseGraph a -> Int
numVerticesDG DenseGraph a
gr)
    (DenseGraph a -> Vector a
forall a. DenseGraph a -> Vector a
adjacentDG DenseGraph a
gr)
{-# INLINE adjDG #-}

createDenseGraph :: (U.Unbox a) => Int -> (forall s. UM.MVector s a -> ST s ()) -> DenseGraph a
createDenseGraph :: forall a.
Unbox a =>
Int -> (forall s. MVector s a -> ST s ()) -> DenseGraph a
createDenseGraph Int
n forall s. MVector s a -> ST s ()
f = (forall s. ST s (DenseGraph a)) -> DenseGraph a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (DenseGraph a)) -> DenseGraph a)
-> (forall s. ST s (DenseGraph a)) -> DenseGraph a
forall a b. (a -> b) -> a -> b
$ do
  MVector s a
buf <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
UM.unsafeNew (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n)
  MVector s a -> ST s ()
forall s. MVector s a -> ST s ()
f MVector s a
buf
  Int -> Vector a -> DenseGraph a
forall a. Int -> Vector a -> DenseGraph a
DenseGraph Int
n (Vector a -> DenseGraph a)
-> ST s (Vector a) -> ST s (DenseGraph a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) a -> ST s (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s a
MVector (PrimState (ST s)) a
buf

fromListNDG :: (U.Unbox a) => Int -> [[a]] -> DenseGraph a
fromListNDG :: forall a. Unbox a => Int -> [[a]] -> DenseGraph a
fromListNDG Int
n [[a]]
xss =
  DenseGraph
    { numVerticesDG :: Int
numVerticesDG = Int
n
    , adjacentDG :: Vector a
adjacentDG = Int -> [a] -> Vector a
forall a. Unbox a => Int -> [a] -> Vector a
U.fromListN (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) ([a] -> Vector a) -> [a] -> Vector a
forall a b. (a -> b) -> a -> b
$ [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[a]]
xss
    }

fromListDG :: (U.Unbox a) => [[a]] -> DenseGraph a
fromListDG :: forall a. Unbox a => [[a]] -> DenseGraph a
fromListDG = [[a]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([[a]] -> Int)
-> (Int -> [[a]] -> DenseGraph a) -> [[a]] -> DenseGraph a
forall a b. ([[a]] -> a) -> (a -> [[a]] -> b) -> [[a]] -> b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> [[a]] -> DenseGraph a
forall a. Unbox a => Int -> [[a]] -> DenseGraph a
fromListNDG

toListDG :: (U.Unbox a) => DenseGraph a -> [[a]]
toListDG :: forall a. Unbox a => DenseGraph a -> [[a]]
toListDG DenseGraph a
gr = [Vector a -> [a]
forall a. Unbox a => Vector a -> [a]
U.toList (Vector a -> [a]) -> Vector a -> [a]
forall a b. (a -> b) -> a -> b
$ DenseGraph a -> Int -> Vector a
forall a. Unbox a => DenseGraph a -> Int -> Vector a
adjDG DenseGraph a
gr Int
i | Int
i <- [Int
0 .. DenseGraph a -> Int
forall a. DenseGraph a -> Int
numVerticesDG DenseGraph a
gr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]