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