content: {-# LANGUAGE BangPatterns #-}
  
  module Data.Benchmarks.UnorderedCollections.Distributions
    ( makeRandomData
    , makeRandomVariateData
      -- * Workloads
    , insertWorkload
    , deleteWorkload
    , uniformLookupWorkload
    , exponentialLookupWorkload
  
    , loadOnly
    , loadAndUniformLookup
    , loadAndSkewedLookup
    , loadAndDeleteAll
    , loadAndDeleteSome
    , uniformlyMixed
    ) where
  
  
  import qualified Control.Concurrent.Thread as Th
  import           Control.DeepSeq
  import           Control.Monad
  import           Control.Monad.Reader
  import           Control.Monad.Trans (liftIO)
  import           Data.Benchmarks.UnorderedCollections.Types
  import qualified Data.Vector as V
  import qualified Data.Vector.Mutable as MV
  import qualified Data.Vector.Unboxed as VU
  import           Data.Vector (Vector)
  import qualified Data.Vector.Algorithms.Shuffle as V
  import           GHC.Conc (numCapabilities)
  import           Statistics.Distribution
  import           Statistics.Distribution.Exponential
  import           System.Random.MWC
  
  import           Criterion.Collection.Types
  
  
  ------------------------------------------------------------------------------
  debug :: (MonadIO m) => String -> m ()
  debug s = liftIO $ putStrLn s
  
  
  ------------------------------------------------------------------------------
  makeRandomData :: (NFData k) =>
                    (GenIO -> IO k)
                 -> Int
                 -> WorkloadMonad (Vector (k,Int))
  makeRandomData !genFunc !n = do
      rng <- getRNG
      debug $ "making " ++ show n ++ " data items"
      keys <- liftIO $ vreplicateM n rng genFunc
      let !v = keys `V.zip` vals
      let !_ = forceVector v
      debug $ "made " ++ show n ++ " data items"
      return $! v
  
    where
      vals      = V.enumFromN 0 n
  
  
  ------------------------------------------------------------------------------
  makeRandomVariateData :: (Ord k, NFData k, Variate k) =>
                           Int
                        -> WorkloadMonad (Vector (k,Int))
  makeRandomVariateData = makeRandomData uniform
  
  
  ------------------------------------------------------------------------------
  insertWorkload :: (NFData k) => Vector (k,Int) -> Vector (Operation k)
  insertWorkload = mapForce $ \(k,v) -> Insert k v
  
  
  ------------------------------------------------------------------------------
  deleteWorkload :: (NFData k) => Vector (k,Int) -> Vector (Operation k)
  deleteWorkload = mapForce $ \(k,_) -> Delete k
  
  
  ------------------------------------------------------------------------------
  uniformLookupWorkload :: (NFData k) =>
                           Vector (k,Int)
                        -> Int
                        -> WorkloadMonad (Vector (Operation k))
  uniformLookupWorkload !vec !ntimes = do
      rng <- getRNG
      debug $ "uniformLookupWorkload: generating " ++ show ntimes ++ " lookups"
      v <- liftIO $ vreplicateM ntimes rng f
      debug $ "uniformLookupWorkload: done"
      return v
  
    where
      !n = V.length vec
      f r = do
          idx <- pick
          let (k,_) = V.unsafeIndex vec idx
          return $ Lookup k
        where
          pick = uniformR (0,n-1) r
  
  
  ------------------------------------------------------------------------------
  exponentialLookupWorkload :: (NFData k) =>
                               Double
                            -> Vector (k,Int)
                            -> Int
                            -> WorkloadMonad (Vector (Operation k))
  exponentialLookupWorkload !lambda !vec !ntimes = do
      rng <- getRNG
      liftIO $ vreplicateM ntimes rng f
    where
      !dist = exponential lambda
      !n    = V.length vec
      !n1   = n-1
      !nd   = fromIntegral n
  
      f r = do
          x <- uniformR (0.1, 7.0) r
          let idx = max 0 . min n1 . round $ nd * density dist x
          let (k,_) = V.unsafeIndex vec idx
          return $! Lookup k
  
  
  ------------------------------------------------------------------------------
  loadOnly :: (NFData k) =>
              (GenIO -> IO k)     -- ^ rng for keys
           -> WorkloadGenerator (Operation k)
  loadOnly !genFunc !n = return $ Workload V.empty f
    where
      f _ = liftM insertWorkload $ makeRandomData genFunc n
  
  
  ------------------------------------------------------------------------------
  loadAndUniformLookup :: (NFData k) =>
                          (GenIO -> IO k)  -- ^ rng for keys
                       -> WorkloadGenerator (Operation k)
  loadAndUniformLookup !genFunc !n = do
      !vals <- makeRandomData genFunc n
      let !inserts = insertWorkload vals
  
      return $! Workload inserts $ uniformLookupWorkload vals
  
  
  ------------------------------------------------------------------------------
  loadAndSkewedLookup :: (NFData k) =>
                         (GenIO -> IO k)  -- ^ rng for keys
                      -> WorkloadGenerator (Operation k)
  loadAndSkewedLookup !genFunc !n = do
      !vals <- makeRandomData genFunc n
      let !inserts = insertWorkload vals
      return $! Workload inserts $ exponentialLookupWorkload 1.5 vals
  
  
  ------------------------------------------------------------------------------
  loadAndDeleteAll :: (NFData k) =>
                      (GenIO -> IO k)     -- ^ key generator
                   -> WorkloadGenerator (Operation k)
  loadAndDeleteAll !genFunc !n = do
      rng <- getRNG
      !vals <- makeRandomData genFunc n
      let !inserts = insertWorkload vals
      let !deletes = deleteWorkload $ V.shuffle rng vals
  
      return $ Workload inserts (const $ return deletes)
  
  
  ------------------------------------------------------------------------------
  loadAndDeleteSome :: (NFData k) =>
                       (GenIO -> IO k)
                    -> WorkloadGenerator (Operation k)
  loadAndDeleteSome !genFunc !n = do
      !vals <- makeRandomData genFunc n
      let !inserts = insertWorkload vals
  
      return $ Workload inserts $ f vals
  
    where
      f vals k = do
          rng <- getRNG
          return $ deleteWorkload $ V.take k $ V.shuffle rng vals
  
  
  ------------------------------------------------------------------------------
  uniformlyMixed :: (NFData k) =>
                    (GenIO -> IO k)
                 -> Double
                 -> Double
                 -> WorkloadGenerator (Operation k)
  uniformlyMixed !genFunc !lookupPercentage !deletePercentage !n = do
      let !numLookups = ceiling (fromIntegral n * lookupPercentage)
      let !numDeletes = ceiling (fromIntegral n * deletePercentage)
  
      !vals <- makeRandomData genFunc n
      let !inserts = insertWorkload vals
      !lookups <- uniformLookupWorkload vals numLookups
  
      rng <- getRNG
      let !deletes = deleteWorkload $ V.take numDeletes $ V.shuffle rng vals
      let !out = V.shuffle rng $ V.concat [inserts, lookups, deletes]
  
      return $! Workload V.empty $ const $ return $ forceVector out
  
  
  ------------------------------------------------------------------------------
  -- utilities
  ------------------------------------------------------------------------------
  forceVector :: (NFData k) => Vector k -> Vector k
  forceVector !vec = V.foldl' force () vec `seq` vec
    where
      force x v = x `deepseq` v `deepseq` ()
  
  
  mapForce :: (NFData b) => (a -> b) -> Vector a -> Vector b
  mapForce !f !vIn = let !vOut = V.map f vIn
                     in forceVector vOut
  
  
  -- split a GenIO
  splitGenIO :: GenIO -> IO GenIO
  splitGenIO rng = VU.replicateM 256 (uniform rng) >>= initialize
  
  
  -- vector replicateM is slow as dogshit.
  vreplicateM :: Int -> GenIO -> (GenIO -> IO a) -> IO (Vector a)
  vreplicateM n origRng act = do
      rngs <- replicateM numCapabilities (splitGenIO origRng)
      mv <- MV.new n
      let actions = map (f mv) (parts `zip` rngs)
      results <- liftM (map snd) $ mapM Th.forkIO actions
      _ <- sequence results
      V.unsafeFreeze mv
  
    where
      parts = partition (n-1) numCapabilities
  
      f mv ((low,high),rng) = do
          f' low
        where
          f' !idx | idx > high = return ()
                  | otherwise = do
                                  x <- act rng
                                  MV.unsafeWrite mv idx x
                                  f' (idx+1)
  
  
  partition :: Int -> Int -> [(Int,Int)]
  partition n k = ys `zip` xs
    where
      xs = map f [1..k]
      ys = 0:(map (+1) xs)
      f i = (i * n) `div` k
  