Improve performance of uniformSubSet and reduce duplication

Author: lehins

libs/cardano-ledger-core/testlib/Test/Cardano/Ledger/Core/Arbitrary.hs

@@ -629,13 +629,24 @@ uniformSubSet mSubSetSize inputSet gen = do
   subSetSize <- case mSubSetSize of
     Nothing -> uniformRM (0, Set.size inputSet) gen
     Just n -> pure $ max 0 $ min (Set.size inputSet) n
-  go inputSet Set.empty subSetSize
+  if subSetSize < Set.size inputSet `div` 2
+    then
+      goAdd inputSet Set.empty subSetSize
+    else
+      goDelete inputSet (Set.size inputSet - subSetSize)
   where
-    go !s !acc !i
+    -- Constructing a new Set is faster when less then a half of original Set will be used
+    goAdd !s !acc !i
       | i <= 0 = pure acc
       | otherwise = do
           ix <- uniformRM (0, Set.size s - 1) gen
-          go (Set.deleteAt ix s) (Set.insert (Set.elemAt ix s) acc) (i - 1)
+          goAdd (Set.deleteAt ix s) (Set.insert (Set.elemAt ix s) acc) (i - 1)
+    -- Deleting is faster when more items need to be retained in the Set
+    goDelete !acc !i
+      | i <= 0 = pure acc
+      | otherwise = do
+          ix <- uniformRM (0, Set.size acc - 1) gen
+          goDelete (Set.deleteAt ix acc) (i - 1)
 
 uniformSubMap ::
   (StatefulGen g m, Ord k) =>

libs/cardano-ledger-test/src/Test/Cardano/Ledger/Constrained/Combinators.hs

@@ -14,7 +14,7 @@ import Data.Set (Set)
 import qualified Data.Set as Set
 import GHC.Stack (HasCallStack)
 import Test.Cardano.Ledger.Binary.Random (QC (..))
-import Test.Cardano.Ledger.Core.Arbitrary (uniformSubMap)
+import Test.Cardano.Ledger.Core.Arbitrary (uniformSubMap, uniformSubSet)
 import Test.QuickCheck hiding (Fixed, total)
 
 -- ==========================================================================
@@ -122,18 +122,7 @@ subsetFromSetWithSize mess set n
             ++ show (Set.size set)
         )
         mess
-  -- It is faster to remove extra elements, when we need to drop more than a half
-  | Set.size set `div` 2 < n = fst <$> help (flip const) set Set.empty (Set.size set - n)
-  -- It is faster to pick out random elements when we only need under a half of the original set.
-  | otherwise = snd <$> help seq set Set.empty n
-  where
-    help optSeq !source target count
-      | count <= 0 = pure (source, target)
-      | otherwise = do
-          (item, source') <- itemFromSet (("subsetFromSetWithSize " ++ show n) : mess) source
-          -- To avoid a memory leak we only force the target when it is the set we want to keep.
-          let target' = Set.insert item target
-          target' `optSeq` help optSeq source' target' (count - 1)
+  | otherwise = uniformSubSet (Just n) set QC
 
 -- | Generate a larger map, from a smaller map 'subset'. The new larger map, should have all the
 --   keys and values of the smaller map.