From 55a80914de235ebe06e237e0d03799fb01b231c9 Mon Sep 17 00:00:00 2001
From: Gustavo Grieco <31542053+ggrieco-tob@users.noreply.github.com>
Date: Tue, 28 May 2024 19:29:12 +0200
Subject: [PATCH] Improved shrinking removing reverts from reproducers (#1250)

* remove reverted sequences from reproducers

* fixes

* concat NoCalls

* clean useless no calls when delay is zero

* avoid reversing transaction sequence
---
 lib/Echidna/Shrink.hs   | 67 ++++++++++++++++++++++++++++-------------
 lib/Echidna/Types/Tx.hs | 19 ++++++++++++
 2 files changed, 65 insertions(+), 21 deletions(-)

diff --git a/lib/Echidna/Shrink.hs b/lib/Echidna/Shrink.hs
index fb00d3e05..a8641d5a1 100644
--- a/lib/Echidna/Shrink.hs
+++ b/lib/Echidna/Shrink.hs
@@ -8,14 +8,15 @@ import Control.Monad.State.Strict (MonadIO)
 import Control.Monad.ST (RealWorld)
 import Data.Set qualified as Set
 import Data.List qualified as List
+import Data.Maybe (mapMaybe)
 
-import EVM.Types (VM, VMType(Concrete))
+import EVM.Types (VM, VMType(..))
 
 import Echidna.Exec
 import Echidna.Transaction
 import Echidna.Types.Solidity (SolConf(..))
 import Echidna.Types.Test (TestValue(..), EchidnaTest(..), TestState(..), isOptimizationTest)
-import Echidna.Types.Tx (Tx(..))
+import Echidna.Types.Tx (Tx(..), hasReverted, isUselessNoCall, catNoCalls, TxCall(..))
 import Echidna.Types.Config
 import Echidna.Types.Campaign (CampaignConf(..))
 import Echidna.Test (getResultFromVM, checkETest)
@@ -31,24 +32,47 @@ shrinkTest vm test = do
     Large i | i >= env.cfg.campaignConf.shrinkLimit && not (isOptimizationTest test) ->
       pure $ Just test { state = Solved }
     Large i ->
-      if length test.reproducer > 1 || any canShrinkTx test.reproducer then do
-        maybeShrunk <- shrinkSeq vm (checkETest test) test.value test.reproducer
-        pure $ case maybeShrunk of
-          Just (txs, val, vm') -> do
-            Just test { state = Large (i + 1)
-                 , reproducer = txs
-                 , vm = Just vm'
-                 , result = getResultFromVM vm'
-                 , value = val }
-          Nothing ->
-            -- No success with shrinking this time, just bump trials
-            Just test { state = Large (i + 1) }
-      else
-        pure $ Just test { state = if isOptimizationTest test
-                                 then Large (i + 1)
-                                 else Solved }
+      do  repro <- removeReverts vm test.reproducer
+          let rr = removeUselessNoCalls $ catNoCalls repro 
+          if length rr > 1 || any canShrinkTx rr then do
+            maybeShrunk <- shrinkSeq vm (checkETest test) test.value rr
+            pure $ case maybeShrunk of
+              Just (txs, val, vm') -> do
+                Just test { state = Large (i + 1)
+                    , reproducer = txs
+                    , vm = Just vm'
+                    , result = getResultFromVM vm'
+                    , value = val }
+              Nothing ->
+                -- No success with shrinking this time, just bump trials
+                Just test { state = Large (i + 1), reproducer = rr}
+          else
+            pure $ Just test { state = if isOptimizationTest test
+                                    then Large (i + 1)
+                                    else Solved }
     _ -> pure Nothing
 
+replaceByNoCall :: Tx -> Tx
+replaceByNoCall tx = tx { call = NoCall }
+
+removeUselessNoCalls :: [Tx] -> [Tx]
+removeUselessNoCalls = mapMaybe f
+  where f tx = if isUselessNoCall tx then Nothing else Just tx 
+
+removeReverts :: (MonadIO m, MonadReader Env m, MonadThrow m) => VM Concrete RealWorld -> [Tx] -> m [Tx]
+removeReverts vm txs = do
+  let (itxs, le) = (init txs, last txs)
+  ftxs <- removeReverts' vm itxs []
+  return (ftxs ++ [le])
+
+removeReverts' :: (MonadIO m, MonadReader Env m, MonadThrow m) => VM Concrete RealWorld -> [Tx] -> [Tx] -> m [Tx]
+removeReverts' _ [] ftxs = return $ reverse ftxs
+removeReverts' vm (t:txs) ftxs = do
+  (_, vm') <- execTx vm t
+  if hasReverted vm'
+  then removeReverts' vm' txs (replaceByNoCall t: ftxs) 
+  else removeReverts' vm' txs (t:ftxs) 
+
 -- | Given a call sequence that solves some Echidna test, try to randomly
 -- generate a smaller one that still solves that test.
 shrinkSeq
@@ -60,11 +84,12 @@ shrinkSeq
   -> m (Maybe ([Tx], TestValue, VM Concrete RealWorld))
 shrinkSeq vm f v txs = do
   txs' <- uniform =<< sequence [shorten, shrunk]
-  (value, vm') <- check txs' vm
+  let txs'' = removeUselessNoCalls txs' 
+  (value, vm') <- check txs'' vm
   -- if the test passed it means we didn't shrink successfully
   pure $ case (value,v) of
-    (BoolValue False, _)              -> Just (txs', value, vm')
-    (IntValue x, IntValue y) | x >= y -> Just (txs', value, vm')
+    (BoolValue False, _)              -> Just (txs'', value, vm')
+    (IntValue x, IntValue y) | x >= y -> Just (txs'', value, vm')
     _                                 -> Nothing
   where
     check [] vm' = f vm'
diff --git a/lib/Echidna/Types/Tx.hs b/lib/Echidna/Types/Tx.hs
index fe820580a..11c1372cc 100644
--- a/lib/Echidna/Types/Tx.hs
+++ b/lib/Echidna/Types/Tx.hs
@@ -9,6 +9,7 @@ module Echidna.Types.Tx where
 import Prelude hiding (Word)
 
 import Control.Applicative ((<|>))
+import Control.Monad.ST (RealWorld)
 import Data.Aeson (FromJSON, ToJSON, parseJSON, toJSON, object, withObject, (.=), (.:))
 import Data.Aeson.TH (deriveJSON, defaultOptions)
 import Data.Aeson.Types (Parser)
@@ -199,6 +200,24 @@ data TxConf = TxConf
   -- ^ Maximum value to use in transactions
   }
 
+hasReverted :: VM Concrete RealWorld -> Bool
+hasReverted vm = let r = vm.result in
+  case r of
+    (Just (VMSuccess _)) -> False
+    _                    -> True
+
+isUselessNoCall :: Tx -> Bool
+isUselessNoCall tx = tx.call == NoCall && tx.delay == (0, 0)
+
+catNoCalls :: [Tx] -> [Tx]
+catNoCalls [] = []
+catNoCalls [tx] = [tx]
+catNoCalls (tx1:tx2:xs) = 
+    case (tx1.call, tx2.call) of 
+      (NoCall, NoCall) -> catNoCalls (nc:xs)
+      _                -> tx1 : catNoCalls (tx2:xs)
+  where nc = tx1 { delay = (fst tx1.delay + fst tx2.delay, snd tx1.delay + snd tx2.delay) }  
+
 -- | Transform a VMResult into a more hash friendly sum type
 getResult :: VMResult Concrete s -> TxResult
 getResult = \case