From 924dfb220f5ae97fd8ffca55cd56322e96a67103 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           | 18 ++++++++----------
 lib/Echidna/Output/Source.hs  | 13 ++++++-------
 lib/Echidna/Types/Coverage.hs |  3 +++
 3 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/lib/Echidna/Exec.hs b/lib/Echidna/Exec.hs
index 7dc036f81..94aba192b 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)
@@ -287,23 +287,20 @@ execTxWithCov tx = do
       addCoverage !vm = do
         let (pc, opIx, depth) = currentCovLoc vm
             contract = currentContract vm
+            contractSize = BS.length . forceBuf . fromJust . view bytecode $ contract
 
         maybeCovVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp env.coverageRef $ do
-          let size = BS.length . forceBuf . fromJust . view bytecode $ contract
-          if size == 0 then pure Nothing else do
+          if contractSize == 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 contractSize (-1, 0, 0)
             pure $ Just vec
 
         statsRef <- getTLS env.statsRef
         maybeStatsVec <- lookupUsingCodehashOrInsert env.codehashMap contract env.dapp statsRef $ do
-          let size = BS.length . forceBuf . fromJust . view bytecode $ contract
-          if size == 0 then pure Nothing else do
+          if contractSize == 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 contractSize (0, 0)
             pure $ Just vec
 
         case maybeCovVec of
@@ -320,7 +317,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..d232d759a 100644
--- a/lib/Echidna/Output/Source.hs
+++ b/lib/Echidna/Output/Source.hs
@@ -9,7 +9,7 @@ import Data.ByteString qualified as BS
 import Data.Foldable
 import Data.IORef (readIORef, IORef)
 import Data.List (nub, sort)
-import Data.Maybe (fromMaybe, mapMaybe, isJust, fromJust)
+import Data.Maybe (fromMaybe, mapMaybe)
 import Data.Map (Map)
 import Data.Map qualified as Map
 import Data.Sequence qualified as Seq
@@ -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
@@ -199,8 +199,7 @@ srcMapCov sc covMap statMap contracts = do
                   updateLine (Just (r, q)) = Just ((<> unpackTxResults txResults) r, max q execQty)
                   updateLine Nothing = Just (unpackTxResults txResults, execQty)
                   fileStats = Map.lookup c.runtimeCodehash statMap
-                  idxStats | isJust fileStats = fromJust fileStats U.! opIx
-                           | otherwise        = (0, 0)
+                  idxStats = maybe (0, 0) (U.! opIx) fileStats
                   execQty = fst idxStats
                 Nothing -> acc
             Nothing -> acc
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