improve serialize performance

haskell-igraph.cabal

@@ -60,6 +60,7 @@ library
     , bytestring >= 0.9
     , bytestring-lexing >= 0.5
     , cereal
+    , cereal-conduit
     , colour
     , conduit >= 1.3.0
     , primitive

src/IGraph.hs

@@ -5,8 +5,7 @@ module IGraph
     , U(..)
     , D(..)
     , Graph(..)
---    , encodeC
- --   , decodeC
+    , decodeC
     , empty
     , mkGraph
     , fromLabeledEdges
@@ -32,10 +31,13 @@ module IGraph
 
 import           Conduit
 import           Control.Arrow             ((***))
-import           Control.Monad             (forM, forM_, liftM, unless, replicateM)
+import           Control.Monad             (forM, forM_, liftM, replicateM,
+                                            unless)
 import           Control.Monad.Primitive
 import           Control.Monad.ST          (runST)
 import qualified Data.ByteString           as B
+import           Data.Conduit.Cereal
+import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
 import           Data.Hashable             (Hashable)
 import qualified Data.HashMap.Strict       as M
 import qualified Data.HashSet              as S
@@ -43,11 +45,12 @@ import           Data.List                 (sortBy)
 import           Data.Maybe
 import           Data.Ord                  (comparing)
 import           Data.Serialize
-import           Foreign                   (with)
+import           Foreign                   (with, castPtr)
 import           System.IO.Unsafe          (unsafePerformIO)
 
 import           IGraph.Internal.Attribute
 import           IGraph.Internal.Constants
+import           IGraph.Internal.Data
 import           IGraph.Internal.Graph
 import           IGraph.Internal.Selector
 import           IGraph.Mutable
@@ -146,20 +149,14 @@ instance (Graph d, Serialize v, Serialize e, Hashable v, Eq v)
             es <- replicateM ne get
             return $ mkGraph nds es
 
-        {-
-encodeC :: (Monad m, Graph d, Serialize v, Serialize e, Hashable v, Eq v)
-        => LGraph d v e -> ConduitT i B.ByteString m ()
-encodeC gr = do
-    sourcePut $ put (M.toList $ _labelToNode gr)
-    yieldMany (edges gr) .| mapC (\e -> (e, edgeLab gr e)) .| conduitPut put
-
 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
-    labelToId <- M.fromList <$> sinkGet get
-    conduitGet2 get .| deserializeGraphFromEdges 10000 labelToId
-    -}
+    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
@@ -186,19 +183,20 @@ fromLabeledEdges es = mkGraph labels es'
 
 -- | Deserialize a graph.
 fromLabeledEdges' :: (PrimMonad m, Graph d, Hashable v, Serialize v, Eq v, Serialize e)
-                  => Int -- ^ buffer size
-                  -> a    -- ^ Input, usually a file
+                  => a    -- ^ Input, usually a file
                   -> (a -> ConduitT () ((v, v), e) m ())  -- ^ deserialize the input into a stream of edges
                   -> m (LGraph d v e)
-fromLabeledEdges' bufferN input mkConduit = do
-    (labelToId, _) <- runConduit $ mkConduit input .| foldlC f (M.empty, 0::Int)
+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 bufferN
-            (fst $ unzip $ sortBy (comparing snd) $ M.toList labelToId)
+        deserializeGraph (fst $ unzip $ sortBy (comparing snd) $ M.toList labelToId) ne
   where
-    f acc ((v1, v2), _) = add v1 $ add v2 acc
+    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)
@@ -206,24 +204,25 @@ fromLabeledEdges' bufferN input mkConduit = do
 
 deserializeGraph :: ( PrimMonad m, Graph d, Hashable v, Serialize v
                     , Eq v, Serialize e )
-                 => Int -- ^ buffer size
-                 -> [v]
+                 => [v]
+                 -> Int   -- ^ The number of edges
                  -> ConduitT (LEdge e) o m (LGraph d v e)
