From 839d2d32ce9463f6b09aac5cdc8d64112ae998e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Emilio=20L=C3=B3pez?= <emilio.lopez@trailofbits.com>
Date: Thu, 11 Jul 2024 14:35:20 +0200
Subject: [PATCH] Address review comments

---
 lib/Echidna/Exec.hs           | 11 +++++------
 lib/Echidna/Output/Source.hs  |  8 ++++----
 lib/Echidna/Types/Coverage.hs |  3 +++
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/lib/Echidna/Exec.hs b/lib/Echidna/Exec.hs
index 7dc036f81..f1170d9fa 100644
--- a/lib/Echidna/Exec.hs
+++ b/lib/Echidna/Exec.hs
@@ -7,7 +7,7 @@ module Echidna.Exec where
 import Optics.Core
 import Optics.State.Operators
 
-import Control.Monad (when, forM_)
+import Control.Monad (when)
 import Control.Monad.Catch (MonadThrow(..))
 import Control.Monad.State.Strict (MonadState(get, put), execState, runStateT, MonadIO(liftIO), gets, modify', execStateT)
 import Control.Monad.Reader (MonadReader, ask, asks)
@@ -292,9 +292,8 @@ execTxWithCov tx = do
           let size = BS.length . forceBuf . fromJust . view bytecode $ contract
           if size == 0 then pure Nothing else do
             -- IO for making a new vec
-            vec <- VMut.new size
             -- We use -1 for opIx to indicate that the location was not covered
-            forM_ [0..size-1] $ \i -> VMut.write vec i (-1, 0, 0)
+            vec <- VMut.replicate size (-1, 0, 0)
             pure $ Just vec
 
         statsRef <- getTLS env.statsRef
@@ -302,8 +301,7 @@ execTxWithCov tx = do
           let size = BS.length . forceBuf . fromJust . view bytecode $ contract
           if size == 0 then pure Nothing else do
             -- IO for making a new vec
-            vec <- VMut.new size
-            forM_ [0..size-1] $ \i -> VMut.write vec i (0, 0)
+            vec <- VMut.replicate size (0, 0)
             pure $ Just vec
 
         case maybeCovVec of
@@ -320,7 +318,8 @@ execTxWithCov tx = do
                 (_, depths, results) | depth < 64 && not (depths `testBit` depth) -> do
                   VMut.write vec pc (opIx, depths `setBit` depth, results `setBit` fromEnum Stop)
                   writeIORef covContextRef (True, Just (vec, pc))
-                _ -> modifyIORef' covContextRef $ \(new, _) -> (new, Just (vec, pc))
+                _ ->
+                  modifyIORef' covContextRef $ \(new, _) -> (new, Just (vec, pc))
 
       -- | Get the VM's current execution location
       currentCovLoc vm = (vm.state.pc, fromMaybe 0 $ vmOpIx vm, length vm.frames)
diff --git a/lib/Echidna/Output/Source.hs b/lib/Echidna/Output/Source.hs
index 0923caf11..ca5e7dbc0 100644
--- a/lib/Echidna/Output/Source.hs
+++ b/lib/Echidna/Output/Source.hs
@@ -42,15 +42,15 @@ zipSumStats v1 v2 = do
   vec2 <- v2
   return [(exec1 + exec2, revert1 + revert2) | (exec1, revert1) <- vec1 | (exec2, revert2) <- vec2]
 
-mvToList :: (VU.Unbox a) => VU.IOVector a -> IO [a]
-mvToList = fmap U.toList . U.freeze
-
 combineStats :: TLS (IORef StatsMap) -> IO StatsMapV
 combineStats statsRef = do
   threadStats' <- allTLS statsRef
   threadStats <-  mapM readIORef threadStats' :: IO [StatsMap]
-  statsLists <- pure $ map (\(m :: StatsMap) -> Map.map (\(x :: VU.IOVector StatsInfo) -> mvToList x) m) threadStats :: IO [Map EVM.Types.W256 (IO [StatsInfo])]
+  let statsLists = map (Map.map mvToList) threadStats :: [Map EVM.Types.W256 (IO [StatsInfo])]
   traverse (U.fromList <$>) $ Map.unionsWith zipSumStats statsLists
+  where
+    mvToList :: (VU.Unbox a) => VU.IOVector a -> IO [a]
+    mvToList = fmap U.toList . U.freeze
 
 saveCoverages
   :: Env
diff --git a/lib/Echidna/Types/Coverage.hs b/lib/Echidna/Types/Coverage.hs
index 81c61e37f..244b308f7 100644
--- a/lib/Echidna/Types/Coverage.hs
+++ b/lib/Echidna/Types/Coverage.hs
@@ -20,10 +20,13 @@ type CoverageMap = Map W256 (IOVector CoverageInfo)
 
 -- | Map with the statistic information needed for source code printing.
 -- Indexed by contracts' compile-time codehash; see `CodehashMap`.
+-- Used during runtime data collection
 type StatsMap = Map W256 (IOVector StatsInfo)
 
 -- | Map with the statistic information needed for source code printing.
 -- Indexed by contracts' compile-time codehash; see `CodehashMap`.
+-- Used during statistics summarization (combining multiple `StatsMap`)
+-- and coverage report generation.
 type StatsMapV = Map W256 (Vector StatsInfo)
 
 -- | Basic coverage information