perf: cache corpus selection as prefix sums and generate fillers on demand

lib/Echidna/Mutator/Corpus.hs

@@ -1,13 +1,15 @@
 module Echidna.Mutator.Corpus where
 
+import Control.Monad (replicateM)
 import Control.Monad.Random.Strict (MonadRandom, getRandomR)
 import Data.Maybe (maybeToList)
-import Data.Set qualified as Set
+import Data.Vector qualified as V
+import Data.Vector.Unboxed qualified as VU
 
 import Echidna.Mutator.Array
 import Echidna.Transaction (forceMutateTx, mutateTx, shrinkTx)
 import Echidna.Types (MutationConsts)
-import Echidna.Types.Corpus
+import Echidna.Types.Corpus (CorpusSelector(..))
 import Echidna.Types.Random (weighted)
 import Echidna.Types.Tx (Tx)
 
@@ -51,45 +53,66 @@ cutRange _ len = (min 1 len, len)
 selectAndMutate
   :: MonadRandom m
   => TxsMutation
-  -> Corpus
+  -> CorpusSelector
   -> m [Tx]
-selectAndMutate m corpus = do
-  rtxs <- selectFromCorpus corpus
+selectAndMutate m sel = do
+  rtxs <- selectFromCorpus sel
   k <- getRandomR (cutRange m (length rtxs))
   mutator m $ take k rtxs
 
 selectAndCombine
   :: MonadRandom m
   => ([Tx] -> [Tx] -> m [Tx])
   -> Int
-  -> Corpus
-  -> [Tx]
+  -> CorpusSelector
+  -> m Tx
   -> m [Tx]
-selectAndCombine f ql corpus gtxs = do
-  rtxs1 <- selectFromCorpus corpus
-  rtxs2 <- selectFromCorpus corpus
-  txs <- f rtxs1 rtxs2
-  pure . take ql $ txs <> gtxs
-
+selectAndCombine f ql sel genOne = do
+  rtxs1 <- selectFromCorpus sel
+  rtxs2 <- selectFromCorpus sel
+  txs <- take ql <$> f rtxs1 rtxs2
+  gtxs <- replicateM (ql - length txs) genOne
+  pure $ txs <> gtxs
+
+-- | Pick a sequence with probability proportional to its weight: draw a point
+-- in the total weight, then binary-search the cumulative weights for the
+-- sequence whose slice covers it.
 selectFromCorpus
   :: MonadRandom m
-  => Corpus
+  => CorpusSelector
   -> m [Tx]
-selectFromCorpus =
-  weighted . map (\(i, txs) -> (txs, fromIntegral i)) . Set.toDescList
-
+selectFromCorpus sel = do
+  r <- getRandomR (0, VU.last sel.cumWeights - 1)
+  pure $ sel.seqs V.! firstGreater r
+  where
+    -- smallest index whose cumulative weight exceeds r; r < total weight
+    -- guarantees one exists
+    firstGreater r = go 0 (VU.length sel.cumWeights - 1)
+      where
+        go lo hi
+          | lo >= hi = lo
+          | sel.cumWeights VU.! mid > r = go lo mid
+          | otherwise = go (mid + 1) hi
+          where mid = (lo + hi) `div` 2
+
+-- | A corpus mutation takes the target sequence length, the prepared corpus,
+-- and a generator for filler transactions, run only as many times as the
+-- mutated sequence needs topping up to that length.
 getCorpusMutation
   :: MonadRandom m
   => CorpusMutation
