From d6522c509ed71c4b8da6e61bcabd50191e27f896 Mon Sep 17 00:00:00 2001
From: dxo <hi@d-xo.org>
Date: Mon, 8 Jan 2024 16:21:51 +0100
Subject: [PATCH 01/20] Types: Add logical storage idx to AbstractStore

---
 src/EVM.hs            | 13 +++++++------
 src/EVM/Expr.hs       | 10 ++++++++--
 src/EVM/Format.hs     |  4 ++--
 src/EVM/Fuzz.hs       | 13 +++++++------
 src/EVM/SMT.hs        | 27 +++++++++++++++------------
 src/EVM/SymExec.hs    |  2 +-
 src/EVM/Traversals.hs |  4 ++--
 src/EVM/Types.hs      |  4 +++-
 test/test.hs          |  6 +++---
 9 files changed, 48 insertions(+), 35 deletions(-)

diff --git a/src/EVM.hs b/src/EVM.hs
index 63f7a0b0e..2cd7895ef 100644
--- a/src/EVM.hs
+++ b/src/EVM.hs
@@ -167,8 +167,8 @@ makeVm o = do
 unknownContract :: Expr EAddr -> Contract
 unknownContract addr = Contract
   { code        = UnknownCode addr
-  , storage     = AbstractStore addr
-  , origStorage = AbstractStore addr
+  , storage     = AbstractStore addr Nothing
+  , origStorage = AbstractStore addr Nothing
   , balance     = Balance addr
   , nonce       = Nothing
   , codehash    = hashcode (UnknownCode addr)
@@ -181,8 +181,8 @@ unknownContract addr = Contract
 abstractContract :: ContractCode -> Expr EAddr -> Contract
 abstractContract code addr = Contract
   { code        = code
-  , storage     = AbstractStore addr
-  , origStorage = AbstractStore addr
+  , storage     = AbstractStore addr Nothing
+  , origStorage = AbstractStore addr Nothing
   , balance     = Balance addr
   , nonce       = if isCreation code then Just 1 else Just 0
   , codehash    = hashcode code
@@ -1895,9 +1895,10 @@ create self this xSize xGas xValue xs newAddr initCode = do
               resetStorage :: Expr Storage -> Expr Storage
               resetStorage = \case
                   ConcreteStore _ -> ConcreteStore mempty
-                  AbstractStore a -> AbstractStore a
+                  AbstractStore a Nothing -> AbstractStore a Nothing
                   SStore _ _ p -> resetStorage p
-                  GVar _  -> error "unexpected global variable"
+                  AbstractStore _ (Just _) -> internalError "unexpected logical store in EVM.hs"
+                  GVar _  -> internalError "unexpected global variable"
 
             modifying (#env % #contracts % ix newAddr % #storage) resetStorage
             modifying (#env % #contracts % ix newAddr % #origStorage) resetStorage
diff --git a/src/EVM/Expr.hs b/src/EVM/Expr.hs
index 4f088471c..0df37f6b6 100644
--- a/src/EVM/Expr.hs
+++ b/src/EVM/Expr.hs
@@ -612,7 +612,7 @@ readStorage w st = go (simplify w) st
   where
     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 slot s@(AbstractStore _ _) = Just $ SLoad slot s
     go (Lit l) (ConcreteStore s) = Lit <$> Map.lookup l s
     go slot store@(ConcreteStore _) = Just $ SLoad slot store
     go slot s@(SStore prevSlot val prev) = case (prevSlot, slot) of
@@ -756,10 +756,16 @@ writeStorage key val store = SStore key val store
 
 getAddr :: Expr Storage -> Maybe (Expr EAddr)
 getAddr (SStore _ _ p) = getAddr p
-getAddr (AbstractStore a) = Just a
+getAddr (AbstractStore a _) = Just a
 getAddr (ConcreteStore _) = Nothing
 getAddr (GVar _) = internalError "cannot determine addr of a GVar"
 
+getLogicalIdx :: Expr Storage -> Maybe W256
+getLogicalIdx (SStore _ _ p) = getLogicalIdx p
+getLogicalIdx (AbstractStore _ idx) = idx
+getLogicalIdx (ConcreteStore _) = Nothing
+getLogicalIdx (GVar _) = internalError "cannot determine addr of a GVar"
+
 
 -- ** Whole Expression Simplification ** -----------------------------------------------------------
 
diff --git a/src/EVM/Format.hs b/src/EVM/Format.hs
index 19bbbade1..468db4423 100644
--- a/src/EVM/Format.hs
+++ b/src/EVM/Format.hs
@@ -673,8 +673,8 @@ formatExpr = go
         , ")"
         , formatExpr prev
         ]
-      AbstractStore a ->
-        "(AbstractStore " <> formatExpr a <> ")"
+      AbstractStore a idx ->
+        "(AbstractStore " <> formatExpr a <> " " <> T.pack (show idx) <> ")"
       ConcreteStore s -> if null s
         then "(ConcreteStore <empty>)"
         else T.unlines
diff --git a/src/EVM/Fuzz.hs b/src/EVM/Fuzz.hs
index 0da0dcacc..304067ff5 100644
--- a/src/EVM/Fuzz.hs
+++ b/src/EVM/Fuzz.hs
@@ -81,7 +81,8 @@ substituteStores ::  Map (Expr 'EAddr) (Map W256 W256) -> Prop -> Prop
 substituteStores valMap p = mapProp go p
   where
     go :: Expr a -> Expr a
-    go (AbstractStore a) = case valMap !? a of
+    -- TODO: Is this OK??
+    go (AbstractStore a _) = case valMap !? a of
                                   Just m -> ConcreteStore m
                                   Nothing -> ConcreteStore mempty
     go a = a
@@ -141,7 +142,7 @@ findBufProp p = mapPropM go p
       e -> pure e
 
 --- Store extraction
-data CollectStorage = CollectStorage { addrs :: Set.Set (Expr EAddr)
+data CollectStorage = CollectStorage { addrs :: Set.Set (Expr EAddr, Maybe W256)
                                      , keys :: Set.Set W256
                                      , vals :: Set.Set W256
                                      }
@@ -181,9 +182,9 @@ findStoragePropInner p = mapPropM go p
   where
     go :: forall a. Expr a -> State CollectStorage (Expr a)
     go = \case
-      e@(AbstractStore a) -> do
+      e@(AbstractStore a idx) -> do
         s <- get
-        put s {addrs=Set.insert a s.addrs}
+        put s {addrs=Set.insert (a, idx) s.addrs}
         pure e
       e@(SLoad (Lit val) _) -> do
         s <- get
@@ -219,9 +220,9 @@ getvals vars = do
 getStores :: CollectStorage -> Gen (Map (Expr EAddr) (Map W256 W256))
 getStores storesLoads = go (Set.toList storesLoads.addrs) mempty
   where
-    go :: [Expr EAddr] -> Map (Expr EAddr) (Map W256 W256) -> Gen (Map (Expr EAddr) (Map W256 W256))
+    go :: [(Expr EAddr, Maybe W256)] -> Map (Expr EAddr) (Map W256 W256) -> Gen (Map (Expr EAddr) (Map W256 W256))
     go [] addrToValsMap = pure addrToValsMap
-    go (addr:ax) addrToValsMap = do
+    go ((addr, _):ax) addrToValsMap = do
       -- number of elements inserted into storage
       numElems :: Int <- frequency [(1, pure 0)
                                    ,(10, choose (1, 10))
diff --git a/src/EVM/SMT.hs b/src/EVM/SMT.hs
index acc01d21d..96b60b231 100644
--- a/src/EVM/SMT.hs
+++ b/src/EVM/SMT.hs
@@ -55,7 +55,7 @@ data CexVars = CexVars
     -- | buffer names and guesses at their maximum size
   , buffers      :: Map Text (Expr EWord)
     -- | reads from abstract storage
-  , storeReads   :: Map (Expr EAddr) (Set (Expr EWord))
+  , storeReads   :: Map (Expr EAddr, Maybe W256) (Set (Expr EWord))
     -- | the names of any block context variables
   , blockContext :: [Text]
     -- | the names of any tx context variables
@@ -290,7 +290,7 @@ referencedAbstractStores :: TraversableTerm a => a -> Set Builder
 referencedAbstractStores term = foldTerm go mempty term
   where
     go = \case
-      AbstractStore s -> Set.singleton (storeName s)
+      AbstractStore s idx -> Set.singleton (storeName s idx)
       _ -> mempty
 
 referencedWAddrs :: TraversableTerm a => a -> Set Builder
@@ -347,16 +347,16 @@ referencedBlockContext expr = nubOrd $ foldTerm go [] expr
 -- the store (e.g, SLoad addr idx (SStore addr idx val AbstractStore)).
 -- However, we expect that most of such reads will have been
 -- simplified away.
-findStorageReads :: Prop -> Map (Expr EAddr) (Set (Expr EWord))
+findStorageReads :: Prop -> Map (Expr EAddr, Maybe W256) (Set (Expr EWord))
 findStorageReads p = Map.fromListWith (<>) $ foldProp go mempty p
   where
-    go :: Expr a -> [(Expr EAddr, Set (Expr EWord))]
+    go :: Expr a -> [((Expr EAddr, Maybe W256), Set (Expr EWord))]
     go = \case
       SLoad slot store ->
-        [((fromMaybe (error $ "Internal Error: could not extract address from: " <> show store) (Expr.getAddr store)), Set.singleton slot) | containsNode isAbstractStore store]
+        [((fromMaybe (internalError $ "could not extract address from: " <> show store) (Expr.getAddr store), Expr.getLogicalIdx store), Set.singleton slot) | containsNode isAbstractStore store]
       _ -> []
 
-    isAbstractStore (AbstractStore _) = True
+    isAbstractStore (AbstractStore _ _) = True
     isAbstractStore _ = False
 
 findBufferAccess :: TraversableTerm a => [a] -> [(Expr EWord, Expr EWord, Expr Buf)]
@@ -815,8 +815,10 @@ exprToSMT = \case
     "(writeWord " <> encIdx `sp` encVal `sp` encPrev <> ")"
   CopySlice srcIdx dstIdx size src dst ->
     copySlice srcIdx dstIdx size (exprToSMT src) (exprToSMT dst)
+
+  -- we need to do a bit of processing here.
   ConcreteStore s -> encodeConcreteStore s
-  AbstractStore a -> storeName a
+  AbstractStore a idx -> storeName a idx
   SStore idx val prev ->
     let encIdx = exprToSMT idx
         encVal = exprToSMT val
@@ -936,8 +938,9 @@ encodeConcreteStore s = foldl encodeWrite "((as const Storage) #x000000000000000
         encVal = exprToSMT (Lit val)
       in "(store " <> prev `sp` encKey `sp` encVal <> ")"
 
-storeName :: Expr EAddr -> Builder
-storeName a = fromString ("baseStore_") <> formatEAddr a
+storeName :: Expr EAddr -> Maybe W256 -> Builder
+storeName a Nothing = fromString ("baseStore_") <> formatEAddr a
+storeName a (Just idx) = fromString ("baseStore_") <> formatEAddr a <> "_" <> (fromString $ show idx)
 
 formatEAddr :: Expr EAddr -> Builder
 formatEAddr = \case
@@ -1097,11 +1100,11 @@ getBufs getVal bufs = foldM getBuf mempty bufs
 -- concretized storage
 getStore
   :: (Text -> IO Text)
-  -> Map (Expr EAddr) (Set (Expr EWord))
+  -> Map (Expr EAddr, Maybe W256) (Set (Expr EWord))
   -> IO (Map (Expr EAddr) (Map W256 W256))
 getStore getVal abstractReads =
-  fmap Map.fromList $ forM (Map.toList abstractReads) $ \(addr, slots) -> do
-    let name = toLazyText (storeName addr)
+  fmap Map.fromList $ forM (Map.toList abstractReads) $ \((addr, idx), slots) -> do
+    let name = toLazyText (storeName addr idx)
     raw <- getVal name
     let parsed = case parseCommentFreeFileMsg getValueRes (T.toStrict raw) of
                    Right (ResSpecific (valParsed :| [])) -> valParsed
diff --git a/src/EVM/SymExec.hs b/src/EVM/SymExec.hs
index 5bc342dd7..8e8984626 100644
--- a/src/EVM/SymExec.hs
+++ b/src/EVM/SymExec.hs
@@ -1000,7 +1000,7 @@ subStores model b = Map.foldlWithKey subStore b model
       where
         go :: Expr a -> Expr a
         go = \case
-          v@(AbstractStore a)
+          v@(AbstractStore a _)
             -> if a == var
                then ConcreteStore val
                else v
diff --git a/src/EVM/Traversals.hs b/src/EVM/Traversals.hs
index d4c950f79..3535a7ee3 100644
--- a/src/EVM/Traversals.hs
+++ b/src/EVM/Traversals.hs
@@ -185,7 +185,7 @@ foldExpr f acc expr = acc <> (go expr)
       -- storage
 
       e@(ConcreteStore _) -> f e
-      e@(AbstractStore _) -> f e
+      e@(AbstractStore _ _) -> f e
       e@(SLoad a b) -> f e <> (go a) <> (go b)
       e@(SStore a b c) -> f e <> (go a) <> (go b) <> (go c)
 
@@ -512,7 +512,7 @@ mapExprM f expr = case expr of
   -- storage
 
   ConcreteStore b -> f (ConcreteStore b)
-  AbstractStore a -> f (AbstractStore a)
+  AbstractStore a b -> f (AbstractStore a b)
   SLoad a b -> do
     a' <- mapExprM f a
     b' <- mapExprM f b
diff --git a/src/EVM/Types.hs b/src/EVM/Types.hs
index 008075ad7..f113b4197 100644
--- a/src/EVM/Types.hs
+++ b/src/EVM/Types.hs
@@ -326,7 +326,9 @@ data Expr (a :: EType) where
   -- storage
 
   ConcreteStore  :: (Map W256 W256) -> Expr Storage
-  AbstractStore  :: Expr EAddr -> 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
 
   SLoad          :: Expr EWord         -- key
                  -> Expr Storage       -- storage
diff --git a/test/test.hs b/test/test.hs
index 733049ddc..e02e5e353 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -285,7 +285,7 @@ tests = testGroup "hevm"
   , testGroup "StorageTests"
     [ test "read-from-sstore" $ assertEqualM ""
         (Lit 0xab)
-        (Expr.readStorage' (Lit 0x0) (SStore (Lit 0x0) (Lit 0xab) (AbstractStore (LitAddr 0x0))))
+        (Expr.readStorage' (Lit 0x0) (SStore (Lit 0x0) (Lit 0xab) (AbstractStore (LitAddr 0x0) Nothing)))
     , test "read-from-concrete" $ assertEqualM ""
         (Lit 0xab)
         (Expr.readStorage' (Lit 0x0) (ConcreteStore $ Map.fromList [(0x0, 0xab)]))
@@ -3742,7 +3742,7 @@ genStorageWrites :: Gen (Expr Storage)
 genStorageWrites = do
   toSlot <- genSlot
   val <- genLit (maxBound :: W256)
-  store <- frequency [ (3, pure $ AbstractStore (SymAddr ""))
+  store <- frequency [ (3, pure $ AbstractStore (SymAddr "") Nothing)
                      , (2, genStorageWrites)
                      ]
   pure $ SStore toSlot val store
@@ -4031,7 +4031,7 @@ genBuf bound sz = oneof
 
 genStorage :: Int -> Gen (Expr Storage)
 genStorage 0 = oneof
-  [ fmap AbstractStore arbitrary
+  [ liftM2 AbstractStore arbitrary arbitrary
   , fmap ConcreteStore arbitrary
   ]
 genStorage sz = liftM3 SStore subWord subWord subStore

From 4aa4ad4865bfd84c0a6830b5fb81ec0143525d1a Mon Sep 17 00:00:00 2001
From: dxo <hi@d-xo.org>
Date: Mon, 8 Jan 2024 17:38:45 +0100
Subject: [PATCH 02/20] Expr: rewrite logical base stores

---
 src/EVM/Expr.hs | 70 +++++++++++++++++++++++++++++++++++++++++++++++++
 src/EVM/SMT.hs  |  9 ++++++-
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/src/EVM/Expr.hs b/src/EVM/Expr.hs
index 0df37f6b6..b6bbb7ad8 100644
--- a/src/EVM/Expr.hs
+++ b/src/EVM/Expr.hs
@@ -7,6 +7,8 @@
 -}
 module EVM.Expr where
 
+import Debug.Trace
+
 import Prelude hiding (LT, GT)
 import Control.Monad.ST
 import Data.Bits hiding (And, Xor)
@@ -770,6 +772,74 @@ getLogicalIdx (GVar _) = internalError "cannot determine addr of a GVar"
 -- ** Whole Expression Simplification ** -----------------------------------------------------------
 
 
+-- | Splits storage into logical substores (solc based heuristic) if it safe to do so
+-- safe if
+--   - all base stores must be one of:
+--     - concrete + empty
+--     - concrete + all keys < 256
+--     - abstract
+--   - all read / write indices must be one of:
+--     - Lit < 256
+--     - MappingSlot
+--     - ArraySlotWithOffset
+--     - ArraySlotZero
+--
+decomposeStorage :: Expr a -> Maybe (Expr a)
+decomposeStorage = mapExprM go
+  where
+    go :: Expr a -> Maybe (Expr a)
+    go (SLoad idx@(Keccak _) store) = case inferLogicalIdx idx of
+      Just (logicalIdx, slot) -> do
+        base <- setLogicalBase logicalIdx store
+        pure (SLoad slot base)
+      Nothing -> trace "1" Nothing
+    go e@(SLoad (Lit idx) _) | idx < 256 = Just e
+    go (SLoad _ _) = trace "2" Nothing
+
+    go s@(SStore idx@(Keccak _) val store) = trace "INSIDE SSTORE" $ case inferLogicalIdx idx of
+      Just (logicalIdx, slot) -> do
+        traceM "DOING A REWRITE"
+        traceShowM s
+        base <- setLogicalBase logicalIdx store
+        pure (SStore slot val base)
+      Nothing -> trace "3" Nothing
+    go e@(SStore (Lit idx) _ _) | idx < 256 = Just e
+    go s@(SStore _ _ _) = trace (show s) Nothing
+
+    go e = Just e
+
+    inferLogicalIdx :: Expr EWord -> Maybe (W256, Expr EWord)
+    inferLogicalIdx = \case
+      (MappingSlot idx key) -> Just (idxToWord idx, key)
+      (ArraySlotWithOffset idx offset) -> Just (idxToWord idx, offset)
+      (ArraySlotZero idx) -> Just (idxToWord idx, Lit 0)
+      _ -> trace "5" Nothing
+
+    idxToWord :: ByteString -> W256
+    idxToWord = W256 . word256 . BS.takeEnd 32
+
+    -- Updates the logical base store of the given expression if it is safe to do so
+    setLogicalBase :: W256 -> Expr Storage -> Maybe (Expr Storage)
+
+    -- abstract bases get their logical idx set to the new value
+    setLogicalBase new (AbstractStore addr _) = Just $ AbstractStore addr (Just new)
+    setLogicalBase new (SStore i v s) = do
+      b <- setLogicalBase new s
+      Just $ SStore i v b
+
+    -- empty concrete base is safe to reuse without any rewriting
+    setLogicalBase _ s@(ConcreteStore m) | Map.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)
+      then Just (ConcreteStore mempty)
+      else trace "6" Nothing
+    setLogicalBase _ (GVar _) = internalError "Unexpected GVar"
+
+
+
 -- | Simple recursive match based AST simplification
 -- Note: may not terminate!
 simplify :: Expr a -> Expr a
diff --git a/src/EVM/SMT.hs b/src/EVM/SMT.hs
index 96b60b231..73ac19704 100644
--- a/src/EVM/SMT.hs
+++ b/src/EVM/SMT.hs
@@ -9,6 +9,8 @@ module EVM.SMT where
 
 import Prelude hiding (LT, GT)
 
+import Debug.Trace
+
 import Control.Monad
 import Data.Containers.ListUtils (nubOrd)
 import Data.ByteString (ByteString)
@@ -207,7 +209,12 @@ assertProps :: Config -> [Prop] -> SMT2
 -- simplify to rewrite sload/sstore combos
 -- notice: it is VERY important not to concretize, because Keccak assumptions
 --         will be generated by assertPropsNoSimp, and that needs unconcretized Props
-assertProps conf ps = assertPropsNoSimp conf $ Expr.simplifyProps ps
+assertProps conf ps = assertPropsNoSimp conf withSepStore
+  where
+    simplifiedProps = trace "YO" $ Expr.simplifyProps ps
+    withSepStore = trace "OI" $ case mapM (mapPropM Expr.decomposeStorage) simplifiedProps of
+                     Just ps' -> trace "REWROTE" ps'
+                     Nothing -> trace "NO REWRITE" simplifiedProps
 
 -- Note: we need a version that does NOT call simplify,
 -- because we make use of it to verify the correctness of our simplification

From 2add0a0f866daa66d4195c560be35804aa31ff89 Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Mon, 8 Jan 2024 18:06:48 +0100
Subject: [PATCH 03/20] Adding test to quickly debug

---
 test/test.hs | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/test/test.hs b/test/test.hs
index e02e5e353..62dfecd49 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -205,6 +205,23 @@ tests = testGroup "hevm"
         |]
        expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "transfer(uint256,uint256)" [AbiUIntType 256, AbiUIntType 256])) [] (defaultVeriOpts { maxIter = Just 5 })
        assertEqualM "Expression is not clean." (badStoresInExpr expr) False
+    , test "different-addr" $ do
+      Just c <- solcRuntime "MyContract"
+        [i|
+        contract MyContract {
+          mapping (address => uint) balances;
+          mapping (uint => bool) auth;
+          mapping (address => mapping (address => uint)) allowance;
+          function prove_mapping_access(address x, address y) public {
+              require(x != y);
+              balances[x] = 1;
+              balances[y] = 2;
+              assert(balances[x] != balances[y]);
+          }
+        }
+        |]
+      expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mapping_access(address,address)" [AbiAddressType, AbiAddressType])) [] defaultVeriOpts
+      putStrLnM $ T.unpack $ formatExpr expr
     , test "simplify-storage-map-only-static" $ do
        Just c <- solcRuntime "MyContract"
         [i|

From 909f42f3726ac1516b71d0609f9d2ae3d353d1fd Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Mon, 8 Jan 2024 18:13:09 +0100
Subject: [PATCH 04/20] Better

---
 test/test.hs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/test.hs b/test/test.hs
index 62dfecd49..a7433977b 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -221,7 +221,8 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mapping_access(address,address)" [AbiAddressType, AbiAddressType])) [] defaultVeriOpts
-      putStrLnM $ T.unpack $ formatExpr expr
+      let a = mapExprM Expr.decomposeStorage expr
+      putStrLnM $ T.unpack $ formatExpr (fromJust a)
     , test "simplify-storage-map-only-static" $ do
        Just c <- solcRuntime "MyContract"
         [i|

From c01af5b6ffd17770fb012f36bb26a2bbc9894eaf Mon Sep 17 00:00:00 2001
From: dxo <hi@d-xo.org>
Date: Mon, 8 Jan 2024 19:36:18 +0100
Subject: [PATCH 05/20] Expr: decomposeStorage: fix ordering bug and tidy up

---
 src/EVM/Expr.hs | 22 +++++++++-------------
 src/EVM/SMT.hs  |  9 ++-------
 2 files changed, 11 insertions(+), 20 deletions(-)

diff --git a/src/EVM/Expr.hs b/src/EVM/Expr.hs
index b6bbb7ad8..b27f9f85d 100644
--- a/src/EVM/Expr.hs
+++ b/src/EVM/Expr.hs
@@ -7,8 +7,6 @@
 -}
 module EVM.Expr where
 
-import Debug.Trace
-
 import Prelude hiding (LT, GT)
 import Control.Monad.ST
 import Data.Bits hiding (And, Xor)
@@ -783,28 +781,26 @@ getLogicalIdx (GVar _) = internalError "cannot determine addr of a GVar"
 --     - MappingSlot
 --     - ArraySlotWithOffset
 --     - ArraySlotZero
---
 decomposeStorage :: Expr a -> Maybe (Expr a)
-decomposeStorage = mapExprM go
+decomposeStorage = go
   where
     go :: Expr a -> Maybe (Expr a)
     go (SLoad idx@(Keccak _) store) = case inferLogicalIdx idx of
       Just (logicalIdx, slot) -> do
         base <- setLogicalBase logicalIdx store
         pure (SLoad slot base)
-      Nothing -> trace "1" Nothing
+      Nothing -> Nothing
     go e@(SLoad (Lit idx) _) | idx < 256 = Just e
-    go (SLoad _ _) = trace "2" Nothing
+    go (SLoad _ _) = Nothing
 
-    go s@(SStore idx@(Keccak _) val store) = trace "INSIDE SSTORE" $ case inferLogicalIdx idx of
+    go (SStore idx@(Keccak _) val store) = case inferLogicalIdx idx of
       Just (logicalIdx, slot) -> do
-        traceM "DOING A REWRITE"
-        traceShowM s
+        -- TODO: eliminate writes to other logical indices
         base <- setLogicalBase logicalIdx store
         pure (SStore slot val base)
-      Nothing -> trace "3" Nothing
+      Nothing -> Nothing
     go e@(SStore (Lit idx) _ _) | idx < 256 = Just e
-    go s@(SStore _ _ _) = trace (show s) Nothing
+    go (SStore _ _ _) = Nothing
 
     go e = Just e
 
@@ -813,7 +809,7 @@ decomposeStorage = mapExprM go
       (MappingSlot idx key) -> Just (idxToWord idx, key)
       (ArraySlotWithOffset idx offset) -> Just (idxToWord idx, offset)
       (ArraySlotZero idx) -> Just (idxToWord idx, Lit 0)
-      _ -> trace "5" Nothing
+      _ -> Nothing
 
     idxToWord :: ByteString -> W256
     idxToWord = W256 . word256 . BS.takeEnd 32
@@ -835,7 +831,7 @@ decomposeStorage = mapExprM go
     setLogicalBase _ (ConcreteStore store) =
       if all (< 256) (Map.keys store)
       then Just (ConcreteStore mempty)
-      else trace "6" Nothing
+      else Nothing
     setLogicalBase _ (GVar _) = internalError "Unexpected GVar"
 
 
diff --git a/src/EVM/SMT.hs b/src/EVM/SMT.hs
index 73ac19704..76518549c 100644
--- a/src/EVM/SMT.hs
+++ b/src/EVM/SMT.hs
@@ -9,8 +9,6 @@ module EVM.SMT where
 
 import Prelude hiding (LT, GT)
 
-import Debug.Trace
-
 import Control.Monad
 import Data.Containers.ListUtils (nubOrd)
 import Data.ByteString (ByteString)
@@ -209,12 +207,9 @@ assertProps :: Config -> [Prop] -> SMT2
 -- simplify to rewrite sload/sstore combos
 -- notice: it is VERY important not to concretize, because Keccak assumptions
 --         will be generated by assertPropsNoSimp, and that needs unconcretized Props
-assertProps conf ps = assertPropsNoSimp conf withSepStore
+assertProps conf ps = assertPropsNoSimp conf (decompose . Expr.simplifyProps $ ps)
   where
-    simplifiedProps = trace "YO" $ Expr.simplifyProps ps
-    withSepStore = trace "OI" $ case mapM (mapPropM Expr.decomposeStorage) simplifiedProps of
-                     Just ps' -> trace "REWROTE" ps'
-                     Nothing -> trace "NO REWRITE" simplifiedProps
+    decompose props = fromMaybe props (mapM (mapPropM Expr.decomposeStorage) props)
 
 -- Note: we need a version that does NOT call simplify,
 -- because we make use of it to verify the correctness of our simplification

From ce1edac620e751b5db1d076bbd3e3775f52a28cd Mon Sep 17 00:00:00 2001
From: dxo <hi@d-xo.org>
Date: Mon, 8 Jan 2024 19:40:36 +0100
Subject: [PATCH 06/20] changelog: update

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ed10e0ef0..491fa7a4c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Keccak concretization is now done only after all simplifications are performed. This helps with simplification pre-concretization
 
 ## Added
+- Optimized smt queries that significantly improve performance when dealing with solidity mappings and arrays
 - New solc-specific simplification rules that should make the final Props a lot more readable
 - Prop is now correctly ordered, better BufLength and Max simplifications of Expr,
   and further solc-specific simplifications of Expr

From 65b1f7c7787d64f8db670e8d7b5b33bf53056fd8 Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Tue, 9 Jan 2024 18:18:16 +0100
Subject: [PATCH 07/20] Fixing some issues with the decomposer

---
 src/EVM/Expr.hs | 104 +++++++++++++++++++++++++++++++++++++-----------
 src/EVM/SMT.hs  |   1 +
 test/test.hs    |  61 +++++++++++++++++++++++++---
 3 files changed, 138 insertions(+), 28 deletions(-)

diff --git a/src/EVM/Expr.hs b/src/EVM/Expr.hs
index b27f9f85d..3b2b0832c 100644
--- a/src/EVM/Expr.hs
+++ b/src/EVM/Expr.hs
@@ -781,47 +781,106 @@ getLogicalIdx (GVar _) = internalError "cannot determine addr of a GVar"
 --     - MappingSlot
 --     - ArraySlotWithOffset
 --     - ArraySlotZero
+--  - mixing within the same store & load is not permitted
+data StorageType = SmallSlot | Array | Map | Mixed | UNK
+  deriving (Show, Eq)
+
+newtype CollectStorageType = CollectStorageType {storeType :: StorageType }
+  deriving (Show)
+
+initCollectStorageType :: CollectStorageType
+initCollectStorageType = CollectStorageType {storeType = UNK }
+
+decomposeCheckTypes :: Expr a -> Bool
+decomposeCheckTypes inp = result.storeType /= Mixed
+  where
+    result = execState (decomposeCheckTypesRunner inp) initCollectStorageType
+
+    decomposeCheckTypesRunner :: forall a. Expr a -> State CollectStorageType (Expr a)
+    decomposeCheckTypesRunner a = go a
+
+    go :: forall b. Expr b -> State CollectStorageType (Expr b)
+    go e@(SLoad (MappingSlot {}) _) = setMap e
+    go e@(SLoad (ArraySlotWithOffset {}) _) = setArray e
+    go e@(SLoad (ArraySlotZero {}) _) = setArray e
+    go e@(SLoad (Lit x) _) | x < 256 = setArray e
+    go e@(SLoad _ _) = setMixed e
+    go e@(SStore (MappingSlot {}) _ _) = setMap e
+    go e@(SStore (ArraySlotZero {}) _ _) = setArray e
+    go e@(SStore (ArraySlotWithOffset {}) _ _) = setArray e
+    go e@(SStore (Lit x) _ _) | x < 256 = setSmall e
+    go e@(SStore _ _ _) = setMixed e
+    go e = pure e
+
+    -- Helper functions for detecting mixed load/store
+    setMixed e = do
+      put CollectStorageType {storeType = Mixed}
+      pure e
+    setMap e = do
+      s <- get
+      case s.storeType of
+        Array -> put s {storeType = Mixed}
+        SmallSlot -> put s {storeType = Mixed}
+        UNK -> put s {storeType = Map}
+        _ -> pure ()
+      pure e
+    setArray e = do
+      s <- get
+      case s.storeType of
+        Map -> put s {storeType = Mixed}
+        SmallSlot -> put s {storeType = Mixed}
+        UNK -> put s {storeType = Array}
+        _ -> pure ()
+      pure e
+    setSmall e = do
+      s <- get
+      case s.storeType of
+        Map -> put s {storeType = Mixed}
+        Array -> put s {storeType = Mixed}
+        UNK -> put s {storeType = SmallSlot}
+        _ -> pure ()
+      pure e
+
 decomposeStorage :: Expr a -> Maybe (Expr a)
 decomposeStorage = go
   where
+    -- NOTE: it's a bad idea to rewrite SStore on its own, we should rewrite SLoad, which calls SStore
+    --       this ensures that the rewrite happens to the whole SLoad... chain. If we rewrite SStore,
+    --       we will impact the returned "final state" of the system that will then be incorrect since
+    --       the final state actually contains all the Keccak-s
     go :: Expr a -> Maybe (Expr a)
-    go (SLoad idx@(Keccak _) store) = case inferLogicalIdx idx of
-      Just (logicalIdx, slot) -> do
-        base <- setLogicalBase logicalIdx store
-        pure (SLoad slot base)
-      Nothing -> Nothing
-    go e@(SLoad (Lit idx) _) | idx < 256 = Just e
-    go (SLoad _ _) = Nothing
-
-    go (SStore idx@(Keccak _) val store) = case inferLogicalIdx idx of
-      Just (logicalIdx, slot) -> do
-        -- TODO: eliminate writes to other logical indices
-        base <- setLogicalBase logicalIdx store
-        pure (SStore slot val base)
-      Nothing -> Nothing
-    go e@(SStore (Lit idx) _ _) | idx < 256 = Just e
-    go (SStore _ _ _) = Nothing
+    go e@(SLoad origKey store) = if Prelude.not (decomposeCheckTypes e) then Nothing else case inferLogicalIdx origKey of
+      Just (idx, key) -> do
+        base <- setLogicalBase idx store
+        pure (SLoad key base)
+      _ -> Nothing
 
     go e = Just e
 
     inferLogicalIdx :: Expr EWord -> Maybe (W256, Expr EWord)
     inferLogicalIdx = \case
+      Lit a | a >= 256 -> Nothing
+      Lit a -> Just (0, Lit a)
       (MappingSlot idx key) -> Just (idxToWord idx, key)
-      (ArraySlotWithOffset idx offset) -> Just (idxToWord idx, offset)
-      (ArraySlotZero idx) -> Just (idxToWord idx, Lit 0)
+      (ArraySlotWithOffset idx offset) | BS.length idx == 32 -> Just (idxToWord64 idx, offset)
+      (ArraySlotZero idx) | BS.length idx == 32 -> Just (idxToWord64 idx, Lit 0)
       _ -> Nothing
 
     idxToWord :: ByteString -> W256
     idxToWord = W256 . word256 . BS.takeEnd 32
+    -- Arrays take the whole `id` and keccak it. It's supposed to be 64B
+    idxToWord64 :: ByteString -> W256
+    idxToWord64 = W256 . word256 . BS.takeEnd 64
 
     -- Updates the logical base store of the given expression if it is safe to do so
     setLogicalBase :: W256 -> Expr Storage -> Maybe (Expr Storage)
 
     -- abstract bases get their logical idx set to the new value
-    setLogicalBase new (AbstractStore addr _) = Just $ AbstractStore addr (Just new)
-    setLogicalBase new (SStore i v s) = do
-      b <- setLogicalBase new s
-      Just $ SStore i v b
+    setLogicalBase idx (AbstractStore addr _) = Just $ AbstractStore addr (Just idx)
+    setLogicalBase idx (SStore i v s) = do
+      b <- setLogicalBase idx s
+      (idx2, key2) <- inferLogicalIdx i
+      if idx == idx2 then Just (SStore key2 v b) else Nothing
 
     -- empty concrete base is safe to reuse without any rewriting
     setLogicalBase _ s@(ConcreteStore m) | Map.null m = Just s
@@ -835,7 +894,6 @@ decomposeStorage = go
     setLogicalBase _ (GVar _) = internalError "Unexpected GVar"
 
 
-
 -- | Simple recursive match based AST simplification
 -- Note: may not terminate!
 simplify :: Expr a -> Expr a
diff --git a/src/EVM/SMT.hs b/src/EVM/SMT.hs
index 76518549c..e6a602c81 100644
--- a/src/EVM/SMT.hs
+++ b/src/EVM/SMT.hs
@@ -209,6 +209,7 @@ assertProps :: Config -> [Prop] -> SMT2
 --         will be generated by assertPropsNoSimp, and that needs unconcretized Props
 assertProps conf ps = assertPropsNoSimp conf (decompose . Expr.simplifyProps $ ps)
   where
+    decompose :: [Prop] -> [Prop]
     decompose props = fromMaybe props (mapM (mapPropM Expr.decomposeStorage) props)
 
 -- Note: we need a version that does NOT call simplify,
diff --git a/test/test.hs b/test/test.hs
index a7433977b..e5e508dba 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -205,13 +205,11 @@ tests = testGroup "hevm"
         |]
        expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "transfer(uint256,uint256)" [AbiUIntType 256, AbiUIntType 256])) [] (defaultVeriOpts { maxIter = Just 5 })
        assertEqualM "Expression is not clean." (badStoresInExpr expr) False
-    , test "different-addr" $ do
+    , test "decompose-1" $ do
       Just c <- solcRuntime "MyContract"
         [i|
         contract MyContract {
           mapping (address => uint) balances;
-          mapping (uint => bool) auth;
-          mapping (address => mapping (address => uint)) allowance;
           function prove_mapping_access(address x, address y) public {
               require(x != y);
               balances[x] = 1;
@@ -221,8 +219,61 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mapping_access(address,address)" [AbiAddressType, AbiAddressType])) [] defaultVeriOpts
-      let a = mapExprM Expr.decomposeStorage expr
-      putStrLnM $ T.unpack $ formatExpr (fromJust a)
+      let simpExpr = mapExprM Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
+      assertEqualM "Expression is not clean." (isJust simpExpr) True
+    , test "decompose-2" $ do
+      Just c <- solcRuntime "MyContract"
+        [i|
+        contract MyContract {
+          mapping (address => uint) balances;
+          function prove_mixed_symoblic_concrete_writes(address x, uint v) public {
+              balances[x] = v;
+              balances[address(0)] = balances[x];
+              assert(balances[address(0)] == v);
+          }
+        }
+        |]
+      expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mixed_symoblic_concrete_writes(address,uint256)" [AbiAddressType, AbiUIntType 256])) [] defaultVeriOpts
+      let simpExpr = mapExprM Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
+      assertEqualM "Expression is not clean." (isJust simpExpr) True
+    , test "decompose-3" $ do
+      Just c <- solcRuntime "MyContract"
+        [i|
+        contract MyContract {
+          uint[] a;
+          function prove_array(uint x, uint v1, uint y, uint v2) public {
+              require(v1 != v2);
+              a[x] = v1;
+              a[y] = v2;
+              assert(a[x] == a[y]);
+          }
+        }
+        |]
+      expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_array(uint256,uint256,uint256,uint256)" [AbiUIntType 256, AbiUIntType 256, AbiUIntType 256, AbiUIntType 256])) [] defaultVeriOpts
+      let simpExpr = mapExprM Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
+      assertEqualM "Expression is not clean." (isJust simpExpr) True
+    , test "decompose-4-mixed" $ do
+      Just c <- solcRuntime "MyContract"
+        [i|
+        contract MyContract {
+          uint[] a;
+          mapping( uint => uint) balances;
+          function prove_array(uint x, uint v1, uint y, uint v2) public {
+              require(v1 != v2);
+              balances[x] = v1+1;
+              balances[y] = v1+2;
+              a[x] = v1;
+              assert(balances[x] != balances[y]);
+          }
+        }
+        |]
+      expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_array(uint256,uint256,uint256,uint256)" [AbiUIntType 256, AbiUIntType 256, AbiUIntType 256, AbiUIntType 256])) [] defaultVeriOpts
+      let simpExpr = mapExprM Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
+      assertEqualM "Expression is not clean." (isJust simpExpr) True
     , test "simplify-storage-map-only-static" $ do
        Just c <- solcRuntime "MyContract"
         [i|

From 38d90d0e3119162f9c84bd7b2f41cbe0f8df5da0 Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Tue, 9 Jan 2024 18:50:04 +0100
Subject: [PATCH 08/20] Adding one more test

---
 test/test.hs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/test/test.hs b/test/test.hs
index e5e508dba..bb1a0f9ce 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -274,6 +274,27 @@ tests = testGroup "hevm"
       let simpExpr = mapExprM Expr.decomposeStorage expr
       -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
       assertEqualM "Expression is not clean." (isJust simpExpr) True
+    , test "decompose-5-mixed" $ do
+      Just c <- solcRuntime "MyContract"
+        [i|
+        contract MyContract {
+          mapping (address => uint) balances;
+          mapping (uint => bool) auth;
+          function prove_mixed(address x, address y, uint val) public {
+            require(x != y);
+            balances[x] = val;
+            balances[y] = val+2;
+            auth[val+1] = true;
+            if (balances[y] == balances[y]) {
+                assert(balances[y] == val);
+            }
+          }
+        }
+        |]
+      expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mixed(address,address,uint256)" [AbiAddressType, AbiAddressType, AbiUIntType 256])) [] defaultVeriOpts
+      let simpExpr = mapExprM Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
+      assertEqualM "Expression is not clean." (isJust simpExpr) True
     , test "simplify-storage-map-only-static" $ do
        Just c <- solcRuntime "MyContract"
         [i|

From 645555a5e98cdd017f8d0deb31630858a2e39e31 Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Tue, 9 Jan 2024 18:50:49 +0100
Subject: [PATCH 09/20] Adding more more complicated test

---
 test/test.hs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/test/test.hs b/test/test.hs
index bb1a0f9ce..1068fcdef 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -280,10 +280,14 @@ tests = testGroup "hevm"
         contract MyContract {
           mapping (address => uint) balances;
           mapping (uint => bool) auth;
+          uint a;
+          uint b;
           function prove_mixed(address x, address y, uint val) public {
+            b = val+1;
             require(x != y);
             balances[x] = val;
             balances[y] = val+2;
+            a = val;
             auth[val+1] = true;
             if (balances[y] == balances[y]) {
                 assert(balances[y] == val);

From fb1b92d7939faf6b51d578daaee8ccc8a7d1d26e Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Tue, 9 Jan 2024 18:53:18 +0100
Subject: [PATCH 10/20] Adding one more test

---
 test/test.hs | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/test/test.hs b/test/test.hs
index 1068fcdef..acb9e1234 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -280,15 +280,17 @@ tests = testGroup "hevm"
         contract MyContract {
           mapping (address => uint) balances;
           mapping (uint => bool) auth;
+          uint[] arr;
           uint a;
           uint b;
           function prove_mixed(address x, address y, uint val) public {
             b = val+1;
             require(x != y);
             balances[x] = val;
-            balances[y] = val+2;
             a = val;
+            arr[val] = 5;
             auth[val+1] = true;
+            balances[y] = val+2;
             if (balances[y] == balances[y]) {
                 assert(balances[y] == val);
             }

From 4c63d564efc2f1eb14960930a5ef73aeab3e5f36 Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Tue, 9 Jan 2024 18:55:23 +0100
Subject: [PATCH 11/20] Fixing small bug

---
 src/EVM/Expr.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/EVM/Expr.hs b/src/EVM/Expr.hs
index 3b2b0832c..1b6026abe 100644
--- a/src/EVM/Expr.hs
+++ b/src/EVM/Expr.hs
@@ -803,7 +803,7 @@ decomposeCheckTypes inp = result.storeType /= Mixed
     go e@(SLoad (MappingSlot {}) _) = setMap e
     go e@(SLoad (ArraySlotWithOffset {}) _) = setArray e
     go e@(SLoad (ArraySlotZero {}) _) = setArray e
-    go e@(SLoad (Lit x) _) | x < 256 = setArray e
+    go e@(SLoad (Lit x) _) | x < 256 = setSmall e
     go e@(SLoad _ _) = setMixed e
     go e@(SStore (MappingSlot {}) _ _) = setMap e
     go e@(SStore (ArraySlotZero {}) _ _) = setArray e

From 6d44dc7e001caeb2817ef5b448911e2f14efd09f Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Tue, 9 Jan 2024 18:55:44 +0100
Subject: [PATCH 12/20] Cleaner order

---
 src/EVM/Expr.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/EVM/Expr.hs b/src/EVM/Expr.hs
index 1b6026abe..fd7efd1f3 100644
--- a/src/EVM/Expr.hs
+++ b/src/EVM/Expr.hs
@@ -801,8 +801,8 @@ decomposeCheckTypes inp = result.storeType /= Mixed
 
     go :: forall b. Expr b -> State CollectStorageType (Expr b)
     go e@(SLoad (MappingSlot {}) _) = setMap e
-    go e@(SLoad (ArraySlotWithOffset {}) _) = setArray e
     go e@(SLoad (ArraySlotZero {}) _) = setArray e
+    go e@(SLoad (ArraySlotWithOffset {}) _) = setArray e
     go e@(SLoad (Lit x) _) | x < 256 = setSmall e
     go e@(SLoad _ _) = setMixed e
     go e@(SStore (MappingSlot {}) _ _) = setMap e

From 4ae9ed704913280f97342af96b60f821c316d869 Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Wed, 10 Jan 2024 13:11:39 +0100
Subject: [PATCH 13/20] Rewrite what we can, rather than bailing out on
 everything

---
 src/EVM/Expr.hs | 12 ++++++----
 src/EVM/SMT.hs  |  2 +-
 test/test.hs    | 64 +++++++++++++++++++++++++++++++++++++------------
 3 files changed, 57 insertions(+), 21 deletions(-)

diff --git a/src/EVM/Expr.hs b/src/EVM/Expr.hs
index fd7efd1f3..e88fb386b 100644
--- a/src/EVM/Expr.hs
+++ b/src/EVM/Expr.hs
@@ -841,22 +841,24 @@ decomposeCheckTypes inp = result.storeType /= Mixed
         _ -> pure ()
       pure e
 
-decomposeStorage :: Expr a -> Maybe (Expr a)
+decomposeStorage :: Expr a -> Expr a
 decomposeStorage = go
   where
     -- NOTE: it's a bad idea to rewrite SStore on its own, we should rewrite SLoad, which calls SStore
     --       this ensures that the rewrite happens to the whole SLoad... chain. If we rewrite SStore,
     --       we will impact the returned "final state" of the system that will then be incorrect since
     --       the final state actually contains all the Keccak-s
-    go :: Expr a -> Maybe (Expr a)
-    go e@(SLoad origKey store) = if Prelude.not (decomposeCheckTypes e) then Nothing else case inferLogicalIdx origKey of
+    go :: Expr a -> Expr a
+    go e@(SLoad origKey store) = if Prelude.not (decomposeCheckTypes e) then e else fromMaybe e (tryRewrite origKey store)
+    go e = e
+
+    tryRewrite :: Expr EWord -> Expr Storage -> Maybe (Expr EWord)
+    tryRewrite origKey store = case inferLogicalIdx origKey of
       Just (idx, key) -> do
         base <- setLogicalBase idx store
         pure (SLoad key base)
       _ -> Nothing
 
-    go e = Just e
-
     inferLogicalIdx :: Expr EWord -> Maybe (W256, Expr EWord)
     inferLogicalIdx = \case
       Lit a | a >= 256 -> Nothing
diff --git a/src/EVM/SMT.hs b/src/EVM/SMT.hs
index e6a602c81..8ad06482a 100644
--- a/src/EVM/SMT.hs
+++ b/src/EVM/SMT.hs
@@ -210,7 +210,7 @@ assertProps :: Config -> [Prop] -> SMT2
 assertProps conf ps = assertPropsNoSimp conf (decompose . Expr.simplifyProps $ ps)
   where
     decompose :: [Prop] -> [Prop]
-    decompose props = fromMaybe props (mapM (mapPropM Expr.decomposeStorage) props)
+    decompose props = map (mapProp Expr.decomposeStorage) props
 
 -- Note: we need a version that does NOT call simplify,
 -- because we make use of it to verify the correctness of our simplification
diff --git a/test/test.hs b/test/test.hs
index acb9e1234..a8544c973 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -219,9 +219,9 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mapping_access(address,address)" [AbiAddressType, AbiAddressType])) [] defaultVeriOpts
-      let simpExpr = mapExprM Expr.decomposeStorage expr
-      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
-      assertEqualM "Expression is not clean." (isJust simpExpr) True
+      let simpExpr = mapExpr Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr simpExpr
+      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
     , test "decompose-2" $ do
       Just c <- solcRuntime "MyContract"
         [i|
@@ -235,9 +235,9 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mixed_symoblic_concrete_writes(address,uint256)" [AbiAddressType, AbiUIntType 256])) [] defaultVeriOpts
-      let simpExpr = mapExprM Expr.decomposeStorage expr
-      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
-      assertEqualM "Expression is not clean." (isJust simpExpr) True
+      let simpExpr = mapExpr Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr simpExpr
+      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
     , test "decompose-3" $ do
       Just c <- solcRuntime "MyContract"
         [i|
@@ -252,9 +252,9 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_array(uint256,uint256,uint256,uint256)" [AbiUIntType 256, AbiUIntType 256, AbiUIntType 256, AbiUIntType 256])) [] defaultVeriOpts
-      let simpExpr = mapExprM Expr.decomposeStorage expr
-      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
-      assertEqualM "Expression is not clean." (isJust simpExpr) True
+      let simpExpr = mapExpr Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr simpExpr
+      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
     , test "decompose-4-mixed" $ do
       Just c <- solcRuntime "MyContract"
         [i|
@@ -271,9 +271,9 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_array(uint256,uint256,uint256,uint256)" [AbiUIntType 256, AbiUIntType 256, AbiUIntType 256, AbiUIntType 256])) [] defaultVeriOpts
-      let simpExpr = mapExprM Expr.decomposeStorage expr
-      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
-      assertEqualM "Expression is not clean." (isJust simpExpr) True
+      let simpExpr = mapExpr Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr simpExpr
+      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
     , test "decompose-5-mixed" $ do
       Just c <- solcRuntime "MyContract"
         [i|
@@ -298,9 +298,26 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mixed(address,address,uint256)" [AbiAddressType, AbiAddressType, AbiUIntType 256])) [] defaultVeriOpts
-      let simpExpr = mapExprM Expr.decomposeStorage expr
-      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
-      assertEqualM "Expression is not clean." (isJust simpExpr) True
+      let simpExpr = mapExpr Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr simpExpr
+      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
+    -- TODO check what's going on here. Likely the "arbitrary write through array" is the reason why we fail
+    , expectFail $ test "decompose-6-fail" $ do
+      Just c <- solcRuntime "MyContract"
+        [i|
+        contract MyContract {
+          uint[] arr;
+          function prove_mixed(uint val) public {
+            arr[val] = 5;
+            arr[val+1] = val+5;
+            assert(arr[val] == arr[val+1]);
+          }
+        }
+        |]
+      expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mixed(uint256)" [AbiUIntType 256])) [] defaultVeriOpts
+      let simpExpr = mapExpr Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr simpExpr
+      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
     , test "simplify-storage-map-only-static" $ do
        Just c <- solcRuntime "MyContract"
         [i|
@@ -4088,6 +4105,23 @@ badStoresInExpr expr = bad
         pure e
       _ -> pure e
 
+keccakStoreInExpr :: Expr a -> Bool
+keccakStoreInExpr expr = bad
+  where
+    FoundBad bad = execState (mapExprM findBadLoad expr) initFoundBad
+    findBadLoad :: Expr a-> State FoundBad (Expr a)
+    findBadLoad e = case e of
+      SLoad key store -> do
+        _ <- findBad key e
+        findBad store e
+      _ -> pure e
+    findBad e orig = if isNothing $ mapExprM keccakFinder e then do
+          put (FoundBad { bad = True })
+          pure orig
+        else pure orig
+    keccakFinder (Keccak {}) = Nothing
+    keccakFinder e = Just e
+
 defaultBuf :: Int -> Gen (Expr Buf)
 defaultBuf = genBuf (4_000_000)
 

From 05edded828f1a05373309c8e2f9c24af4a67bf92 Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Wed, 10 Jan 2024 16:32:00 +0100
Subject: [PATCH 14/20] Updated fuzzer for storage load

---
 test/test.hs | 67 +++++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 61 insertions(+), 6 deletions(-)

diff --git a/test/test.hs b/test/test.hs
index a8544c973..93e585b38 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -500,6 +500,16 @@ tests = testGroup "hevm"
     , testProperty "store-simplification" $ \(expr :: Expr Storage) -> prop $ do
         let simplified = Expr.simplify expr
         checkEquiv expr simplified
+    , testProperty "load-simplification" $ \(GenWriteStorageLoad expr) -> prop $ do
+        let simplified = Expr.simplify expr
+        checkEquiv expr simplified
+    , testProperty "load-decompose" $ \(GenWriteStorageLoad expr) -> prop $ do
+        putStrLnM $ T.unpack $ formatExpr expr
+        let decomposed = Expr.decomposeStorage (Expr.simplify expr)
+        putStrLnM $ "-----------------------------------------"
+        putStrLnM $ T.unpack $ formatExpr decomposed
+        putStrLnM $ "\n\n\n\n"
+        checkEquiv expr decomposed
     , testProperty "byte-simplification" $ \(expr :: Expr Byte) -> prop $ do
         let simplified = Expr.simplify expr
         checkEquiv expr simplified
@@ -3927,6 +3937,11 @@ genEnd sz = oneof
     subEnd = genEnd (sz `div` 2)
     subProp = genProps False (sz `div` 2)
 
+genSmallLit :: W256 -> Gen (Expr EWord)
+genSmallLit m = do
+  val :: W256 <- arbitrary
+  pure $ Lit (val `mod` m)
+
 genWord :: Int -> Int -> Gen (Expr EWord)
 genWord litFreq 0 = frequency
   [ (litFreq, do
@@ -4160,14 +4175,54 @@ genBuf bound sz = oneof
     subBuf = genBuf bound (sz `div` 10)
 
 genStorage :: Int -> Gen (Expr Storage)
-genStorage 0 = oneof
-  [ liftM2 AbstractStore arbitrary arbitrary
-  , fmap ConcreteStore arbitrary
+genStorage sz = genStorageMap True sz
+
+genStorageMap :: Bool -> Int -> Gen (Expr Storage)
+genStorageMap withMap 0 = oneof
+  [ liftM2 AbstractStore arbitrary (if withMap then arbitrary else pure Nothing)
+  , fmap ConcreteStore (if withMap then arbitrary else pure mempty)
   ]
-genStorage sz = liftM3 SStore subWord subWord subStore
+genStorageMap withMap sz = liftM3 SStore key val subStore
   where
-    subStore = genStorage (sz `div` 10)
-    subWord = defaultWord (sz `div` 5)
+    subStore = genStorageMap withMap (sz `div` 10)
+    val = defaultWord (sz `div` 5)
+    key = genStorageKey
+
+genStorageKey :: Gen (Expr EWord)
+genStorageKey = frequency
+     -- array slot
+    [ (4, liftM2 Expr.ArraySlotWithOffset (genByteStringKey 32) (genSmallLit 5))
+    , (4, fmap Expr.ArraySlotZero (genByteStringKey 32))
+     -- mapping slot
+    , (8, liftM2 Expr.MappingSlot (genByteStringKey 64) (genSmallLit 5))
+     -- small slot
+    , (4, genLit 20)
+    -- unrecognized slot type
+    , (1, genSmallLit 5)
+    ]
+
+genByteStringKey :: W256 -> Gen (ByteString)
+genByteStringKey len = oneof
+    [ do
+        b :: Word8 <- arbitrary
+        pure $ BS.pack ([ 0 | _ <- [0..(len-2)]] ++ [b `mod` 5])
+    ]
+
+-- GenWriteStorageLoad
+newtype GenWriteStorageLoad = GenWriteStorageLoad (Expr EWord)
+  deriving (Show, Eq)
+
+instance Arbitrary GenWriteStorageLoad where
+  arbitrary = do
+    load <- genStorageLoad 10
+    pure $ GenWriteStorageLoad load
+
+    where
+      genStorageLoad :: Int -> Gen (Expr EWord)
+      genStorageLoad sz = liftM2 SLoad key subStore
+        where
+          subStore = genStorageMap False (sz `div` 10)
+          key = genStorageKey
 
 data Invocation
   = SolidityCall Text [AbiValue]

From de80888c96971c24345b8f1d1b3ea99e7e77b76c Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Wed, 10 Jan 2024 18:29:32 +0100
Subject: [PATCH 15/20] Fixing up the array separation

---
 src/EVM/Expr.hs | 34 +++++++++++++---------
 src/EVM/SMT.hs  |  2 +-
 test/test.hs    | 77 ++++++++++++++++++++-----------------------------
 3 files changed, 52 insertions(+), 61 deletions(-)

diff --git a/src/EVM/Expr.hs b/src/EVM/Expr.hs
index e88fb386b..ee2e28383 100644
--- a/src/EVM/Expr.hs
+++ b/src/EVM/Expr.hs
@@ -644,7 +644,11 @@ readStorage w st = go (simplify w) st
       -- for our purposes. This lets us completely simplify reads from write
       -- chains involving writes to arrays at literal offsets.
       (Lit a, Add (Lit b) (Keccak _) ) | a < 256, b < maxW32 -> go slot prev
-      (Add (Lit a) (Keccak _) , Lit b) | a < 256, b < maxW32 -> go slot prev
+      (Add (Lit a) (Keccak _) , Lit b) | b < 256, a < maxW32 -> go slot prev
+
+      --- NOTE these are needed to succeed in rewriting arrays with a variable index
+      -- (Lit a, Add (Keccak _) (Var _) ) | a < 256 -> go slot prev
+      -- (Add (Keccak _) (Var _) , Lit b) | b < 256 -> go slot prev
 
       -- Finding two Keccaks that are < 256 away from each other should be VERY hard
       -- This simplification allows us to deal with maps of structs
@@ -841,16 +845,16 @@ decomposeCheckTypes inp = result.storeType /= Mixed
         _ -> pure ()
       pure e
 
-decomposeStorage :: Expr a -> Expr a
+decomposeStorage :: Expr a -> Maybe (Expr a)
 decomposeStorage = go
   where
     -- NOTE: it's a bad idea to rewrite SStore on its own, we should rewrite SLoad, which calls SStore
     --       this ensures that the rewrite happens to the whole SLoad... chain. If we rewrite SStore,
     --       we will impact the returned "final state" of the system that will then be incorrect since
     --       the final state actually contains all the Keccak-s
-    go :: Expr a -> Expr a
-    go e@(SLoad origKey store) = if Prelude.not (decomposeCheckTypes e) then e else fromMaybe e (tryRewrite origKey store)
-    go e = e
+    go :: Expr a -> Maybe (Expr a)
+    go e@(SLoad origKey store) = if Prelude.not (decomposeCheckTypes e) then Nothing else tryRewrite origKey store
+    go e = Just e
 
     tryRewrite :: Expr EWord -> Expr Storage -> Maybe (Expr EWord)
     tryRewrite origKey store = case inferLogicalIdx origKey of
@@ -859,26 +863,28 @@ decomposeStorage = go
         pure (SLoad key base)
       _ -> Nothing
 
-    inferLogicalIdx :: Expr EWord -> Maybe (W256, Expr EWord)
+    -- NOTE: we use (Maybe W256) for idx here, because for small slot numbers we want to keep the
+    -- Logical Store value a Nothing
+    inferLogicalIdx :: Expr EWord -> Maybe (Maybe W256, Expr EWord)
     inferLogicalIdx = \case
       Lit a | a >= 256 -> Nothing
-      Lit a -> Just (0, Lit a)
-      (MappingSlot idx key) -> Just (idxToWord idx, key)
-      (ArraySlotWithOffset idx offset) | BS.length idx == 32 -> Just (idxToWord64 idx, offset)
-      (ArraySlotZero idx) | BS.length idx == 32 -> Just (idxToWord64 idx, Lit 0)
+      Lit a -> Just (Nothing, Lit a)
+      (MappingSlot idx key) | BS.length idx == 64 -> Just (Just $ idxToWord idx, key)
+      (ArraySlotWithOffset idx offset) | BS.length idx == 32 -> Just (Just $ idxToWord64 idx, offset)
+      (ArraySlotZero idx) | BS.length idx == 32 -> Just (Just $ idxToWord64 idx, Lit 0)
       _ -> Nothing
 
     idxToWord :: ByteString -> W256
-    idxToWord = W256 . word256 . BS.takeEnd 32
+    idxToWord = W256 . word256 . (BS.takeEnd 32)
     -- Arrays take the whole `id` and keccak it. It's supposed to be 64B
     idxToWord64 :: ByteString -> W256
-    idxToWord64 = W256 . word256 . BS.takeEnd 64
+    idxToWord64 = W256 . word256 . (BS.takeEnd 64)
 
     -- Updates the logical base store of the given expression if it is safe to do so
-    setLogicalBase :: W256 -> Expr Storage -> Maybe (Expr Storage)
+    setLogicalBase :: Maybe W256 -> Expr Storage -> Maybe (Expr Storage)
 
     -- abstract bases get their logical idx set to the new value
-    setLogicalBase idx (AbstractStore addr _) = Just $ AbstractStore addr (Just idx)
+    setLogicalBase idx (AbstractStore addr _) = Just $ AbstractStore addr idx
     setLogicalBase idx (SStore i v s) = do
       b <- setLogicalBase idx s
       (idx2, key2) <- inferLogicalIdx i
diff --git a/src/EVM/SMT.hs b/src/EVM/SMT.hs
index 8ad06482a..e6a602c81 100644
--- a/src/EVM/SMT.hs
+++ b/src/EVM/SMT.hs
@@ -210,7 +210,7 @@ assertProps :: Config -> [Prop] -> SMT2
 assertProps conf ps = assertPropsNoSimp conf (decompose . Expr.simplifyProps $ ps)
   where
     decompose :: [Prop] -> [Prop]
-    decompose props = map (mapProp Expr.decomposeStorage) props
+    decompose props = fromMaybe props (mapM (mapPropM Expr.decomposeStorage) props)
 
 -- Note: we need a version that does NOT call simplify,
 -- because we make use of it to verify the correctness of our simplification
diff --git a/test/test.hs b/test/test.hs
index 93e585b38..96753ae67 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -219,9 +219,9 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mapping_access(address,address)" [AbiAddressType, AbiAddressType])) [] defaultVeriOpts
-      let simpExpr = mapExpr Expr.decomposeStorage expr
-      -- putStrLnM $ T.unpack $ formatExpr simpExpr
-      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
+      let simpExpr = mapExprM Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
+      assertEqualM "Decompose did not succeed." (isJust simpExpr) True
     , test "decompose-2" $ do
       Just c <- solcRuntime "MyContract"
         [i|
@@ -235,10 +235,13 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mixed_symoblic_concrete_writes(address,uint256)" [AbiAddressType, AbiUIntType 256])) [] defaultVeriOpts
-      let simpExpr = mapExpr Expr.decomposeStorage expr
-      -- putStrLnM $ T.unpack $ formatExpr simpExpr
-      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
-    , test "decompose-3" $ do
+      let simpExpr = mapExprM Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
+      assertEqualM "Decompose did not succeed." (isJust simpExpr) True
+    -- NOTE: we can't do the rewrite below, because x could be very large,
+    -- which would allow us to overwrite *anything in the storage*, and that throws
+    -- everything off.
+    , expectFail $ test "decompose-3" $ do
       Just c <- solcRuntime "MyContract"
         [i|
         contract MyContract {
@@ -252,9 +255,9 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_array(uint256,uint256,uint256,uint256)" [AbiUIntType 256, AbiUIntType 256, AbiUIntType 256, AbiUIntType 256])) [] defaultVeriOpts
-      let simpExpr = mapExpr Expr.decomposeStorage expr
-      -- putStrLnM $ T.unpack $ formatExpr simpExpr
-      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
+      let simpExpr = mapExprM Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
+      assertEqualM "Decompose did not succeed." (isJust simpExpr) True
     , test "decompose-4-mixed" $ do
       Just c <- solcRuntime "MyContract"
         [i|
@@ -271,9 +274,9 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_array(uint256,uint256,uint256,uint256)" [AbiUIntType 256, AbiUIntType 256, AbiUIntType 256, AbiUIntType 256])) [] defaultVeriOpts
-      let simpExpr = mapExpr Expr.decomposeStorage expr
-      -- putStrLnM $ T.unpack $ formatExpr simpExpr
-      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
+      let simpExpr = mapExprM Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
+      assertEqualM "Decompose did not succeed." (isJust simpExpr) True
     , test "decompose-5-mixed" $ do
       Just c <- solcRuntime "MyContract"
         [i|
@@ -298,9 +301,9 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mixed(address,address,uint256)" [AbiAddressType, AbiAddressType, AbiUIntType 256])) [] defaultVeriOpts
-      let simpExpr = mapExpr Expr.decomposeStorage expr
-      -- putStrLnM $ T.unpack $ formatExpr simpExpr
-      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
+      let simpExpr = mapExprM Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
+      assertEqualM "Decompose did not succeed." (isJust simpExpr) True
     -- TODO check what's going on here. Likely the "arbitrary write through array" is the reason why we fail
     , expectFail $ test "decompose-6-fail" $ do
       Just c <- solcRuntime "MyContract"
@@ -315,9 +318,9 @@ tests = testGroup "hevm"
         }
         |]
       expr <- withSolvers Z3 1 Nothing $ \s -> getExpr s c (Just (Sig "prove_mixed(uint256)" [AbiUIntType 256])) [] defaultVeriOpts
-      let simpExpr = mapExpr Expr.decomposeStorage expr
-      -- putStrLnM $ T.unpack $ formatExpr simpExpr
-      assertEqualM "Expression is not clean." (keccakStoreInExpr simpExpr) False
+      let simpExpr = mapExprM Expr.decomposeStorage expr
+      -- putStrLnM $ T.unpack $ formatExpr (fromJust simpExpr)
+      assertEqualM "Decompose did not succeed." (isJust simpExpr) True
     , test "simplify-storage-map-only-static" $ do
        Just c <- solcRuntime "MyContract"
         [i|
@@ -503,12 +506,13 @@ tests = testGroup "hevm"
     , testProperty "load-simplification" $ \(GenWriteStorageLoad expr) -> prop $ do
         let simplified = Expr.simplify expr
         checkEquiv expr simplified
-    , testProperty "load-decompose" $ \(GenWriteStorageLoad expr) -> prop $ do
+    , ignoreTest $ testProperty "load-decompose" $ \(GenWriteStorageLoad expr) -> prop $ do
         putStrLnM $ T.unpack $ formatExpr expr
-        let decomposed = Expr.decomposeStorage (Expr.simplify expr)
-        putStrLnM $ "-----------------------------------------"
-        putStrLnM $ T.unpack $ formatExpr decomposed
-        putStrLnM $ "\n\n\n\n"
+        let simp = Expr.simplify expr
+        let decomposed = fromMaybe simp $ mapExprM Expr.decomposeStorage simp
+        -- putStrLnM $ "-----------------------------------------"
+        -- putStrLnM $ T.unpack $ formatExpr decomposed
+        -- putStrLnM $ "\n\n\n\n"
         checkEquiv expr decomposed
     , testProperty "byte-simplification" $ \(expr :: Expr Byte) -> prop $ do
         let simplified = Expr.simplify expr
@@ -4120,23 +4124,6 @@ badStoresInExpr expr = bad
         pure e
       _ -> pure e
 
-keccakStoreInExpr :: Expr a -> Bool
-keccakStoreInExpr expr = bad
-  where
-    FoundBad bad = execState (mapExprM findBadLoad expr) initFoundBad
-    findBadLoad :: Expr a-> State FoundBad (Expr a)
-    findBadLoad e = case e of
-      SLoad key store -> do
-        _ <- findBad key e
-        findBad store e
-      _ -> pure e
-    findBad e orig = if isNothing $ mapExprM keccakFinder e then do
-          put (FoundBad { bad = True })
-          pure orig
-        else pure orig
-    keccakFinder (Keccak {}) = Nothing
-    keccakFinder e = Just e
-
 defaultBuf :: Int -> Gen (Expr Buf)
 defaultBuf = genBuf (4_000_000)
 
@@ -4202,11 +4189,9 @@ genStorageKey = frequency
     ]
 
 genByteStringKey :: W256 -> Gen (ByteString)
-genByteStringKey len = oneof
-    [ do
-        b :: Word8 <- arbitrary
-        pure $ BS.pack ([ 0 | _ <- [0..(len-2)]] ++ [b `mod` 5])
-    ]
+genByteStringKey len = do
+  b :: Word8 <- arbitrary
+  pure $ BS.pack ([ 0 | _ <- [0..(len-2)]] ++ [b `mod` 5])
 
 -- GenWriteStorageLoad
 newtype GenWriteStorageLoad = GenWriteStorageLoad (Expr EWord)

From 9be4127ca5e449f4c6a20d6da0d4b932349d15fa Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Wed, 10 Jan 2024 18:39:48 +0100
Subject: [PATCH 16/20] Better variable name

---
 test/test.hs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/test/test.hs b/test/test.hs
index 96753ae67..e8c3b5288 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -4165,9 +4165,9 @@ genStorage :: Int -> Gen (Expr Storage)
 genStorage sz = genStorageMap True sz
 
 genStorageMap :: Bool -> Int -> Gen (Expr Storage)
-genStorageMap withMap 0 = oneof
-  [ liftM2 AbstractStore arbitrary (if withMap then arbitrary else pure Nothing)
-  , fmap ConcreteStore (if withMap then arbitrary else pure mempty)
+genStorageMap withIdx 0 = oneof
+  [ liftM2 AbstractStore arbitrary (if withIdx then arbitrary else pure Nothing)
+  , fmap ConcreteStore (if withIdx then arbitrary else pure mempty)
   ]
 genStorageMap withMap sz = liftM3 SStore key val subStore
   where

From 27bdc6396acc7b1885f513e28800dc0f3511d91b Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Wed, 10 Jan 2024 18:42:04 +0100
Subject: [PATCH 17/20] Actually, we should NOT generate arbitrary IDX for
 AbstractStore

---
 test/test.hs | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/test/test.hs b/test/test.hs
index e8c3b5288..3eb213c7d 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -4162,16 +4162,13 @@ genBuf bound sz = oneof
     subBuf = genBuf bound (sz `div` 10)
 
 genStorage :: Int -> Gen (Expr Storage)
-genStorage sz = genStorageMap True sz
-
-genStorageMap :: Bool -> Int -> Gen (Expr Storage)
-genStorageMap withIdx 0 = oneof
-  [ liftM2 AbstractStore arbitrary (if withIdx then arbitrary else pure Nothing)
-  , fmap ConcreteStore (if withIdx then arbitrary else pure mempty)
+genStorage 0 = oneof
+  [ liftM2 AbstractStore arbitrary (pure Nothing)
+  , fmap ConcreteStore arbitrary
   ]
-genStorageMap withMap sz = liftM3 SStore key val subStore
+genStorage sz = liftM3 SStore key val subStore
   where
-    subStore = genStorageMap withMap (sz `div` 10)
+    subStore = genStorage (sz `div` 10)
     val = defaultWord (sz `div` 5)
     key = genStorageKey
 

From af2aa62cdee6491475313b1183379e0e0b3e230c Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Thu, 11 Jan 2024 18:19:30 +0100
Subject: [PATCH 18/20] Fixing and cleanup thanks to @d-xo

---
 src/EVM/Expr.hs | 95 ++++++++++++++++++++++++-------------------------
 1 file changed, 46 insertions(+), 49 deletions(-)

diff --git a/src/EVM/Expr.hs b/src/EVM/Expr.hs
index ee2e28383..8982c5034 100644
--- a/src/EVM/Expr.hs
+++ b/src/EVM/Expr.hs
@@ -774,36 +774,25 @@ getLogicalIdx (GVar _) = internalError "cannot determine addr of a GVar"
 -- ** Whole Expression Simplification ** -----------------------------------------------------------
 
 
--- | Splits storage into logical substores (solc based heuristic) if it safe to do so
--- safe if
---   - all base stores must be one of:
---     - concrete + empty
---     - concrete + all keys < 256
---     - abstract
---   - all read / write indices must be one of:
---     - Lit < 256
---     - MappingSlot
---     - ArraySlotWithOffset
---     - ArraySlotZero
---  - mixing within the same store & load is not permitted
 data StorageType = SmallSlot | Array | Map | Mixed | UNK
   deriving (Show, Eq)
 
-newtype CollectStorageType = CollectStorageType {storeType :: StorageType }
-  deriving (Show)
+-- newtype CollectStorageType = CollectStorageType {storeType :: StorageType }
+--   deriving (Show)
+
+-- initCollectStorageType :: CollectStorageType
+-- initCollectStorageType = CollectStorageType {storeType = UNK }
 
-initCollectStorageType :: CollectStorageType
-initCollectStorageType = CollectStorageType {storeType = UNK }
 
-decomposeCheckTypes :: Expr a -> Bool
-decomposeCheckTypes inp = result.storeType /= Mixed
+safeToDecompose :: Expr a -> Bool
+safeToDecompose inp = result /= Mixed
   where
-    result = execState (decomposeCheckTypesRunner inp) initCollectStorageType
+    result = execState (safeToDecomposeRunner inp) UNK
 
-    decomposeCheckTypesRunner :: forall a. Expr a -> State CollectStorageType (Expr a)
-    decomposeCheckTypesRunner a = go a
+    safeToDecomposeRunner :: forall a. Expr a -> State StorageType ()
+    safeToDecomposeRunner a = go a
 
-    go :: forall b. Expr b -> State CollectStorageType (Expr b)
+    go :: forall b. Expr b -> State StorageType ()
     go e@(SLoad (MappingSlot {}) _) = setMap e
     go e@(SLoad (ArraySlotZero {}) _) = setArray e
     go e@(SLoad (ArraySlotWithOffset {}) _) = setArray e
@@ -814,46 +803,54 @@ decomposeCheckTypes inp = result.storeType /= Mixed
     go e@(SStore (ArraySlotWithOffset {}) _ _) = setArray e
     go e@(SStore (Lit x) _ _) | x < 256 = setSmall e
     go e@(SStore _ _ _) = setMixed e
-    go e = pure e
+    go _ = pure ()
 
     -- Helper functions for detecting mixed load/store
-    setMixed e = do
-      put CollectStorageType {storeType = Mixed}
-      pure e
-    setMap e = do
+    setMixed _ = do
+      put Mixed
+      pure ()
+    setMap _ = do
       s <- get
-      case s.storeType of
-        Array -> put s {storeType = Mixed}
-        SmallSlot -> put s {storeType = Mixed}
-        UNK -> put s {storeType = Map}
+      case s of
+        Array -> put Mixed
+        SmallSlot -> put Mixed
+        UNK -> put Map
         _ -> pure ()
-      pure e
-    setArray e = do
+      pure ()
+    setArray _ = do
       s <- get
-      case s.storeType of
-        Map -> put s {storeType = Mixed}
-        SmallSlot -> put s {storeType = Mixed}
-        UNK -> put s {storeType = Array}
+      case s of
+        Map -> put Mixed
+        SmallSlot -> put Mixed
+        UNK -> put Array
         _ -> pure ()
-      pure e
-    setSmall e = do
+      pure ()
+    setSmall _ = do
       s <- get
-      case s.storeType of
-        Map -> put s {storeType = Mixed}
-        Array -> put s {storeType = Mixed}
-        UNK -> put s {storeType = SmallSlot}
+      case s of
+        Map -> put Mixed
+        Array -> put Mixed
+        UNK -> put SmallSlot
         _ -> pure ()
-      pure e
+      pure ()
 
+-- | Splits storage into logical sub-stores if (1) all SLoad->SStore* chains are one of:
+--     (1a) Lit < 256, (1b) MappingSlot, (1c) ArraySlotWithOffset, (1d) ArraySlotZero
+--  and (2) there is no mixing of different types (e.g. Map with Array) within
+--  the same SStore -> SLoad* chain, and (3) there is no mixing of different array/map slots.
+--
+--  Mixing (2) and (3) are attempted to be prevented (if possible) as part of the rewrites
+--  done by the `readStorage` function that is ran before this. If there is still mixing here,
+--  we abort with a Nothing.
+--
+--  We do NOT rewrite stand-alone `SStore`-s (i.e. SStores that are not read), since
+--  they are often used to describe a post-state, and are not dispatched as-is to
+--  the solver
 decomposeStorage :: Expr a -> Maybe (Expr a)
 decomposeStorage = go
   where
-    -- NOTE: it's a bad idea to rewrite SStore on its own, we should rewrite SLoad, which calls SStore
-    --       this ensures that the rewrite happens to the whole SLoad... chain. If we rewrite SStore,
-    --       we will impact the returned "final state" of the system that will then be incorrect since
-    --       the final state actually contains all the Keccak-s
     go :: Expr a -> Maybe (Expr a)
-    go e@(SLoad origKey store) = if Prelude.not (decomposeCheckTypes e) then Nothing else tryRewrite origKey store
+    go e@(SLoad origKey store) = if Prelude.not (safeToDecompose e) then Nothing else tryRewrite origKey store
     go e = Just e
 
     tryRewrite :: Expr EWord -> Expr Storage -> Maybe (Expr EWord)

From 1e9bcbab4005ec0e232b57eb2967dd642772ccfc Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Thu, 11 Jan 2024 18:20:54 +0100
Subject: [PATCH 19/20] No need for these

---
 src/EVM/Expr.hs | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/src/EVM/Expr.hs b/src/EVM/Expr.hs
index 8982c5034..a14439dd1 100644
--- a/src/EVM/Expr.hs
+++ b/src/EVM/Expr.hs
@@ -777,13 +777,6 @@ getLogicalIdx (GVar _) = internalError "cannot determine addr of a GVar"
 data StorageType = SmallSlot | Array | Map | Mixed | UNK
   deriving (Show, Eq)
 
--- newtype CollectStorageType = CollectStorageType {storeType :: StorageType }
---   deriving (Show)
-
--- initCollectStorageType :: CollectStorageType
--- initCollectStorageType = CollectStorageType {storeType = UNK }
-
-
 safeToDecompose :: Expr a -> Bool
 safeToDecompose inp = result /= Mixed
   where

From 93760068293ca51d23b192627ee3cc0b41b8adb7 Mon Sep 17 00:00:00 2001
From: Mate Soos <mate.soos@ethereum.org>
Date: Tue, 16 Jan 2024 13:14:31 +0100
Subject: [PATCH 20/20] Fixing build

---
 test/test.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test.hs b/test/test.hs
index 3eb213c7d..8a63a411e 100644
--- a/test/test.hs
+++ b/test/test.hs
@@ -4203,7 +4203,7 @@ instance Arbitrary GenWriteStorageLoad where
       genStorageLoad :: Int -> Gen (Expr EWord)
       genStorageLoad sz = liftM2 SLoad key subStore
         where
-          subStore = genStorageMap False (sz `div` 10)
+          subStore = genStorage (sz `div` 10)
           key = genStorageKey
 
 data Invocation