-deserializeGraph bufferN nds = mkChunks bufferN .| buildGraph
-  where
-    buildGraph = do
-        gr <- new 0
-        addLNodes nds gr
-        mapM_C (\es -> addLEdges es gr)
-        unsafeFreeze gr
-    mkChunks n = do
-        isEmpty <- nullC
-        unless isEmpty $ do
-            go 0 >>= yield
-            mkChunks n
-      where
-        go i | i >= n = return []
-             | otherwise = await >>= maybe (return []) (\x -> fmap (x :) $ go (i+1))
+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
+            unsafeUseAsCStringLen (encode attr) $ \bs -> with (BSLen bs) $ \ptr ->
+                bsvectorSet bsvec i $ castPtr ptr
+            return $ i + 1
+    foldMC f 0
+    gr@(MLGraph g) <- new 0
+    addLNodes nds gr
+    unsafePrimToPrim $ withEdgeAttr $ \eattr -> with (mkStrRec eattr bsvec) $ \ptr -> do
+            vptr <- fromPtrs [castPtr ptr]
+            withVectorPtr vptr (igraphAddEdges g evec . castPtr)
+    unsafeFreeze gr
 {-# INLINE deserializeGraph #-}
 
 unsafeFreeze :: (Hashable v, Eq v, Serialize v, PrimMonad m)

src/IGraph/Internal/Clique.chs

@@ -10,7 +10,7 @@ import Foreign.C.Types
 {#import IGraph.Internal.Graph #}
 {#import IGraph.Internal.Data #}
 
-#include "igraph/igraph.h"
+#include "haskell_igraph.h"
 
 {#fun igraph_cliques as ^ { `IGraph', `VectorPtr', `Int', `Int' } -> `Int' #}
 

src/IGraph/Internal/Community.chs

@@ -9,7 +9,7 @@ import Foreign.C.Types
 {#import IGraph.Internal.Data #}
 {#import IGraph.Internal.Constants #}
 
-#include "igraph/igraph.h"
+#include "haskell_igraph.h"
 
 {#fun igraph_community_spinglass as ^
 { `IGraph'

src/IGraph/Internal/Data.chs

@@ -29,6 +29,7 @@ module IGraph.Internal.Data
     , BSVector(..)
     , withBSVector
     , bsvectorNew
+    , bsvectorSet
     , toBSVector
 
     , Matrix(..)

src/IGraph/Internal/Graph.chs

@@ -57,7 +57,16 @@ igraphNew n directed _ = igraphNew' n directed
 
 {# fun igraph_add_edge as ^ { `IGraph', `Int', `Int' } -> `()' #}
 
-{# fun igraph_add_edges as ^ { `IGraph', `Vector', id `Ptr ()' } -> `()' #}
+-- | The edges are given in a vector, the first two elements define the first
+-- edge (the order is from , to for directed graphs). The vector should
+-- contain even number of integer numbers between zero and the number of
+-- vertices in the graph minus one (inclusive). If you also want to add
+-- new vertices, call igraph_add_vertices() first.
+{# fun igraph_add_edges as ^
+    { `IGraph'     -- ^ The graph to which the edges will be added.
+    , `Vector'     -- ^ The edges themselves.
+    , id `Ptr ()'  -- ^ The attributes of the new edges.
+    } -> `()' #}
 
 
 -- generators
@@ -73,5 +82,4 @@ igraphNew n directed _ = igraphNew' n directed
 {#fun igraph_rewire as ^ { `IGraph', `Int', `Rewiring' } -> `Int' #}
 
 
-
 {#fun igraph_isoclass_create as ^ { +, `Int', `Int', `Bool' } -> `IGraph' #}

src/IGraph/Internal/Isomorphism.chs

@@ -7,7 +7,7 @@ import Foreign.C.Types
 {#import IGraph.Internal.Graph #}
 {#import IGraph.Internal.Data #}
 
-#include "igraph/igraph.h"
+#include "haskell_igraph.h"
 
 {#fun igraph_get_subisomorphisms_vf2 as ^ { `IGraph', `IGraph',
     id `Ptr ()', id `Ptr ()', id `Ptr ()', id `Ptr ()', `VectorPtr',

src/IGraph/Internal/Motif.chs

@@ -11,7 +11,7 @@ import Foreign.C.Types
 {#import IGraph.Internal.Constants #}
 {#import IGraph.Internal.Data #}
 
-#include "igraph/igraph.h"
+#include "haskell_igraph.h"
 
 {#fun igraph_triad_census as ^ { `IGraph'
                                , `Vector' } -> `Int' #}

src/IGraph/Mutable.hs

@@ -6,6 +6,8 @@ module IGraph.Mutable
     , setNodeAttr
     , edgeAttr
     , vertexAttr
+    , withVertexAttr
+    , withEdgeAttr
     )where
 
 import           Control.Monad                  (when, forM)
@@ -86,7 +88,6 @@ instance MGraph U where
         esptr <- igraphEsVector vptr
         igraphDeleteEdges g esptr
         return ()
-      where
 
 instance MGraph D where
     new n = unsafePrimToPrim $ igraphInit >>= igraphNew n True >>= return . MLGraph
@@ -97,7 +98,6 @@ instance MGraph D where
         esptr <- igraphEsVector vptr
         igraphDeleteEdges g esptr
         return ()
-      where
 
 setNodeAttr :: (PrimMonad m, Serialize v)
             => Int   -- ^ Node id

tests/Test/Attributes.hs

@@ -54,7 +54,7 @@ serializeTest = testCase "serialize test" $ do
             Left msg -> error msg
             Right r  -> r
         es' = map (\(a,b) -> ((nodeLab gr' a, nodeLab gr' b), edgeLab gr' (a,b))) $ edges gr'
-    gr'' <- runConduit $ encodeC gr .| decodeC :: IO (LGraph D NodeAttr EdgeAttr)
+    gr'' <- runConduit $ (yield $ encode gr) .| decodeC :: IO (LGraph D NodeAttr EdgeAttr)
     let es'' = map (\(a,b) -> ((nodeLab gr'' a, nodeLab gr'' b), edgeLab gr'' (a,b))) $ edges gr''
     assertBool "" $ sort (map show es) == sort (map show es') &&
         sort (map show es) == sort (map show es'')

tests/Test/Basic.hs

@@ -50,7 +50,7 @@ graphCreationLabeled = testGroup "Graph creation -- with labels"
     n = length $ nubSort $ concatMap (\((a,b),_) -> [a,b]) edgeList
     m = length edgeList
     gr = fromLabeledEdges edgeList :: LGraph D String Int
-    gr' = runST $ fromLabeledEdges' 10 edgeList yieldMany :: LGraph D String Int
+    gr' = runST $ fromLabeledEdges' edgeList yieldMany :: LGraph D String Int
 
 graphEdit :: TestTree
 graphEdit = testGroup "Graph editing"