Use HashMap for ConcreteStorage

hevm.cabal

@@ -186,6 +186,7 @@ library
     split                             >= 0.2.3 && < 0.3,
     template-haskell                  >= 2.19.0 && < 3,
     extra                             >= 1.7.14 && < 2,
+    hashable                          >= 1.4.4.0 && < 2,
   hs-source-dirs:
     src
 
@@ -247,6 +248,7 @@ common test-base
     tasty-expected-failure,
     temporary,
     text,
+    unordered-containers,
     vector,
     witherable,
     operational,
@@ -290,6 +292,7 @@ test-suite test
     binary,
     data-dword,
     extra,
+    hashable,
     here,
     time,
     regex

src/EVM/Expr.hs

@@ -14,8 +14,9 @@ import Data.Bits hiding (And, Xor)
 import Data.ByteString (ByteString)
 import Data.ByteString qualified as BS
 import Data.DoubleWord (Int256, Word256(Word256), Word128(Word128))
+import Data.HashMap.Lazy qualified as LHashMap
+import Data.HashMap.Strict qualified as HashMap
 import Data.List
-import Data.Map qualified as LMap
 import Data.Map.Strict qualified as Map
 import Data.Maybe (mapMaybe, isJust, fromMaybe)
 import Data.Semigroup (Any, Any(..), getAny)
@@ -631,7 +632,7 @@ readStorage w st = go (simplifyNoLitToKeccak w) (simplifyNoLitToKeccak st)
     go :: Expr EWord -> Expr Storage -> Maybe (Expr EWord)
     go _ (GVar _) = internalError "Can't read from a GVar"
     go slot s@(AbstractStore _ _) = Just $ SLoad slot s
-    go (Lit l) (ConcreteStore s) = Lit <$> Map.lookup l s
+    go (Lit l) (ConcreteStore s) = Lit <$> HashMap.lookup l s
     go slot store@(ConcreteStore _) = Just $ SLoad slot store
     go slot s@(SStore prevSlot val prev) = case (prevSlot, slot) of
       -- if address and slot match then we return the val in this write
@@ -770,7 +771,7 @@ litToArrayPreimage val = go preImages
 -- ConcreteStore, otherwise we add a new write to the storage expression.
 writeStorage :: Expr EWord -> Expr EWord -> Expr Storage -> Expr Storage
 writeStorage k@(Lit key) v@(Lit val) store = case store of
-  ConcreteStore s -> ConcreteStore (Map.insert key val s)
+  ConcreteStore s -> ConcreteStore (HashMap.insert key val s)
   _ -> SStore k v store
 writeStorage key val store@(SStore key' val' prev)
      = if key == key'
@@ -948,12 +949,12 @@ decomposeStorage = go
       else setLogicalBase idx b
 
     -- empty concrete base is safe to reuse without any rewriting
-    setLogicalBase _ s@(ConcreteStore m) | Map.null m = Just s
+    setLogicalBase _ s@(ConcreteStore m) | HashMap.null m = Just s
 
     -- if the existing base is concrete but we have writes to only keys < 256
     -- then we can safely rewrite the base to an empty ConcreteStore (safe because we assume keccack(x) > 256)
     setLogicalBase _ (ConcreteStore store) =
-      if all (< 256) (Map.keys store)
+      if all (< 256) (HashMap.keys store)
       then Just (ConcreteStore mempty)
       else Nothing
     setLogicalBase _ (GVar _) = internalError "Unexpected GVar"
@@ -1758,7 +1759,7 @@ maybeLitAddrSimp e = case concKeccakSimpExpr e of
   LitAddr a -> Just a
   _ -> Nothing
 
-maybeConcStoreSimp :: Expr Storage -> Maybe (LMap.Map W256 W256)
+maybeConcStoreSimp :: Expr Storage -> Maybe (LHashMap.HashMap W256 W256)
 maybeConcStoreSimp (ConcreteStore s) = Just s
 maybeConcStoreSimp e = case concKeccakSimpExpr e of
   ConcreteStore s -> Just s

src/EVM/Format.hs

