diff --git a/test/state/state.cpp b/test/state/state.cpp
index 619c8ae20..4038a8401 100644
--- a/test/state/state.cpp
+++ b/test/state/state.cpp
@@ -142,10 +142,9 @@ StateDiff State::build_diff(evmc_revision rev) const
 
         // Output only the new code.
         // TODO: Output also the code hash. It will be needed for DB update and MPT hash.
+        // TODO: Only use "code_changed" flag to indicate the code change.
         if ((m.just_created && !m.code.empty()) || m.code_changed)
             a.code = m.code;
-        if (m.code_changed && m.code.empty())
-            a.code_cleared = true;
 
         for (const auto& [k, v] : m.storage)
         {
diff --git a/test/state/state_diff.hpp b/test/state/state_diff.hpp
index 29af3f420..426880475 100644
--- a/test/state/state_diff.hpp
+++ b/test/state/state_diff.hpp
@@ -30,11 +30,8 @@ struct StateDiff
         /// TODO: Currently it is not guaranteed the value is different from the initial one.
         uint256 balance;
 
-        /// New account code. If empty, it means the code has not changed.
-        bytes code;
-
-        /// If true, account code was emptied due to resetting EIP-7702 delegation.
-        bool code_cleared = false;
+        /// New or modified account code. If bytes are empty, it means the code has been cleared.
+        std::optional<bytes> code;
 
         /// The list of the account's storage modifications: key => new value.
         /// The value 0 means the storage entry is deleted.
diff --git a/test/state/test_state.cpp b/test/state/test_state.cpp
index 060e106be..fc8300842 100644
--- a/test/state/test_state.cpp
+++ b/test/state/test_state.cpp
@@ -34,10 +34,8 @@ void TestState::apply(const state::StateDiff& diff)
         auto& a = (*this)[m.addr];
         a.nonce = m.nonce;
         a.balance = m.balance;
-        if (!m.code.empty())
-            a.code = m.code;  // TODO: Consider taking rvalue ref to avoid code copy.
-        if (m.code_cleared)
-            a.code.clear();
+        if (m.code.has_value())
+            a.code = *m.code;  // TODO: Consider taking rvalue ref to avoid code copy.
         for (const auto& [k, v] : m.modified_storage)
         {
             if (v)