Optimize genDelay and genValue

lib/Echidna/ABI.hs

@@ -19,10 +19,11 @@ import Data.Hashable (Hashable(..))
 import Data.HashMap.Strict (HashMap)
 import Data.HashMap.Strict qualified as M
 import Data.HashSet (HashSet, fromList, union)
-import Data.HashSet qualified as H
+import Data.Set (Set)
+import Data.Set qualified as Set
 import Data.List (intercalate)
 import Data.List.NonEmpty qualified as NE
-import Data.Maybe (fromMaybe, catMaybes)
+import Data.Maybe (fromMaybe, catMaybes, mapMaybe)
 import Data.Text (Text)
 import Data.Text qualified as T
 import Data.Text.Encoding (encodeUtf8)
@@ -108,28 +109,20 @@ data GenDict = GenDict { _pSynthA    :: Float
                          -- ^ Default seed to use if one is not provided in EConfig
                        , _rTypes     :: Text -> Maybe AbiType
                          -- ^ Return types of any methods we scrape return values from
+                       , _dictValues :: Set W256
                        }
 
 makeLenses 'GenDict
 
 hashMapBy :: (Hashable k, Hashable a, Eq k, Ord a) => (a -> k) -> [a] -> HashMap k (HashSet a)
 hashMapBy f = M.fromListWith union . fmap (\v -> (f v, fromList [v]))
 
-gaddConstants :: [AbiValue] -> GenDict -> GenDict
-gaddConstants l = constants <>~ hashMapBy abiValueType l
-
 gaddCalls :: [SolCall] -> GenDict -> GenDict
 gaddCalls c = wholeCalls <>~ hashMapBy (fmap $ fmap abiValueType) c
 
 defaultDict :: GenDict
 defaultDict = mkGenDict 0 [] [] 0 (const Nothing)
 
-dictValues :: GenDict -> [W256]
-dictValues g = catMaybes $ concatMap (\(_,h) -> map fromValue $ H.toList h) $ M.toList $ g ^. constants
-  where fromValue (AbiUInt _ n) = Just (fromIntegral n)
-        fromValue (AbiInt  _ n) = Just (fromIntegral n)
-        fromValue _             = Nothing
-
 deriving anyclass instance Hashable AbiType
 deriving anyclass instance Hashable AbiValue
 deriving anyclass instance Hashable Addr
@@ -143,7 +136,14 @@ mkGenDict :: Float      -- ^ Percentage of time to mutate instead of synthesize.
           -> (Text -> Maybe AbiType)
           -- ^ A return value typing rule
           -> GenDict
-mkGenDict p vs cs = GenDict p (hashMapBy abiValueType vs) (hashMapBy (fmap $ fmap abiValueType) cs)
+mkGenDict p vs cs s tr =
+  GenDict p (hashMapBy abiValueType vs) (hashMapBy (fmap $ fmap abiValueType) cs) s tr (mkDictValues vs)
+
+mkDictValues :: [AbiValue] -> Set W256
+mkDictValues vs = Set.fromList $ mapMaybe fromValue vs
+  where fromValue (AbiUInt _ n) = Just (fromIntegral n)
+        fromValue (AbiInt  _ n) = Just (fromIntegral n)
+        fromValue _             = Nothing
 
 -- Generation (synthesis)
 

lib/Echidna/Campaign.hs

@@ -253,6 +253,7 @@ callseq ic v w ql = do
       additions = H.unionWith S.union diffs results
   -- append to the constants dictionary
   modifying (hasLens . genDict . constants) . H.unionWith S.union $ additions
+  modifying (hasLens . genDict . dictValues) . DS.union $ mkDictValues $ S.toList $ S.unions $ H.elems additions
   where
     -- Given a list of transactions and a return typing rule, this checks whether we know the return
     -- type for each function called, and if we do, tries to parse the return value as a value of that

lib/Echidna/Transaction.hs

