Adapt storage capturing in OpcodeTracer for modularized execution

hedera-mirror-web3/src/main/java/com/hedera/mirror/web3/evm/contracts/execution/traceability/OpcodeTracer.java

@@ -42,11 +42,9 @@
 import jakarta.inject.Named;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.Set;
 import java.util.TreeMap;
 import java.util.stream.Collectors;
 import lombok.CustomLog;
@@ -269,35 +267,25 @@ private Bytes formatRevertReason(final Bytes revertReason) {
         return BytesDecoder.getAbiEncodedRevertReason(revertReason);
     }
 
-    private Optional<Map<Bytes, Bytes>> getStorageUpdates(Address accountAddress) {
-        Map<Bytes, Bytes> storageUpdates = new HashMap<>();
+    private Map<Bytes, Bytes> getModularizedUpdatedStorage(Address accountAddress) {
+        Map<Bytes, Bytes> storageUpdates = new TreeMap<>();
         MapWritableStates states = (MapWritableStates) mirrorNodeState.getWritableStates(ContractService.NAME);
 
         try {
+            var accountContractID = EntityIdUtils.toContractID(accountAddress);
             var storageState = states.get(ContractStorageReadableKVState.KEY);
-            Set<SlotKey> modifiedKeys = storageState.modifiedKeys().stream()
+            storageState.modifiedKeys().stream()
                     .filter(SlotKey.class::isInstance)
                     .map(SlotKey.class::cast)
-                    .filter(SlotKey::hasContractID)
-                    .collect(Collectors.toSet());
-
-            if (modifiedKeys.isEmpty()) {
-                return Optional.empty();
-            }
-
-            var accountContractID = EntityIdUtils.toContractID(accountAddress);
-            for (SlotKey slotKey : modifiedKeys) {
-                if (!slotKey.contractID().equals(accountContractID)) {
-                    continue;
-                }
-
-                SlotValue slotValue = (SlotValue) storageState.get(slotKey);
-                if (slotValue != null) {
-                    storageUpdates.put(
-                            Bytes.wrap(slotKey.key().toByteArray()),
-                            Bytes.wrap(slotValue.value().toByteArray()));
-                }
-            }
+                    .filter(slotKey -> slotKey.contractID().equals(accountContractID))
+                    .forEach(slotKey -> {
+                        SlotValue slotValue = (SlotValue) storageState.get(slotKey);
+                        if (slotValue != null) {
+                            storageUpdates.put(
+                                    Bytes.wrap(slotKey.key().toByteArray()),
+                                    Bytes.wrap(slotValue.value().toByteArray()));
+                        }
+                    });
         } catch (IllegalArgumentException e) {
             log.warn(
                     "Failed to retrieve modified storage keys for service: {}, key: {}",
@@ -306,10 +294,6 @@ private Optional<Map<Bytes, Bytes>> getStorageUpdates(Address accountAddress) {
                     e);
         }
 
-        return storageUpdates.isEmpty() ? Optional.empty() : Optional.of(storageUpdates);
-    }
-
-    private Map<Bytes, Bytes> getModularizedUpdatedStorage(Address accountAddress) {
-        return getStorageUpdates(accountAddress).map(TreeMap::new).orElseGet(TreeMap::new);
+        return storageUpdates;
     }
 }

hedera-mirror-web3/src/test/java/com/hedera/mirror/web3/evm/contracts/execution/traceability/OpcodeTracerTest.java

@@ -548,32 +548,6 @@ void shouldReturnEmptyStorageWhenThereAreNoUpdates() {
         assertThat(opcode.storage()).isEmpty();
     }
 
-    @Test
-    @DisplayName(
-            "given storage is enabled in tracer options, should skip slotKey when hasContractID is false for modularized services")
-    void shouldSkipSlotKeyWhenHasContractIDIsFalse() {
-        // Given
-        tracerOptions = tracerOptions.toBuilder().storage(true).build();
-        when(mirrorNodeEvmProperties.isModularizedServices()).thenReturn(true);
-        frame = setupInitialFrame(tracerOptions);
-
-        MapWritableStates mockStates = mock(MapWritableStates.class);
-        when(mirrorNodeState.getWritableStates(CONTRACT_SERVICE)).thenReturn(mockStates);
-        WritableKVState<SlotKey, SlotValue> mockStorageState = mock(WritableKVState.class);
-        doReturn(mockStorageState).when(mockStates).get(ContractStorageReadableKVState.KEY);
-
-        SlotKey slotKeyWithoutContractID = mock(SlotKey.class);
-        when(slotKeyWithoutContractID.hasContractID()).thenReturn(false);
-
-        when(mockStorageState.modifiedKeys()).thenReturn(Set.of(slotKeyWithoutContractID));
-
-        // When
-        final Opcode opcode = executeOperation(frame);
-
-        // Then
-        assertThat(opcode.storage()).isEmpty();
-    }
-
     @Test
     @DisplayName(
             "given storage is enabled in tracer options, should skip slotKey when contract address does not match for modularized services")
@@ -1055,8 +1029,6 @@ private ContractAction contractAction(
     private SlotKey createMockSlotKey(Address contractAddress) {
         SlotKey slotKey = mock(SlotKey.class);
 
-        when(slotKey.hasContractID()).thenReturn(true);
-
         ContractID testContractId = com.hedera.hapi.node.base.ContractID.newBuilder()
                 .contractNum(EntityIdUtils.numFromEvmAddress(contractAddress.toArray()))
                 .build();