@@ -57,6 +57,7 @@ import Data.ByteString.Lazy (toStrict, fromStrict)
 import Data.Char qualified as Char
 import Data.DoubleWord (signedWord)
 import Data.Foldable (toList)
+import Data.HashMap.Strict qualified as HashMap
 import Data.List (isPrefixOf, sort)
 import Data.Map qualified as Map
 import Data.Maybe (catMaybes, fromMaybe, listToMaybe)
@@ -738,7 +739,7 @@ formatExpr = go
           [ "(ConcreteStore"
           , indent 2 $ T.unlines
             [ "vals:"
-            , indent 2 $ T.unlines $ fmap (T.pack . show) $ Map.toList s
+            , indent 2 $ T.unlines $ fmap (T.pack . show) $ HashMap.toList s
             ]
           , ")"
           ]

src/EVM/Fuzz.hs

@@ -7,7 +7,8 @@ module EVM.Fuzz where
 
 import Control.Monad (replicateM)
 import Control.Monad.State (State, get, put, execState)
-import Data.Map.Strict as Map (fromList, Map, (!), (!?), insert)
+import Data.HashMap.Strict as HashMap (fromList, HashMap)
+import Data.Map.Strict as Map (Map, (!), (!?), insert)
 import Data.Maybe (fromMaybe)
 import Data.Set as Set (insert, Set, empty, toList, fromList)
 import Data.ByteString qualified as BS
@@ -76,7 +77,7 @@ substituteBuf valMap p = mapProp go p
                                 Nothing -> orig
     go a = a
 
-substituteStores ::  Map (Expr 'EAddr) (Map W256 W256) -> Prop -> Prop
+substituteStores ::  Map (Expr 'EAddr) (HashMap W256 W256) -> Prop -> Prop
 substituteStores valMap p = mapProp go p
   where
     go :: Expr a -> Expr a
@@ -216,10 +217,10 @@ getvals vars = do
       go ax (Map.insert a val valMap)
 
 -- Storage value generation
-getStores :: CollectStorage -> Gen (Map (Expr EAddr) (Map W256 W256))
+getStores :: CollectStorage -> Gen (Map (Expr EAddr) (HashMap W256 W256))
 getStores storesLoads = go (Set.toList storesLoads.addrs) mempty
   where
-    go :: [(Expr EAddr, Maybe W256)] -> Map (Expr EAddr) (Map W256 W256) -> Gen (Map (Expr EAddr) (Map W256 W256))
+    go :: [(Expr EAddr, Maybe W256)] -> Map (Expr EAddr) (HashMap W256 W256) -> Gen (Map (Expr EAddr) (HashMap W256 W256))
     go [] addrToValsMap = pure addrToValsMap
     go ((addr, _):ax) addrToValsMap = do
       -- number of elements inserted into storage
@@ -228,7 +229,7 @@ getStores storesLoads = go (Set.toList storesLoads.addrs) mempty
                                    ,(1, choose (11, 100))
                                    ]
       l <- replicateM numElems oneWrite
-      go ax (Map.insert addr (Map.fromList l) addrToValsMap)
+      go ax (Map.insert addr (HashMap.fromList l) addrToValsMap)
         where
           oneWrite :: Gen (W256, W256)
           oneWrite = do

src/EVM/SMT.hs

@@ -19,6 +19,8 @@ import Data.String.Here
 import Data.Maybe (fromJust, fromMaybe, isJust, listToMaybe)
 import Data.Either.Extra (fromRight')
 import Data.Foldable (fold)
+import Data.HashMap.Strict (HashMap)
+import Data.HashMap.Strict qualified as HashMap
 import Data.Map.Strict (Map)
 import Data.Map.Strict qualified as Map
 import Data.Set (Set)
@@ -928,8 +930,8 @@ writeBytes bytes buf =  do
           idxSMT <- exprToSMT . Lit . unsafeInto $ idx
           pure (idx + 1, "(store " <> inner `sp` idxSMT `sp` byteSMT <> ")")
 
-encodeConcreteStore :: Map W256 W256 -> Err Builder
-encodeConcreteStore s = foldM encodeWrite ("((as const Storage) #x0000000000000000000000000000000000000000000000000000000000000000)") (Map.toList s)
+encodeConcreteStore :: HashMap W256 W256 -> Err Builder
+encodeConcreteStore s = foldM encodeWrite ("((as const Storage) #x0000000000000000000000000000000000000000000000000000000000000000)") (HashMap.toList s)
   where
     encodeWrite :: Builder -> (W256, W256) -> Err Builder
     encodeWrite prev (key, val) = do
@@ -1109,7 +1111,7 @@ getBufs getVal bufs = foldM getBuf mempty bufs
 getStore
   :: (Text -> IO Text)
   -> Map (Expr EAddr, Maybe W256) (Set (Expr EWord))
-  -> IO (Map (Expr EAddr) (Map W256 W256))
+  -> IO (Map (Expr EAddr) (HashMap W256 W256))
 getStore getVal abstractReads =
   fmap Map.fromList $ forM (Map.toList abstractReads) $ \((addr, idx), slots) -> do
     let name = toLazyText (storeName addr idx)
@@ -1128,7 +1130,7 @@ getStore getVal abstractReads =
     -- then create a map by adding only the locations that are read by the program
     store <- foldM (\m slot -> do
       slot' <- queryValue getVal slot
-      pure $ Map.insert slot' (fun slot') m) Map.empty slots
+      pure $ HashMap.insert slot' (fun slot') m) HashMap.empty slots
     pure (addr, store)
 
 -- | Ask the solver to give us the concrete value of an arbitrary abstract word