@@ -14,6 +14,7 @@ import Data.Has (Has(..))
 import Data.HashMap.Strict qualified as M
 import Data.Map (Map, toList)
 import Data.Maybe (mapMaybe)
+import Data.Set (Set)
 import Data.Vector qualified as V
 import EVM hiding (value)
 import EVM.ABI (abiValueType)
@@ -52,7 +53,7 @@ genTxM memo m = do
   World ss hmm lmm ps _ <- view hasLens
   genDict <- use hasLens
   mm <- getSignatures hmm lmm
-  let ns = dictValues genDict
+  let ns = _dictValues genDict
   s' <- rElem ss
   r' <- rElem $ NE.fromList (mapMaybe (toContractA mm) (toList m))
   c' <- genInteractionsM genDict (snd r')
@@ -66,26 +67,21 @@ genTxM memo m = do
       let metadata = lookupBytecodeMetadata memo bc
       (addr,) <$> M.lookup metadata mm
 
-genDelay :: MonadRandom m => W256 -> [W256] -> m W256
+genDelay :: MonadRandom m => W256 -> Set W256 -> m W256
 genDelay mv ds = do
-  let ds' = (`mod` (mv + 1)) <$> ds
-  x <- rElem (0 NE.:| ds')
-  y <- randValue
-  fromIntegral <$> oftenUsually x (fromIntegral y)
-  where randValue = getRandomR (1 :: Integer, fromIntegral mv)
+  join $ oftenUsually fromDict randValue
+  where randValue = fromIntegral <$> getRandomR (1 :: Integer, fromIntegral mv)
+        fromDict = (`mod` (mv + 1)) <$> rElem' ds
 
-genValue :: MonadRandom m => W256 -> [W256] -> [FunctionHash] -> SolCall -> m W256
+genValue :: MonadRandom m => W256 -> Set W256 -> [FunctionHash] -> SolCall -> m W256
 genValue mv ds ps sc =
-  if sig `elem` ps then do
-    let ds' = (`mod` (mv + 1)) <$> ds
-    x <- rElem (0 NE.:| ds')
-    y <- randValue
-    fromIntegral <$> oftenUsually x (fromIntegral y)
+  if sig `elem` ps then
+    join $ oftenUsually fromDict randValue
   else do
-    g <- usuallyVeryRarely (pure 0) randValue -- once in a while, this will generate value in a non-payable function
-    fromIntegral <$> g
-  where randValue = getRandomR (0 :: Integer, fromIntegral mv)
+    join $ usuallyVeryRarely (pure 0) randValue -- once in a while, this will generate value in a non-payable function
+  where randValue = fromIntegral <$> getRandomR (0 :: Integer, fromIntegral mv)
         sig = (hashSig . encodeSig . signatureCall) sc
+        fromDict = (`mod` (mv + 1)) <$> rElem' ds
 
 -- | Check if a 'Transaction' is as \"small\" (simple) as possible (using ad-hoc heuristics).
 canShrinkTx :: Tx -> Bool

lib/Echidna/Types/Random.hs

@@ -3,13 +3,19 @@ module Echidna.Types.Random where
 import Prelude hiding ((!!))
 import Control.Monad.Random.Strict (MonadRandom, getRandomR, weighted)
 import Data.List.NonEmpty ((!!), NonEmpty(..))
+import Data.Set (Set)
+import Data.Set qualified as S
 
 type Seed = Int
 
 -- | Get a random element of a non-empty list.
 rElem :: MonadRandom m => NonEmpty a -> m a
 rElem l  = (l !!) <$> getRandomR (0, length l - 1)
 
+-- | Get a random element of a Set
+rElem' :: MonadRandom m => Set a -> m a
+rElem' v = (`S.elemAt` v) <$> getRandomR (0, length v - 1)
+
 oftenUsually :: MonadRandom m => a -> a -> m a
 oftenUsually u r = weighted [(u, 10), (r, 1)]