{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module IGraph
( Graph(..)
, LGraph(..)
, U
, D
, decodeC
, empty
, mkGraph
, fromLabeledEdges
, fromLabeledEdges'
, unsafeFreeze
, freeze
, unsafeThaw
, thaw
, neighbors
, pre
, suc
, nmap
, emap
, nfilter
, efilter
) where
import Conduit
import Control.Arrow ((&&&))
import Control.Monad (forM, forM_, liftM, replicateM)
import Control.Monad.Primitive
import Control.Monad.ST (runST)
import qualified Data.ByteString as B
import Data.Conduit.Cereal
import Data.Either (fromRight)
import Data.Hashable (Hashable)
import qualified Data.HashMap.Strict as M
import qualified Data.HashSet as S
import Data.List (sortBy)
import Data.Ord (comparing)
import Data.Serialize
import Foreign (castPtr)
import System.IO.Unsafe (unsafePerformIO)
import IGraph.Internal
import IGraph.Internal.Constants
import IGraph.Mutable
import IGraph.Types
class MGraph d => Graph d where
isDirected :: LGraph d v e -> Bool
isD :: d -> Bool
nNodes :: LGraph d v e -> Int
nNodes (LGraph g _) = unsafePerformIO $ igraphVcount g
{-# INLINE nNodes #-}
nodes :: LGraph d v e -> [Node]
nodes gr = [0 .. nNodes gr - 1]
{-# INLINE nodes #-}
labNodes :: Serialize v => LGraph d v e -> [LNode v]
labNodes gr = map (\i -> (i, nodeLab gr i)) $ nodes gr
{-# INLINE labNodes #-}
nEdges :: LGraph d v e -> Int
nEdges (LGraph g _) = unsafePerformIO $ igraphEcount g
{-# INLINE nEdges #-}
edges :: LGraph d v e -> [Edge]
edges gr = map (getEdgeByEid gr) [0 .. nEdges gr - 1]
{-# INLINE edges #-}
labEdges :: Serialize e => LGraph d v e -> [LEdge e]
labEdges gr = map (getEdgeByEid gr &&& getEdgeLabByEid gr) [0 .. nEdges gr - 1]
{-# INLINE labEdges #-}
hasEdge :: LGraph d v e -> Edge -> Bool
hasEdge (LGraph g _) (fr, to) = unsafePerformIO $ do
i <- igraphGetEid g fr to True False
return $ i >= 0
{-# INLINE hasEdge #-}
nodeLab :: Serialize v => LGraph d v e -> Node -> v
nodeLab (LGraph g _) i = unsafePerformIO $
igraphHaskellAttributeVAS g vertexAttr i >>= bsToByteString >>=
return . fromRight (error "decode failed") . decode
{-# INLINE nodeLab #-}
getNodes :: (Hashable v, Eq v) => LGraph d v e -> v -> [Node]
getNodes gr x = M.lookupDefault [] x $ _labelToNode gr
{-# INLINE getNodes #-}
edgeLab :: Serialize e => LGraph d v e -> Edge -> e
edgeLab (LGraph g _) (fr,to) = unsafePerformIO $
igraphGetEid g fr to True True >>=
igraphHaskellAttributeEAS g edgeAttr >>= bsToByteString >>=
return . fromRight (error "decode failed") . decode
{-# INLINE edgeLab #-}
getEdgeByEid :: LGraph d v e -> Int -> Edge
getEdgeByEid (LGraph g _) i = unsafePerformIO $ igraphEdge g i
{-# INLINE getEdgeByEid #-}
getEdgeLabByEid :: Serialize e => LGraph d v e -> Int -> e
getEdgeLabByEid (LGraph g _) i = unsafePerformIO $
igraphHaskellAttributeEAS g edgeAttr i >>= bsToByteString >>=
return . fromRight (error "decode failed") . decode
{-# INLINE getEdgeLabByEid #-}
instance Graph U where
isDirected = const False
isD = const False
instance Graph D where
isDirected = const True
isD = const True
data LGraph d v e = LGraph
{ _graph :: IGraph
, _labelToNode :: M.HashMap v [Node]
}
instance (Graph d, Serialize v, Serialize e, Hashable v, Eq v)
=> Serialize (LGraph d v e) where
put gr = do
put $ nNodes gr
go (nodeLab gr) (nNodes gr) 0
put $ nEdges gr
go (\i -> (getEdgeByEid gr i, getEdgeLabByEid gr i)) (nEdges gr) 0
where
go f n i | i >= n = return ()
| otherwise = put (f i) >> go f n (i+1)
get = do
nn <- get
nds <- replicateM nn get
ne <- get
es <- replicateM ne get
return $ mkGraph nds es
decodeC :: ( PrimMonad m, MonadThrow m, Graph d
, Serialize v, Serialize e, Hashable v, Eq v )
=> ConduitT B.ByteString o m (LGraph d v e)
decodeC = do
nn <- sinkGet get
nds <- replicateM nn $ sinkGet get
ne <- sinkGet get
conduitGet2 get .| deserializeGraph nds ne
empty :: (Graph d, Hashable v, Serialize v, Eq v, Serialize e)
=> LGraph d v e
empty = runST $ new 0 >>= unsafeFreeze
mkGraph :: (Graph d, Hashable v, Serialize v, Eq v, Serialize e)
=> [v]
-> [LEdge e]
-> LGraph d v e
mkGraph vattr es = runST $ do
g <- new 0
addLNodes vattr g
addLEdges es g
unsafeFreeze g
fromLabeledEdges :: (Graph d, Hashable v, Serialize v, Eq v, Serialize e)
=> [((v, v), e)] -> LGraph d v e
fromLabeledEdges es = mkGraph labels es'
where
es' = flip map es $ \((fr, to), x) -> ((f fr, f to), x)
where f x = M.lookupDefault undefined x labelToId
labels = S.toList $ S.fromList $ concat [ [a,b] | ((a,b),_) <- es ]
labelToId = M.fromList $ zip labels [0..]
fromLabeledEdges' :: (PrimMonad m, Graph d, Hashable v, Serialize v, Eq v, Serialize e)
=> a
-> (a -> ConduitT () ((v, v), e) m ())
-> m (LGraph d v e)
fromLabeledEdges' input mkConduit = do
(labelToId, _, ne) <- runConduit $ mkConduit input .|
foldlC f (M.empty, 0::Int, 0::Int)
let getId x = M.lookupDefault undefined x labelToId
runConduit $ mkConduit input .|
mapC (\((v1, v2), e) -> ((getId v1, getId v2), e)) .|
deserializeGraph (fst $ unzip $ sortBy (comparing snd) $ M.toList labelToId) ne
where
f (vs, nn, ne) ((v1, v2), _) =
let (vs', nn') = add v1 $ add v2 (vs, nn)
in (vs', nn', ne+1)
where
add v (m, i) = if v `M.member` m
then (m, i)
else (M.insert v i m, i + 1)
deserializeGraph :: ( PrimMonad m, Graph d, Hashable v, Serialize v
, Eq v, Serialize e )
=> [v]
-> Int
-> ConduitT (LEdge e) o m (LGraph d v e)
deserializeGraph nds ne = do
evec <- unsafePrimToPrim $ igraphVectorNew $ 2 * ne
bsvec <- unsafePrimToPrim $ bsvectorNew ne
let f i ((fr, to), attr) = unsafePrimToPrim $ do
igraphVectorSet evec (i*2) $ fromIntegral fr
igraphVectorSet evec (i*2+1) $ fromIntegral to
bsvectorSet bsvec i $ encode attr
return $ i + 1
_ <- foldMC f 0
gr@(MLGraph g) <- new 0
addLNodes nds gr
unsafePrimToPrim $ withAttr edgeAttr bsvec $ \ptr -> do
vptr <- fromPtrs [castPtr ptr]
withVectorPtr vptr (igraphAddEdges g evec . castPtr)
unsafeFreeze gr
{-# INLINE deserializeGraph #-}
freeze :: (Hashable v, Eq v, Serialize v, PrimMonad m)
=> MLGraph (PrimState m) d v e -> m (LGraph d v e)
freeze (MLGraph g) = do
g' <- unsafePrimToPrim $ igraphCopy g
unsafeFreeze (MLGraph g')
unsafeFreeze :: (Hashable v, Eq v, Serialize v, PrimMonad m)
=> MLGraph (PrimState m) d v e -> m (LGraph d v e)
unsafeFreeze (MLGraph g) = unsafePrimToPrim $ do
nV <- igraphVcount g
labels <- forM [0 .. nV - 1] $ \i ->
igraphHaskellAttributeVAS g vertexAttr i >>= bsToByteString >>=
return . fromRight (error "decode failed") . decode
return $ LGraph g $ M.fromListWith (++) $ zip labels $ map return [0..nV-1]
where
thaw :: (PrimMonad m, Graph d) => LGraph d v e -> m (MLGraph (PrimState m) d v e)
thaw (LGraph g _) = unsafePrimToPrim . liftM MLGraph . igraphCopy $ g
unsafeThaw :: PrimMonad m => LGraph d v e -> m (MLGraph (PrimState m) d v e)
unsafeThaw (LGraph g _) = return $ MLGraph g
neighbors :: LGraph d v e -> Node -> [Node]
neighbors gr i = unsafePerformIO $ do
vs <- igraphVsAdj i IgraphAll
vit <- igraphVitNew (_graph gr) vs
vitToList vit
suc :: LGraph D v e -> Node -> [Node]
suc gr i = unsafePerformIO $ do
vs <- igraphVsAdj i IgraphOut
vit <- igraphVitNew (_graph gr) vs
vitToList vit
pre :: LGraph D v e -> Node -> [Node]
pre gr i = unsafePerformIO $ do
vs <- igraphVsAdj i IgraphIn
vit <- igraphVitNew (_graph gr) vs
vitToList vit
nmap :: (Graph d, Serialize v1, Serialize v2, Hashable v2, Eq v2)
=> (LNode v1 -> v2) -> LGraph d v1 e -> LGraph d v2 e
nmap f gr = runST $ do
(MLGraph gptr) <- thaw gr
let gr' = MLGraph gptr
forM_ (nodes gr) $ \x -> setNodeAttr x (f (x, nodeLab gr x)) gr'
unsafeFreeze gr'
emap :: (Graph d, Serialize e1, Serialize e2, Hashable v, Eq v, Serialize v)
=> (LEdge e1 -> e2) -> LGraph d v e1 -> LGraph d v e2
emap f gr = runST $ do
(MLGraph gptr) <- thaw gr
let gr' = MLGraph gptr
forM_ [0 .. nEdges gr - 1] $ \i -> do
let lab = f (getEdgeByEid gr i, getEdgeLabByEid gr i)
setEdgeAttr i lab gr'
unsafeFreeze gr'
nfilter :: (Hashable v, Eq v, Serialize v, Graph d)
=> (LNode v -> Bool) -> LGraph d v e -> LGraph d v e
nfilter f gr = runST $ do
let deleted = fst $ unzip $ filter (not . f) $ labNodes gr
gr' <- thaw gr
delNodes deleted gr'
unsafeFreeze gr'
efilter :: (Hashable v, Eq v, Serialize v, Serialize e, Graph d)
=> (LEdge e -> Bool) -> LGraph d v e -> LGraph d v e
efilter f gr = runST $ do
let deleted = fst $ unzip $ filter (not . f) $ labEdges gr
gr' <- thaw gr
delEdges deleted gr'
unsafeFreeze gr'