Graph500-0.4.0: src/Gen.hs

-- |
-- Copyright : (C) 2013 Parallel Scientific Labs, LLC.
-- License   : GPLv2
--
-- A program to generate various test sets.

module Main where

import Control.Monad (mplus, forM_, when)
import Control.Monad.State
import Control.Monad.Error
import Data.Bits (shiftL, shiftR)
import Data.Array.IO (IOUArray, newArray, writeArray, readArray, getBounds, hPutArray)
import Data.List (stripPrefix)
import Data.Word (Word8)
import System.Environment (getArgs)
import System.Exit (exitSuccess, exitFailure)
import System.IO (openBinaryFile, IOMode(WriteMode), hClose)

import G500
import G500.GenerateFile

-------------------------------------------------------------------------------
-- Command line arguments handling.

-- |Type of graph generated.
-- @Graph500@ is a R-MAT graph that is useful for Graph500 benchmark.
--
-- @Simple@ is a simple graph where each i-th index is connected to
-- (i+[1..edgeFactor])-th index. In the case of overflow (dest index is
-- bigger than 2^scale-1) we just connect to max index allowed.
-- This type of graph is for verification purposes.
--
type MatchM a = ErrorT String (State [String]) a

matchArgs' :: [String] -> Either String (String, GraphType, Int, Int, Maybe (Int, String))
matchArgs' args = flip evalState args $ runErrorT $ do
	filename <- getArg
	graphType <- getArg
	ty <- matchGraphType graphType
	scale <- match "scale"
	edgeFactor <- match "edgeFactor" `mplus` return 16
	actionsInfo <- liftM Just matchActionsInfo `mplus` return Nothing
	return (filename, ty, scale, edgeFactor, actionsInfo)
	where
		getArg :: MatchM String
		getArg = do
			as <- get
			case as of
				(a : as) -> do
					put as
					return a
				_ -> fail "no args."
		match :: Read x => String -> MatchM x
		match prefix = do
			opt <- getArg
			case stripPrefix (prefix ++ "=") opt of
				Just suffix -> return (read suffix)
				Nothing -> fail $ prefix++" required."
		matchActionsInfo :: MatchM (Int, String)
		matchActionsInfo = do
			n <- match "numActions"
			s <- getArg
			return (n,s)
		matchGraphType :: String -> MatchM GraphType
		matchGraphType "graph500" = return Graph500
		matchGraphType "simple" = return Simple
		matchGraphType ty = fail $ "Bad graph type "++ty

matchArgs :: [String] -> IO (String, GraphType, Int, Int, Maybe (Int, String))
matchArgs args = case matchArgs' args of
	Left err -> do
		putStrLn $ "command line error: "++err
		putStrLn "usage: graph500gen filename (simple|graph500) scale=<scale> [edgeFactor=<edge factor>] [numActions=# actionsFilename]"
		putStrLn ""
		putStrLn "The data will be stored into file filename."
		putStrLn "If edgeFactor argument is omitted edgeFactor=16 assumed."
		exitFailure
	Right r -> return r

-------------------------------------------------------------------------------
-- Main driver.

main = do
	args <- getArgs
	(fn,ty,scale,edgeFactor, actions) <- matchArgs args
	generateWriteFile fn ty scale edgeFactor
	return ()