From 49c83c994c9e72cd514845c432a0719f2a1ffcc9 Mon Sep 17 00:00:00 2001
From: A_A <21040751+Otto-AA@users.noreply.github.com>
Date: Sun, 9 Jun 2024 11:03:50 +0200
Subject: [PATCH] fix: only add the first stack access to the same index

---
 tests/e2e/test_sample_traces_analysis.py         |  1 -
 tests/parser/instructions/test_instructions.py   | 12 +-----------
 traces_analyzer/parser/storage/storage_writes.py | 13 ++++++++++++-
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/tests/e2e/test_sample_traces_analysis.py b/tests/e2e/test_sample_traces_analysis.py
index 12ebcbf..f6ffc85 100644
--- a/tests/e2e/test_sample_traces_analysis.py
+++ b/tests/e2e/test_sample_traces_analysis.py
@@ -68,7 +68,6 @@ def test_sample_traces_analysis_e2e(
             tod_source_analyzer,
             instruction_input_analyzer,
         ],
-        # TODO: why is the reverse one first? check this and document it
         transactions=(transactions_actual, transactions_reverse),
     )
 
diff --git a/tests/parser/instructions/test_instructions.py b/tests/parser/instructions/test_instructions.py
index 3c95399..f8eed35 100644
--- a/tests/parser/instructions/test_instructions.py
+++ b/tests/parser/instructions/test_instructions.py
@@ -1011,9 +1011,7 @@ def test_codecopy() -> None:
     codecopy = _test_parse_instruction(CODECOPY, env, oracle)
 
     accesses = codecopy.get_accesses()
-    # TODO: should be three
-    # if possible, rewrite access merging to merge stack accesses with the same index and value
-    assert len(accesses.stack) == 4
+    assert len(accesses.stack) == 3
 
     writes = codecopy.get_writes()
     assert len(writes.memory) == 1
@@ -1123,8 +1121,6 @@ def test_returndatacopy() -> None:
     assert writes.memory[0].value.depends_on_instruction_indexes() == {1234}
 
 
-# TODO: update call enter tests
-# TODO: add call (immediate) exists tests
 def test_call_enter() -> None:
     call_context = _test_root()
     env = mock_env(
@@ -1181,8 +1177,6 @@ def test_call_enter() -> None:
     assert call.get_data()["updates_storage_address"]
 
 
-# TODO: update staticcall enter tests
-# TODO: add staticcall (immediate) exists tests
 def test_staticcall_enter() -> None:
     env = mock_env(
         storage_step_index=1,
@@ -1215,8 +1209,6 @@ def test_staticcall_enter() -> None:
     assert staticcall.get_data()["updates_storage_address"]
 
 
-# TODO: update callcode enter tests
-# TODO: add callcode (immediate) exists tests
 def test_callcode_enter() -> None:
     call_context = _test_root()
     env = mock_env(
@@ -1272,8 +1264,6 @@ def test_callcode_enter() -> None:
     assert not callcode.get_data()["updates_storage_address"]
 
 
-# TODO: update delegatecall enter tests
-# TODO: add delegatecall (immediate) exists tests
 def test_delegatecall_enter() -> None:
     env = mock_env(
         storage_step_index=1,
diff --git a/traces_analyzer/parser/storage/storage_writes.py b/traces_analyzer/parser/storage/storage_writes.py
index ba2a4f5..24254eb 100644
--- a/traces_analyzer/parser/storage/storage_writes.py
+++ b/traces_analyzer/parser/storage/storage_writes.py
@@ -250,7 +250,7 @@ def merge(accesses: list["StorageAccesses"]) -> "StorageAccesses":
             return_data_access = return_data_access or access.return_data
 
         return StorageAccesses(
-            stack=stack_accesses,
+            stack=StorageAccesses.unify_stack_accesses(stack_accesses),
             memory=memory_accesses,
             persistent_storage=persistent_storage_accesses,
             transient_storage=transient_storage_accesses,
@@ -259,3 +259,14 @@ def merge(accesses: list["StorageAccesses"]) -> "StorageAccesses":
             callvalue=callvalue_accesses,
             return_data=return_data_access,
         )
+
+    @staticmethod
+    def unify_stack_accesses(accesses: list[StackAccess]) -> list[StackAccess]:
+        indices = set()
+        result = []
+        for access in accesses:
+            if access.index not in indices:
+                indices.add(access.index)
+                result.append(access)
+
+        return result