micro-bench: factor out frequency generator helper

IntersectMBO · May 6, 2024 · dceebb1 · dceebb1
1 parent 758b458
commit dceebb1
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 32 deletions.
diff --git a/bench/micro/Bench/Database/LSMTree/Internal/Lookup.hs b/bench/micro/Bench/Database/LSMTree/Internal/Lookup.hs
@@ -24,8 +24,8 @@ import           Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 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 (..),
@@ -183,21 +183,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
@@ -13,6 +13,7 @@ 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)
 import           Database.LSMTree.Extras.UTxO
 import           Database.LSMTree.Internal.Entry
 import           Database.LSMTree.Internal.Run (Run)
@@ -279,21 +280,6 @@ lookupsEnv Config {..} = take nentries . List.unfoldr (Just . randomKOp)
           )
         ]
 
-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

diff --git a/src-extras/Database/LSMTree/Extras/Random.hs b/src-extras/Database/LSMTree/Extras/Random.hs
@@ -6,6 +6,8 @@ module Database.LSMTree.Extras.Random (
   , uniformWithReplacement
   , sampleUniformWithoutReplacement
   , sampleUniformWithReplacement
+    -- * Sampling from uniform distributions
+  , frequency
   ) where
 
 import           Data.List (unfoldr)
@@ -57,3 +59,23 @@ 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
+-------------------------------------------------------------------------------}
+
+-- | 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"