{-# LANGUAGE Rank2Types #-}
module Main (main) where
import Control.Monad.ST
import Control.Monad.Error
import Data.Char
import Data.Ord (comparing)
import Data.List (maximumBy)
import Data.Array.Vector
import qualified Data.Array.Vector.Algorithms.Insertion as INS
import qualified Data.Array.Vector.Algorithms.Intro as INT
import qualified Data.Array.Vector.Algorithms.TriHeap as TH
import qualified Data.Array.Vector.Algorithms.Merge as M
import qualified Data.Array.Vector.Algorithms.Radix as R
import System.Environment
import System.Console.GetOpt
import System.Random.Mersenne
import Blocks
-- Does nothing. For testing the speed/heap allocation of the building blocks.
noalgo :: (UA e) => MUArr e s -> ST s ()
noalgo _ = return ()
-- Allocates a temporary buffer, like mergesort for similar purposes as noalgo.
alloc :: (UA e) => MUArr e s -> ST s ()
alloc arr | len <= 4 = arr `seq` return ()
| otherwise = (newMU (len `div` 2) :: ST s (MUArr Int s)) >> return ()
where len = lengthMU arr
displayTime :: String -> Integer -> IO ()
displayTime s elapsed = putStrLn $
s ++ " : " ++ show (fromIntegral elapsed / 1e12) ++ " seconds"
run :: String -> IO Integer -> IO ()
run s t = t >>= displayTime s
sortSuite :: String -> MTGen -> Int -> (forall s. MUArr Int s -> ST s ()) -> IO ()
sortSuite str g n sort = do
putStrLn $ "Testing: " ++ str
run "Random " $ speedTest n (rand g >=> modulo n) sort
run "Sorted " $ speedTest n ascend sort
run "Reverse-sorted " $ speedTest n (descend n) sort
run "Random duplicates " $ speedTest n (rand g >=> modulo 1000) sort
let m = 4 * (n `div` 4)
run "Median killer " $ speedTest m (medianKiller m) sort
partialSortSuite :: String -> MTGen -> Int -> Int
-> (forall s. MUArr Int s -> Int -> ST s ()) -> IO ()
partialSortSuite str g n k sort = sortSuite str g n (\a -> sort a k)
-- -----------------
-- Argument handling
-- -----------------
data Algorithm = DoNothing
| Allocate
| InsertionSort
| IntroSort
| IntroPartialSort
| IntroSelect
| TriHeapSort
| TriHeapPartialSort
| TriHeapSelect
| MergeSort
| RadixSort
deriving (Show, Read, Enum, Bounded)
data Options = O { algos :: [Algorithm], elems :: Int, portion :: Int, usage :: Bool } deriving (Show)
defaultOptions :: Options
defaultOptions = O [] 10000 1000 False
type OptionsT = Options -> Either String Options
options :: [OptDescr OptionsT]
options = [ Option ['A'] ["algorithm"] (ReqArg parseAlgo "ALGO")
("Specify an algorithm to be run. Options:\n" ++ algoOpts)
, Option ['n'] ["num-elems"] (ReqArg parseN "INT")
"Specify the size of arrays in algorithms."
, Option ['k'] ["portion"] (ReqArg parseK "INT")
"Specify the number of elements to partial sort/select in\nrelevant algorithms."
, Option ['?','v'] ["help"] (NoArg $ \o -> Right $ o { usage = True })
"Show options."
]
where
allAlgos :: [Algorithm]
allAlgos = [minBound .. maxBound]
algoOpts = fmt allAlgos
fmt (x:y:zs) = '\t' : pad (show x) ++ show y ++ "\n" ++ fmt zs
fmt [x] = '\t' : show x ++ "\n"
fmt [] = ""
size = (" " ++) . maximumBy (comparing length) . map show $ allAlgos
pad str = zipWith const (str ++ repeat ' ') size
parseAlgo :: String -> Options -> Either String Options
parseAlgo "None" o = Right $ o { algos = [] }
parseAlgo "All" o = Right $ o { algos = [DoNothing .. RadixSort] }
parseAlgo s o = leftMap (\e -> "Unrecognized algorithm `" ++ e ++ "'")
. fmap (\v -> o { algos = v : algos o }) $ readEither s
leftMap :: (a -> b) -> Either a c -> Either b c
leftMap f (Left a) = Left (f a)
leftMap _ (Right c) = Right c
parseNum :: (Int -> Options) -> String -> Either String Options
parseNum f = leftMap (\e -> "Invalid numeric argument `" ++ e ++ "'") . fmap f . readEither
parseN, parseK :: String -> Options -> Either String Options
parseN s o = parseNum (\n -> o { elems = n }) s
parseK s o = parseNum (\k -> o { portion = k }) s
readEither :: Read a => String -> Either String a
readEither s = case reads s of
[(x,t)] | all isSpace t -> Right x
_ -> Left s
runTest :: MTGen -> Int -> Int -> Algorithm -> IO ()
runTest g n k alg = case alg of
DoNothing -> sortSuite "no algorithm" g n noalgo
Allocate -> sortSuite "allocate" g n alloc
InsertionSort -> sortSuite "insertion sort" g n INS.sort
IntroSort -> sortSuite "introsort" g n INT.sort
IntroPartialSort -> partialSortSuite "partial introsort" g n k INT.partialSort
IntroSelect -> partialSortSuite "introselect" g n k INT.select
TriHeapSort -> sortSuite "tri-heap sort" g n TH.sort
TriHeapPartialSort -> partialSortSuite "partial tri-heap sort" g n k TH.partialSort
TriHeapSelect -> partialSortSuite "tri-heap select" g n k TH.select
MergeSort -> sortSuite "merge sort" g n M.sort
RadixSort -> sortSuite "radix sort" g n R.sort
_ -> putStrLn $ "Currently unsupported algorithm: " ++ show alg
main :: IO ()
main = do args <- getArgs
gen <- getStdGen
case getOpt Permute options args of
(fs, _, []) -> case foldl (>>=) (Right defaultOptions) fs of
Left err -> putStrLn $ usageInfo err options
Right opts | not (usage opts) ->
mapM_ (runTest gen (elems opts) (portion opts)) (algos opts)
| otherwise -> putStrLn $ usageInfo "uvector-algorithms-bench" options
(_, _, errs) -> putStrLn $ usageInfo (concat errs) options