GraphLoadIntMap.hs

-- -----------------------------------------------------------------------------
--                Immutable version (IntMap)
-- -----------------------------------------------------------------------------

module Main (main) where

import Control.Applicative

import Data.Maybe (fromMaybe)

import qualified Data.IntMap.Strict as M
import Data.IntMap.Strict (IntMap)


import Data.Attoparsec.ByteString.Lazy (parse)
import qualified Data.ByteString.Lazy as BL (getContents)
import Data.Attoparsec.ByteString.Char8 (decimal, char, endOfLine)


data Node = Node
      { nodeIns  :: [Int]
      , nodeOuts :: [Int] } deriving (Show)

type Graph = IntMap Node

main :: IO ()
main = print
     . fmap (useGraph . foldr addEdge M.empty) -- produce and consume the graph
     . parse ( many ( (,) <$> decimal
                          <* char ' '
                          <*> decimal
                          <* endOfLine
                    ))
     =<< BL.getContents

alterNode :: (Node -> Node) -> Int -> Graph -> Graph
alterNode f = M.alter (Just . f . fromMaybe emptyNode)
  where
    emptyNode = Node [] []

addEdge :: (Int, Int) -> Graph -> Graph
addEdge (x, y) = alterNode (\n -> n { nodeOuts = y : nodeOuts n }) x
               . alterNode (\n -> n { nodeIns  = x : nodeIns  n }) y

useGraph :: Graph -> Int
useGraph = M.foldr f 0
  where
    f (Node ins outs) c = c + length ins + length outs

GraphLoadMVector.hs

-- -----------------------------------------------------------------------------
--                Mutable version (MVector)
-- -----------------------------------------------------------------------------
module Main (main) where

import Control.Applicative

import Control.Monad
import Control.Monad.Trans (lift)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Control.Monad.Trans.State.Strict

import Control.Monad.ST.Safe (RealWorld)

import Data.Attoparsec.ByteString.Lazy (parse, Result(..))
import qualified Data.ByteString.Lazy as BL (getContents, null)
import Data.Attoparsec.ByteString.Char8 (decimal, char, endOfLine)

import qualified Data.Vector.Mutable as MV
import qualified Data.Vector as V


data Node = Node
      { nodeIns  :: [Int]
      , nodeOuts :: [Int] }

-- Frozen graph
type Graph = V.Vector Node

-- Mutable graph in monad m (either ST or IO, I use IO here)
type MGraph m = MV.MVector (PrimState m) Node

-- The vector grow on demand, so the reference on it can change
-- StateT allow to pass the vector reference arround
type GraphT m = StateT (MGraph m) m

emptyNode :: Node
emptyNode = Node [] []

runGraphT :: PrimMonad m => Int -> GraphT m a -> m (a, Graph)
runGraphT len0 gt = do
  (a, mv) <- runStateT gt =<< MV.replicate len0 emptyNode
  v <- V.unsafeFreeze mv
  return (a, v)

readNode :: PrimMonad m => Int -> GraphT m Node
readNode idx = do
  v <- growTo idx
  lift $ MV.unsafeRead v idx

alterNode :: PrimMonad m => (Node -> Node) -> Int -> GraphT m Node
alterNode f idx = do
  v <- growTo idx
  lift $ do
    n <- MV.unsafeRead v idx
    let n' = f n
    MV.unsafeWrite v idx n'
    return n'
{-# INLINE alterNode #-}

growTo :: PrimMonad m => Int -> GraphT m (MGraph m)
growTo idx = StateT $ \v -> do
  let len = MV.length v
  v' <- if idx < len
           then return v
           else do
              let len' = head . filter (> idx) . iterate (*2) $ len
              v' <- MV.grow v (len' - len)
              mapM_ (\i -> MV.unsafeWrite v' i emptyNode) [len .. len' - 1]
              return v'
  return (v', v')

addEdge :: PrimMonad m => (Int, Int) -> GraphT m ()
addEdge (x, y) = do
  -- Could do better and check the lenght only once...
  _ <- alterNode (\n -> n { nodeOuts = y : nodeOuts n }) x
  _ <- alterNode (\n -> n { nodeIns  = x : nodeIns  n }) y
  return ()




main :: IO ()
main = do
    pedges <- parse ( many ( (,) <$> decimal <* char ' '
                                 <*> decimal <* endOfLine
                            )) <$> BL.getContents
    case pedges of
        Done t edges | BL.null t -> print .   useGraph . snd
                                          <=< runGraphT 128 . mapM_ addEdge
                                          $   edges
                     | otherwise -> putStrLn $ "Incomplete parse. Remainder : " ++ show t
        Fail _ ctxs err -> mapM_ putStrLn (err:ctxs)

useGraph :: Graph -> Int
useGraph = V.foldr f 0
  where
    f (Node ins outs) c = c + length ins + length outs