-  -> (Int -> Corpus -> [Tx] -> m [Tx])
-getCorpusMutation (RandomAppend m) = \ql ctxs gtxs -> do
-  rtxs' <- selectAndMutate m ctxs
-  pure . take ql $ rtxs' ++ gtxs
-getCorpusMutation (RandomPrepend m) = \ql ctxs gtxs -> do
-  rtxs' <- selectAndMutate m ctxs
+  -> (Int -> CorpusSelector -> m Tx -> m [Tx])
+getCorpusMutation (RandomAppend m) = \ql sel genOne -> do
+  rtxs' <- take ql <$> selectAndMutate m sel
+  gtxs <- replicateM (ql - length rtxs') genOne
+  pure $ rtxs' ++ gtxs
+getCorpusMutation (RandomPrepend m) = \ql sel genOne -> do
+  rtxs' <- selectAndMutate m sel
   k <- getRandomR (0, ql - 1)
-  -- Pad with the remaining fresh transactions so the sequence has ql entries.
-  pure . take ql $ take k gtxs ++ rtxs' ++ drop k gtxs
+  let mid = take (ql - k) rtxs'
+  -- Pad with fresh transactions so the sequence has ql entries.
+  gtxs <- replicateM (ql - length mid) genOne
+  pure $ take k gtxs ++ mid ++ drop k gtxs
 getCorpusMutation RandomSplice = selectAndCombine spliceAtRandom
 getCorpusMutation RandomInterleave = selectAndCombine interleaveAtRandom
 

lib/Echidna/Types/Campaign.hs

@@ -13,6 +13,7 @@ import EVM.Solvers (Solver(..))
 
 import Echidna.ABI (GenDict, emptyDict)
 import Echidna.Types
+import Echidna.Types.Corpus (CorpusSelector)
 import Echidna.Types.Coverage (CoverageFileType, CoverageMap)
 import Echidna.Types.Signature (SolCallPrototype)
 import Echidna.Types.Tx (TxResult(..))
@@ -200,6 +201,10 @@ data WorkerState = WorkerState
     -- ^ Call sequences to bias generation towards, each with the probability
     --   of being used in place of a corpus-mutated sequence. Empty unless
     --   sequences were explicitly injected into this worker.
+  , corpusSelector :: !(Maybe (Int, CorpusSelector))
+    -- ^ Cached corpus selection structure, tagged with the corpus size it was
+    --   built at. The shared corpus only ever grows, so the size doubles as a
+    --   version stamp; see 'Echidna.Worker.Fuzz.randseq'.
   }
 
 initialWorkerState :: WorkerState
@@ -213,6 +218,7 @@ initialWorkerState =
               , runningThreads = []
               , sampledFunctions = Map.empty
               , prioritizedSequences = []
+              , corpusSelector = Nothing
               }
 
 defaultTestLimit :: Int

lib/Echidna/Types/Corpus.hs

@@ -1,10 +1,36 @@
 module Echidna.Types.Corpus where
 
