micro-bench: factor out frequency generator helper

IntersectMBO · May 9, 2024 · 6b7489d · 6b7489d
1 parent 7323091
commit 6b7489d
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 40 deletions.
diff --git a/bench/micro/Bench/Database/LSMTree/Internal/Lookup.hs b/bench/micro/Bench/Database/LSMTree/Internal/Lookup.hs
@@ -25,8 +25,8 @@ import qualified Data.Map.Strict as Map
 import           Data.Maybe (fromMaybe)
 import qualified Data.Vector as V
 import           Database.LSMTree.Extras.Orphans ()
-import           Database.LSMTree.Extras.Random (sampleUniformWithReplacement,
-                     uniformWithoutReplacement)
+import           Database.LSMTree.Extras.Random (frequency,
+                     sampleUniformWithReplacement, uniformWithoutReplacement)
 import           Database.LSMTree.Extras.UTxO
 import           Database.LSMTree.Internal.Entry (Entry (..), NumEntries (..))
 import           Database.LSMTree.Internal.Lookup (BatchSize (..),
@@ -185,21 +185,6 @@ lookupsEnv g nentries npos nneg = do
     assert (length lookups' == npos + nneg) $ pure ()
     pure (entries', lookups')
 
-frequency :: [(Int, StdGen -> (a, StdGen))] -> StdGen -> (a, StdGen)
-frequency xs0 g
-  | any ((< 0) . fst) xs0 = error "frequency: frequencies must be non-negative"
-  | tot == 0              = error "frequency: at least one frequency should be non-zero"
-  | otherwise = pick i xs0
- where
-  (i, g') = uniformR (1, tot) g
-
-  tot = sum (map fst xs0)
-
-  pick n ((k,x):xs)
-    | n <= k    = x g'
-    | otherwise = pick (n-k) xs
-  pick _ _  = error "QuickCheck.pick used with empty list"
-
 -- TODO: tweak distribution
 randomEntry :: StdGen -> (Entry UTxOValue UTxOBlob, StdGen)
 randomEntry g = frequency [

diff --git a/bench/micro/Bench/Database/LSMTree/Internal/WriteBuffer.hs b/bench/micro/Bench/Database/LSMTree/Internal/WriteBuffer.hs
@@ -8,11 +8,11 @@ import           Control.DeepSeq (NFData (..))
 import           Criterion.Main (Benchmark, bench, bgroup)
 import qualified Criterion.Main as Cr
 import           Data.Bifunctor (first)
-import qualified Data.ByteString as BS
 import qualified Data.List as List
 import           Data.Maybe (fromMaybe)
 import           Data.Word (Word64)
 import           Database.LSMTree.Extras.Orphans ()
+import           Database.LSMTree.Extras.Random (frequency, randomByteStringR)
 import           Database.LSMTree.Extras.UTxO
 import           Database.LSMTree.Internal.Entry
 import           Database.LSMTree.Internal.Run (Run)
@@ -30,8 +30,7 @@ import           System.Directory (removeDirectoryRecursive)
 import qualified System.FS.API as FS
 import qualified System.FS.IO as FS
 import           System.IO.Temp
-import qualified System.Random as R
-import           System.Random (StdGen, mkStdGen, uniform, uniformR)
+import           System.Random (StdGen, mkStdGen, uniform)
 
 benchmarks :: Benchmark
 benchmarks = bgroup "Bench.Database.LSMTree.Internal.WriteBuffer" [
@@ -280,23 +279,3 @@ lookupsEnv Config {..} = take nentries . List.unfoldr (Just . randomKOp)
                   in  (Mupdate v, g')
           )
         ]
-
-frequency :: [(Int, Rnd a)] -> Rnd a
-frequency xs0 g
-  | any ((< 0) . fst) xs0 = error "frequency: frequencies must be non-negative"
-  | tot == 0              = error "frequency: at least one frequency should be non-zero"
-  | otherwise = pick i xs0
- where
-  (i, g') = uniformR (1, tot) g
-
-  tot = sum (map fst xs0)
-
-  pick n ((k,x):xs)
-    | n <= k    = x g'
-    | otherwise = pick (n-k) xs
-  pick _ _  = error "frequency: pick used with empty list"
-
-randomByteStringR :: (Int, Int) -> Rnd BS.ByteString
-randomByteStringR range g =
-    let (!l, !g')  = uniformR range g
-    in  R.genByteString l g'
diff --git a/src-extras/Database/LSMTree/Extras/Random.hs b/src-extras/Database/LSMTree/Extras/Random.hs
@@ -6,10 +6,16 @@ module Database.LSMTree.Extras.Random (
   , uniformWithReplacement
   , sampleUniformWithoutReplacement
   , sampleUniformWithReplacement
+    -- * Sampling from multiple distributions
+  , frequency
+    -- * Generators for specific data types
+  , randomByteStringR
   ) where
 
+import qualified Data.ByteString as BS
 import           Data.List (unfoldr)
 import qualified Data.Set as Set
+import qualified System.Random as R
 import           System.Random (StdGen, Uniform, uniform, uniformR)
 import           Text.Printf (printf)
 
@@ -57,3 +63,38 @@ sampleUniformWithReplacement rng0 n xs0 = take n $
       where
         (i, rng') = uniformR (0, Set.size xs - 1) rng
         !x        = Set.elemAt i xs
+
+{-------------------------------------------------------------------------------
+  Sampling from multiple distributions
+-------------------------------------------------------------------------------}
+
+-- | Chooses one of the given generators, with a weighted random distribution.
+-- The input list must be non-empty, weights should be non-negative, and the sum
+-- of weights should be non-zero (i.e., at least one weight should be positive).
+--
+-- Based on the implementation in @QuickCheck@.
+frequency :: [(Int, StdGen -> (a, StdGen))] -> StdGen -> (a, StdGen)
+frequency xs0 g
+  | any ((< 0) . fst) xs0 = error "frequency: frequencies must be non-negative"
+  | tot == 0              = error "frequency: at least one frequency should be non-zero"
+  | otherwise = pick i xs0
+ where
+  (i, g') = uniformR (1, tot) g
+
+  tot = sum (map fst xs0)
+
+  pick n ((k,x):xs)
+    | n <= k    = x g'
+    | otherwise = pick (n-k) xs
+  pick _ _  = error "frequency: pick used with empty list"
+
+{-------------------------------------------------------------------------------
+  Generators for specific data types
+-------------------------------------------------------------------------------}
+
+-- | Generates a random bytestring. Its length is uniformly distributed within
+-- the provided range.
+randomByteStringR :: (Int, Int) -> StdGen -> (BS.ByteString, StdGen)
+randomByteStringR range g =
+    let (!l, !g')  = uniformR range g
+    in  R.genByteString l g'