From 4eb2f33ec3accdf9fabdb187b7cec853cdc83d4a Mon Sep 17 00:00:00 2001
From: Sam Alws <sam.alws@trailofbits.com>
Date: Mon, 9 Sep 2024 10:57:37 -0400
Subject: [PATCH] cleanup

---
 lib/Echidna/Types/Coverage.hs | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/lib/Echidna/Types/Coverage.hs b/lib/Echidna/Types/Coverage.hs
index a09292bc9..f5fff4a58 100644
--- a/lib/Echidna/Types/Coverage.hs
+++ b/lib/Echidna/Types/Coverage.hs
@@ -1,5 +1,6 @@
 module Echidna.Types.Coverage where
 
+import Control.Monad ((>=>))
 import Data.Aeson (ToJSON(toJSON), FromJSON(parseJSON), withText)
 import Data.Bits (testBit)
 import Data.IORef (IORef, readIORef)
@@ -22,6 +23,8 @@ import Echidna.Types.Tx (TxResult)
 -- Indexed by contracts' compile-time codehash; see `CodehashMap`.
 type CoverageMap = Map W256 (IOVector CoverageInfo)
 
+-- | CoverageMap, but using Vectors instead of IOVectors.
+-- IO is not required to access this map's members.
 type FrozenCoverageMap = Map W256 (V.Vector CoverageInfo)
 
 -- | Basic coverage information
@@ -36,17 +39,28 @@ type StackDepths = Word64
 -- | Packed TxResults used for coverage, corresponding bits are set
 type TxResults = Word64
 
+-- | Given the CoverageMaps used for contract init and runtime, produce a single combined coverage map
+-- with op indices from init correctly shifted over (see srcMapForOpLocation in Echidna.Output.Source).
+-- Takes IORef CoverageMap because this is how they are stored in the Env.
 mergeCoverageMaps :: DappInfo -> IORef CoverageMap -> IORef CoverageMap -> IO FrozenCoverageMap
-mergeCoverageMaps dapp initMap runtimeMap = do
-  initMap' <- Map.mapWithKey modifyInitMapEntry <$> (mapM V.freeze =<< readIORef initMap)
-  runtimeMap' <- mapM V.freeze =<< readIORef runtimeMap
-  pure $ Map.unionWith (<>) runtimeMap' initMap'
+mergeCoverageMaps dapp initMap runtimeMap = mergeFrozenCoverageMaps dapp <$> freeze initMap <*> freeze runtimeMap
+  where freeze = readIORef >=> mapM V.freeze
+
+-- | Given the FrozenCoverageMaps used for contract init and runtime, produce a single combined coverage map
+-- with op indices from init correctly shifted over (see srcMapForOpLocation in Echidna.Output.Source).
+-- Helper function for mergeCoverageMaps.
+mergeFrozenCoverageMaps :: DappInfo -> FrozenCoverageMap -> FrozenCoverageMap -> FrozenCoverageMap
+mergeFrozenCoverageMaps dapp initMap runtimeMap = Map.unionWith (<>) runtimeMap initMap'
   where
+    initMap' = Map.mapWithKey modifyInitMapEntry initMap
     -- eta reduced, second argument is a vec
     modifyInitMapEntry hash = V.map $ modifyCoverageInfo $ getOpOffset hash
     modifyCoverageInfo toAdd (op, x, y) = (op + toAdd, x, y)
     getOpOffset hash = maybe 0 (length . (.runtimeSrcmap) . snd) $ Map.lookup hash dapp.solcByHash
 
+-- | Given the CoverageMaps used for contract init and runtime,
+-- return the point coverage and the number of unique contracts hit.
+-- Takes IORef CoverageMap because this is how they are stored in the Env.
 coverageStats :: IORef CoverageMap -> IORef CoverageMap -> IO (Int, Int)
 coverageStats initRef runtimeRef = do
   initMap <- readIORef initRef