-import Data.Set (Set, size)
+import Data.List (scanl')
+import Data.Set (Set, size, toList)
+import Data.Vector qualified as V
+import Data.Vector.Unboxed qualified as VU
 
 import Echidna.Types.Tx (Tx)
 
 type Corpus = Set (Int, [Tx])
 
 corpusSize :: Corpus -> Int
 corpusSize = size
+
+-- | Snapshot of the corpus prepared for weighted random selection: the
+-- sequences in one vector, the cumulative sums of their weights in another.
+-- Building it costs one corpus traversal, after which each draw is a binary
+-- search (see 'Echidna.Mutator.Corpus.selectFromCorpus') instead of
+-- re-listing and re-summing the whole Set. Only rebuilt when the corpus
+-- grows; see 'corpusSelector' in 'Echidna.Types.Campaign.WorkerState'.
+data CorpusSelector = CorpusSelector
+  { seqs       :: !(V.Vector [Tx])
+  , cumWeights :: !(VU.Vector Int)
+    -- ^ inclusive prefix sums of the sequence weights; the last entry is the
+    -- total weight. Weights are the 'ncallseqs' stamps given on insertion,
+    -- so younger sequences are favored.
+  }
+
+mkCorpusSelector :: Corpus -> CorpusSelector
+mkCorpusSelector corpus = CorpusSelector
+  { seqs = V.fromListN n (map snd entries)
+  , cumWeights = VU.fromListN n (drop 1 $ scanl' (+) 0 (map fst entries))
+  }
+  where
+    n = size corpus
+    entries = toList corpus

lib/Echidna/Worker/Fuzz.hs

@@ -8,7 +8,7 @@ import Control.Monad (forM_, replicateM, void)
 import Control.Monad.Catch (MonadThrow)
 import Control.Monad.Random.Strict (MonadRandom, evalRandT, getRandom)
 import Control.Monad.Reader (MonadReader, ask, asks, liftIO)
-import Control.Monad.State.Strict (MonadIO, MonadState, StateT, gets, runStateT)
+import Control.Monad.State.Strict (MonadIO, MonadState, StateT, gets, modify', runStateT)
 import Control.Monad.Trans (lift)
 import Data.IORef (atomicModifyIORef', readIORef)
 import Data.List.NonEmpty qualified as NE
@@ -24,6 +24,7 @@ import Echidna.Shrink (isShrinkable, shrinkWorkerTests)
 import Echidna.Transaction
 import Echidna.Types.Campaign
 import Echidna.Types.Config
+import Echidna.Types.Corpus (Corpus, CorpusSelector, corpusSize, mkCorpusSelector)
 import Echidna.Types.Random (rElem)
 import Echidna.Types.Test
 import Echidna.Types.Test qualified as Test
@@ -147,18 +148,35 @@ genStandardSeq deployedContracts = do
   let
     mutConsts = env.cfg.campaignConf.mutConsts
     seqLen = env.cfg.campaignConf.seqLen
+    genOne = genTx world deployedContracts
 
   -- TODO: include reproducer when optimizing
   --let rs = filter (not . null) $ map (.testReproducer) $ ca._tests
 
-  -- Generate new random transactions
-  randTxs <- replicateM seqLen (genTx world deployedContracts)
-  -- Generate a random mutator
-  cmut <- if seqLen == 1 then seqMutatorsStateless mutConsts
-                         else seqMutatorsStateful mutConsts
-  -- Fetch the mutator
-  let mut = getCorpusMutation cmut
   corpus <- liftIO $ readIORef env.corpusRef
   if null corpus
-    then pure randTxs -- Use the generated random transactions
-    else mut seqLen corpus randTxs -- Apply the mutator
+    then replicateM seqLen genOne -- Use fresh random transactions
+    else do
+      -- Generate a random mutator
+      cmut <- if seqLen == 1 then seqMutatorsStateless mutConsts
+                             else seqMutatorsStateful mutConsts
+      sel <- cachedCorpusSelector corpus
+      -- Apply the mutator, generating fresh transactions only for the part of
+      -- the sequence it doesn't fill from the corpus
+      getCorpusMutation cmut seqLen sel genOne
+
+-- | The corpus prepared for weighted selection, rebuilt only when the corpus
+-- changed. The shared corpus only ever grows an element at a time, so its
+-- size works as a version stamp.
+cachedCorpusSelector
+  :: MonadState WorkerState m
+  => Corpus
+  -> m CorpusSelector
+cachedCorpusSelector corpus = do
+  cached <- gets (.corpusSelector)
+  case cached of
+    Just (sz, sel) | sz == corpusSize corpus -> pure sel
+    _ -> do
+      let !sel = mkCorpusSelector corpus
+      modify' $ \ws -> ws { corpusSelector = Just (corpusSize corpus, sel) }
+      pure sel

lib/Echidna/Worker/Sequence.hs

@@ -115,8 +115,7 @@ callseq vm txSeq isReplaying = do
     -- Even if this takes a bit of time, this is okay as finding new coverage
     -- is expected to be infrequent in the long term
     newSize <- liftIO $ atomicModifyIORef' env.corpusRef $ \corp ->
-      -- Corpus is a bit too lazy, force the evaluation to reduce the memory usage
-      let !corp' = force $ addToCorpus (ncallseqs + 1) results corp
+      let !corp' = addToCorpus (ncallseqs + 1) results corp
       in (corp', corpusSize corp')
 
     (points, numCodehashes) <- liftIO $ coverageStats env.coverageRefInit env.coverageRefRuntime
@@ -207,10 +206,15 @@ callseq vm txSeq isReplaying = do
       getTupleVector (AbiTuple ts) = ts
       getTupleVector _ = error "Not a tuple!"
 
-  -- | Add transactions to the corpus, discarding reverted ones
+  -- | Add transactions to the corpus, discarding reverted ones. The new entry
+  -- is deep-forced so the shared corpus never retains VM results or thunks
+  -- through it; everything already in the corpus was forced on its own
+  -- insertion.
   addToCorpus :: Int -> [(Tx, VMResult Concrete)] -> Corpus -> Corpus
   addToCorpus n res corpus =
-    if null rtxs then corpus else Set.insert (n, rtxs) corpus
+    if null rtxs
+      then corpus
+      else let !entry = force (n, rtxs) in Set.insert entry corpus
     where rtxs = fst <$> res
 
 -- | Fold a sequence of transaction results into the sampling map, updating

src/test/Tests/Mutator.hs

@@ -10,7 +10,7 @@ import Test.Tasty (TestTree, testGroup)
 import Test.Tasty.HUnit (assertBool, testCase, (@?=))
 import Test.Tasty.QuickCheck
   (Gen, Positive(..), arbitrary, choose, counterexample, elements, forAll, listOf1, property,
-   testProperty, vectorOf, (===))
+   testProperty, (===))
 
 import EVM.ABI (AbiValue(..))
 
@@ -20,7 +20,7 @@ import Echidna.Config (defaultConfig)
 import Echidna.Mutator.Corpus (CorpusMutation(..), TxsMutation(..), cutRange, getCorpusMutation)
 import Echidna.Types.Campaign
 import Echidna.Types.Config (EConfig(..), Env(..))
-import Echidna.Types.Corpus (Corpus)
+import Echidna.Types.Corpus (Corpus, CorpusSelector, mkCorpusSelector)
 import Echidna.Types.Tx (Tx(..), TxCall(..))
 import Echidna.Types.Worker (WorkerType(..))
 import Tests.Encoding () -- Arbitrary Tx
@@ -30,10 +30,11 @@ mutatorTests = testGroup "Corpus mutation"
   [ testProperty "every mutation yields exactly seqLen transactions" $
       forAll genCorpus $ \corpus ->
       forAll (choose (1, 8)) $ \ql ->
-      forAll (vectorOf ql arbitrary) $ \gtxs ->
+      forAll arbitrary $ \gtx ->
       forAll (elements allMutations) $ \m ->
       forAll arbitrary $ \seed ->
-        length (evalRand (getCorpusMutation m ql corpus gtxs) (mkStdGen seed)) === ql
+        length (evalRand (getCorpusMutation m ql (mkCorpusSelector corpus) (pure gtx)) (mkStdGen seed))
+          === ql
   , testGroup "cut point"
       [ testCase "identity keeps a strict prefix" $ do
           cutRange Identity 5 @?= (0, 4)
@@ -48,20 +49,19 @@ mutatorTests = testGroup "Corpus mutation"
       [ testProperty "a stored transaction is tweaked or replaced, never replayed" $
           forAll (elements [RandomAppend Mutation, RandomPrepend Mutation]) $ \m ->
           forAll arbitrary $ \seed ->
-            case evalRand (getCorpusMutation m 1 singleton [fresh]) (mkStdGen seed) of
+            case evalRand (getCorpusMutation m 1 singleton (pure fresh)) (mkStdGen seed) of
               [tx] | tx == fresh -> property True
                    | otherwise ->
                        counterexample (show tx) $ tx /= stored && fnName tx == fnName stored
               txs -> counterexample (show txs) False
       , testProperty "a stored transaction with nothing to change is replaced" $
           forAll (elements [mkTx "stored" [], mkTx "stored" [AbiAddress 0]]) $ \s ->
           forAll arbitrary $ \seed ->
-            evalRand (getCorpusMutation (RandomAppend Mutation) 1 (Set.singleton (1, [s])) [fresh])
-                     (mkStdGen seed)
+            evalRand (getCorpusMutation (RandomAppend Mutation) 1 (single s) (pure fresh)) (mkStdGen seed)
               === [fresh]
       , testProperty "identity yields the fresh transaction" $
           forAll arbitrary $ \seed ->
-            evalRand (getCorpusMutation (RandomAppend Identity) 1 singleton [fresh]) (mkStdGen seed)
+            evalRand (getCorpusMutation (RandomAppend Identity) 1 singleton (pure fresh)) (mkStdGen seed)
               === [fresh]
       ]
   , testGroup "forced mutation"
@@ -83,9 +83,12 @@ mutatorTests = testGroup "Corpus mutation"
       ]
   ]
 
--- | A seqLen 1 corpus: one single-transaction sequence.
-singleton :: Corpus
-singleton = Set.singleton (1, [stored])
+-- | A seqLen 1 corpus prepared for selection: one single-transaction sequence.
+singleton :: CorpusSelector
+singleton = single stored
+
+single :: Tx -> CorpusSelector
+single tx = mkCorpusSelector (Set.singleton (1, [tx]))
 
 -- | The stored transaction takes an integer, the argument type the ABI mutators
 -- change least often. The fresh one is distinguishable by name.