1{-# LANGUAGE Rank2Types #-} 2 3module Main (main) where 4 5import Prelude hiding (read, length) 6import qualified Prelude as P 7 8import Control.Monad 9import Control.Monad.ST 10 11import Data.Char 12import Data.Ord (comparing) 13import Data.List (maximumBy) 14 15import Data.Vector.Unboxed.Mutable 16 17import qualified Data.Vector.Algorithms.Insertion as INS 18import qualified Data.Vector.Algorithms.Intro as INT 19import qualified Data.Vector.Algorithms.Heap as H 20import qualified Data.Vector.Algorithms.Merge as M 21import qualified Data.Vector.Algorithms.Radix as R 22import qualified Data.Vector.Algorithms.AmericanFlag as AF 23import qualified Data.Vector.Algorithms.Tim as T 24 25import System.Environment 26import System.Console.GetOpt 27import System.Random.MWC 28 29import Blocks 30 31-- Does nothing. For testing the speed/heap allocation of the building blocks. 32noalgo :: (Unbox e) => MVector RealWorld e -> IO () 33noalgo _ = return () 34 35-- Allocates a temporary buffer, like mergesort for similar purposes as noalgo. 36alloc :: (Unbox e) => MVector RealWorld e -> IO () 37alloc arr | len <= 4 = arr `seq` return () 38 | otherwise = (new (len `div` 2) :: IO (MVector RealWorld Int)) >> return () 39 where len = length arr 40 41displayTime :: String -> Integer -> IO () 42displayTime s elapsed = putStrLn $ 43 s ++ " : " ++ show (fromIntegral elapsed / (1e12 :: Double)) ++ " seconds" 44 45run :: String -> IO Integer -> IO () 46run s t = t >>= displayTime s 47 48sortSuite :: String -> GenIO -> Int -> (MVector RealWorld Int -> IO ()) -> IO () 49sortSuite str g n sort = do 50 arr <- new n 51 putStrLn $ "Testing: " ++ str 52 run "Random " $ speedTest arr n (rand g >=> modulo n) sort 53 run "Sorted " $ speedTest arr n ascend sort 54 run "Reverse-sorted " $ speedTest arr n (descend n) sort 55 run "Random duplicates " $ speedTest arr n (rand g >=> modulo 1000) sort 56 let m = 4 * (n `div` 4) 57 run "Median killer " $ speedTest arr m (medianKiller m) sort 58 59partialSortSuite :: String -> GenIO -> Int -> Int 60 -> (MVector RealWorld Int -> Int -> IO ()) -> IO () 61partialSortSuite str g n k sort = sortSuite str g n (\a -> sort a k) 62 63-- ----------------- 64-- Argument handling 65-- ----------------- 66 67data Algorithm = DoNothing 68 | Allocate 69 | InsertionSort 70 | IntroSort 71 | IntroPartialSort 72 | IntroSelect 73 | HeapSort 74 | HeapPartialSort 75 | HeapSelect 76 | MergeSort 77 | RadixSort 78 | AmericanFlagSort 79 | TimSort 80 deriving (Show, Read, Enum, Bounded) 81 82data Options = O { algos :: [Algorithm], elems :: Int, portion :: Int, usage :: Bool } deriving (Show) 83 84defaultOptions :: Options 85defaultOptions = O [] 10000 1000 False 86 87type OptionsT = Options -> Either String Options 88 89options :: [OptDescr OptionsT] 90options = [ Option ['A'] ["algorithm"] (ReqArg parseAlgo "ALGO") 91 ("Specify an algorithm to be run. Options:\n" ++ algoOpts) 92 , Option ['n'] ["num-elems"] (ReqArg parseN "INT") 93 "Specify the size of arrays in algorithms." 94 , Option ['k'] ["portion"] (ReqArg parseK "INT") 95 "Specify the number of elements to partial sort/select in\nrelevant algorithms." 96 , Option ['?','v'] ["help"] (NoArg $ \o -> Right $ o { usage = True }) 97 "Show options." 98 ] 99 where 100 allAlgos :: [Algorithm] 101 allAlgos = [minBound .. maxBound] 102 algoOpts = fmt allAlgos 103 fmt (x:y:zs) = '\t' : pad (show x) ++ show y ++ "\n" ++ fmt zs 104 fmt [x] = '\t' : show x ++ "\n" 105 fmt [] = "" 106 size = (" " ++) . maximumBy (comparing P.length) . map show $ allAlgos 107 pad str = zipWith const (str ++ repeat ' ') size 108 109parseAlgo :: String -> Options -> Either String Options 110parseAlgo "None" o = Right $ o { algos = [] } 111parseAlgo "All" o = Right $ o { algos = [DoNothing .. AmericanFlagSort] } 112parseAlgo s o = leftMap (\e -> "Unrecognized algorithm `" ++ e ++ "'") 113 . fmap (\v -> o { algos = v : algos o }) $ readEither s 114 115leftMap :: (a -> b) -> Either a c -> Either b c 116leftMap f (Left a) = Left (f a) 117leftMap _ (Right c) = Right c 118 119parseNum :: (Int -> Options) -> String -> Either String Options 120parseNum f = leftMap (\e -> "Invalid numeric argument `" ++ e ++ "'") . fmap f . readEither 121 122parseN, parseK :: String -> Options -> Either String Options 123parseN s o = parseNum (\n -> o { elems = n }) s 124parseK s o = parseNum (\k -> o { portion = k }) s 125 126readEither :: Read a => String -> Either String a 127readEither s = case reads s of 128 [(x,t)] | all isSpace t -> Right x 129 _ -> Left s 130 131runTest :: GenIO -> Int -> Int -> Algorithm -> IO () 132runTest g n k alg = case alg of 133 DoNothing -> sortSuite "no algorithm" g n noalgo 134 Allocate -> sortSuite "allocate" g n alloc 135 InsertionSort -> sortSuite "insertion sort" g n insertionSort 136 IntroSort -> sortSuite "introsort" g n introSort 137 IntroPartialSort -> partialSortSuite "partial introsort" g n k introPSort 138 IntroSelect -> partialSortSuite "introselect" g n k introSelect 139 HeapSort -> sortSuite "heap sort" g n heapSort 140 HeapPartialSort -> partialSortSuite "partial heap sort" g n k heapPSort 141 HeapSelect -> partialSortSuite "heap select" g n k heapSelect 142 MergeSort -> sortSuite "merge sort" g n mergeSort 143 RadixSort -> sortSuite "radix sort" g n radixSort 144 AmericanFlagSort -> sortSuite "flag sort" g n flagSort 145 TimSort -> sortSuite "tim sort" g n timSort 146 147mergeSort :: MVector RealWorld Int -> IO () 148mergeSort v = M.sort v 149{-# NOINLINE mergeSort #-} 150 151introSort :: MVector RealWorld Int -> IO () 152introSort v = INT.sort v 153{-# NOINLINE introSort #-} 154 155introPSort :: MVector RealWorld Int -> Int -> IO () 156introPSort v k = INT.partialSort v k 157{-# NOINLINE introPSort #-} 158 159introSelect :: MVector RealWorld Int -> Int -> IO () 160introSelect v k = INT.select v k 161{-# NOINLINE introSelect #-} 162 163heapSort :: MVector RealWorld Int -> IO () 164heapSort v = H.sort v 165{-# NOINLINE heapSort #-} 166 167heapPSort :: MVector RealWorld Int -> Int -> IO () 168heapPSort v k = H.partialSort v k 169{-# NOINLINE heapPSort #-} 170 171heapSelect :: MVector RealWorld Int -> Int -> IO () 172heapSelect v k = H.select v k 173{-# NOINLINE heapSelect #-} 174 175insertionSort :: MVector RealWorld Int -> IO () 176insertionSort v = INS.sort v 177{-# NOINLINE insertionSort #-} 178 179radixSort :: MVector RealWorld Int -> IO () 180radixSort v = R.sort v 181{-# NOINLINE radixSort #-} 182 183flagSort :: MVector RealWorld Int -> IO () 184flagSort v = AF.sort v 185{-# NOINLINE flagSort #-} 186 187timSort :: MVector RealWorld Int -> IO () 188timSort v = T.sort v 189{-# NOINLINE timSort #-} 190 191main :: IO () 192main = getArgs >>= \args -> withSystemRandom $ \gen -> 193 case getOpt Permute options args of 194 (fs, _, []) -> case foldl (>>=) (Right defaultOptions) fs of 195 Left err -> putStrLn $ usageInfo err options 196 Right opts | not (usage opts) -> 197 mapM_ (runTest gen (elems opts) (portion opts)) (algos opts) 198 | otherwise -> putStrLn $ usageInfo "vector-algorithms-bench" options 199 (_, _, errs) -> putStrLn $ usageInfo (concat errs) options 200 201 202