src/EVM/SymExec.hs

@@ -15,6 +15,8 @@ import Data.ByteString (ByteString)
 import Data.ByteString qualified as BS
 import Data.Containers.ListUtils (nubOrd)
 import Data.DoubleWord (Word256)
+import Data.HashMap.Strict (HashMap)
+import Data.HashMap.Strict qualified as HashMap
 import Data.List (foldl', sortBy, sort)
 import Data.List.NonEmpty qualified as NE
 import Data.Maybe (fromMaybe, listToMaybe, mapMaybe)
@@ -982,7 +984,7 @@ equivalenceCheck' solvers branchesA branchesB create = do
         -- TODO: is this sound? do we need a more sophisticated nonce representation?
         noncesDiffer = PBool (ac.nonce /= bc.nonce)
         storesDiffer = case (ac.storage, bc.storage) of
-          (ConcreteStore as, ConcreteStore bs) | not (as == Map.empty || bs == Map.empty) -> PBool $ as /= bs
+          (ConcreteStore as, ConcreteStore bs) | not (as == HashMap.empty || bs == HashMap.empty) -> PBool $ as /= bs
           (as, bs) -> if as == bs then PBool False else as ./= bs
       in balsDiffer .|| storesDiffer .|| noncesDiffer
 
@@ -1057,7 +1059,7 @@ formatCex cd sig m@(SMTCex _ addrs _ store blockContext txContext) = T.unlines $
           [ "Storage:"
           , indent 2 $ T.unlines $ Map.foldrWithKey (\key val acc ->
               ("Addr " <> (T.pack . show $ key)
-                <> ": " <> (T.pack $ show (Map.toList val))) : acc
+                <> ": " <> (T.pack $ show (HashMap.toList val))) : acc
             ) mempty store
           ]
 
@@ -1223,10 +1225,10 @@ subBufs model b = Map.foldlWithKey subBuf (Right b) model
           Just k -> forceFlattened k
           Nothing -> Left $ show buf
 
-subStores :: Map (Expr EAddr) (Map W256 W256) -> Expr a -> Expr a
+subStores :: Map (Expr EAddr) (HashMap W256 W256) -> Expr a -> Expr a
 subStores model b = Map.foldlWithKey subStore b model
   where
-    subStore :: Expr a -> Expr EAddr -> Map W256 W256 -> Expr a
+    subStore :: Expr a -> Expr EAddr -> HashMap W256 W256 -> Expr a
     subStore x var val = mapExpr go x
       where
         go :: Expr a -> Expr a

src/EVM/Types.hs

@@ -36,6 +36,8 @@ import Data.Int (Int64)
 import Data.Word (Word8, Word32, Word64)
 import Data.DoubleWord
 import Data.DoubleWord.TH
+import Data.Hashable (Hashable, hashWithSalt)
+import Data.HashMap.Strict (HashMap)
 import Data.Map (Map)
 import Data.Map qualified as Map
 import Data.Maybe (fromMaybe)
@@ -333,7 +335,7 @@ data Expr (a :: EType) where
 
   -- storage
 
-  ConcreteStore  :: (Map W256 W256) -> Expr Storage
+  ConcreteStore  :: (HashMap W256 W256) -> Expr Storage
   AbstractStore  :: Expr EAddr -- which contract is this store for?
                  -> Maybe W256 -- which logical store does this refer to? (e.g. solidity mappings / arrays)
                  -> Expr Storage
@@ -1140,7 +1142,7 @@ data SMTCex = SMTCex
   { vars :: Map (Expr EWord) W256
   , addrs :: Map (Expr EAddr) Addr
   , buffers :: Map (Expr Buf) BufModel
-  , store :: Map (Expr EAddr) (Map W256 W256)
+  , store :: Map (Expr EAddr) (HashMap W256 W256)
   , blockContext :: Map (Expr EWord) W256
   , txContext :: Map (Expr EWord) W256
   }
@@ -1299,6 +1301,8 @@ instance ParseFields W256
 instance ParseRecord W256 where
   parseRecord = fmap getOnly parseRecord
 
+instance Hashable W256 where
+    hashWithSalt s (W256 w) = s `hashWithSalt` w
 
 -- Word64 wrapper ----------------------------------------------------------------------------------
 

test/EVM/Test/BlockchainTests.hs

@@ -25,6 +25,8 @@ import Data.Aeson.Types qualified as JSON
 import Data.ByteString qualified as BS
 import Data.ByteString.Lazy qualified as Lazy
 import Data.ByteString.Lazy qualified as LazyByteString
+import Data.HashMap.Strict (HashMap)
+import Data.HashMap.Strict qualified as HashMap
 import Data.List (isInfixOf, isPrefixOf)
 import Data.Map (Map)
 import Data.Map qualified as Map
@@ -229,7 +231,7 @@ splitEithers =
   . (fmap fst &&& fmap snd)
   . (fmap (preview _Left &&& preview _Right))
 
-fromConcrete :: Expr Storage -> Map W256 W256
+fromConcrete :: Expr Storage -> HashMap W256 W256
 fromConcrete (ConcreteStore s) = s
 fromConcrete s = internalError $ "unexpected abstract store: " <> show s
 
@@ -308,7 +310,7 @@ clearOrigStorage = set #origStorage (ConcreteStore mempty)
 
 clearZeroStorage :: Contract -> Contract
 clearZeroStorage c = case c.storage of
-  ConcreteStore m -> let store = Map.filter (/= 0) m
+  ConcreteStore m -> let store = HashMap.filter (/= 0) m
                      in set #storage (ConcreteStore store) c
   _ -> internalError "Internal Error: unexpected abstract store"
 

test/rpc.hs

@@ -6,6 +6,7 @@ import Test.Tasty
 import Test.Tasty.HUnit
 
 import Data.Maybe
+import Data.HashMap.Strict qualified as HashMap
 import Data.Map qualified as Map
 import Data.Text (Text)
 import Data.Vector qualified as V
@@ -89,7 +90,7 @@ tests = testGroup "rpc"
           wethStore' = case wethStore of
             ConcreteStore s -> s
             _ -> internalError "Expecting concrete store"
-          receiverBal = fromJust $ Map.lookup (keccak' (word256Bytes 0xdead <> word256Bytes 0x3)) wethStore'
+          receiverBal = fromJust $ HashMap.lookup (keccak' (word256Bytes 0xdead <> word256Bytes 0x3)) wethStore'
           msg = case postVm.result of
             Just (VMSuccess m) -> m
             _ -> internalError "VMSuccess expected"