From 73533370fd55a5e1bc70e66c4625fa4f64252aa2 Mon Sep 17 00:00:00 2001
From: Steven Gu <asongala@163.com>
Date: Mon, 8 Aug 2022 17:41:02 +0800
Subject: [PATCH 01/14] Implement opcode `BALANCE`.

---
 src/zkevm_specs/evm/execution/__init__.py |   2 +
 src/zkevm_specs/evm/execution/balance.py  |  25 +++++
 tests/evm/test_balance.py                 | 112 ++++++++++++++++++++++
 3 files changed, 139 insertions(+)
 create mode 100644 src/zkevm_specs/evm/execution/balance.py
 create mode 100644 tests/evm/test_balance.py

diff --git a/src/zkevm_specs/evm/execution/__init__.py b/src/zkevm_specs/evm/execution/__init__.py
index d1b4db432..f7b1af51f 100644
--- a/src/zkevm_specs/evm/execution/__init__.py
+++ b/src/zkevm_specs/evm/execution/__init__.py
@@ -10,6 +10,7 @@
 from .add_sub import *
 from .addmod import *
 from .mulmod import *
+from .balance import *
 from .block_ctx import *
 from .call import *
 from .calldatasize import *
@@ -51,6 +52,7 @@
     ExecutionState.MUL: mul_div_mod,
     ExecutionState.NOT: not_opcode,
     ExecutionState.ORIGIN: origin,
+    ExecutionState.BALANCE: balance,
     ExecutionState.CALLER: caller,
     ExecutionState.CALLVALUE: callvalue,
     ExecutionState.CALLDATACOPY: calldatacopy,
diff --git a/src/zkevm_specs/evm/execution/balance.py b/src/zkevm_specs/evm/execution/balance.py
new file mode 100644
index 000000000..1f23d893d
--- /dev/null
+++ b/src/zkevm_specs/evm/execution/balance.py
@@ -0,0 +1,25 @@
+from ...util import EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS, FQ, N_BYTES_ACCOUNT_ADDRESS
+from ..instruction import Instruction, Transition
+from ..opcode import Opcode
+from ..table import AccountFieldTag, CallContextFieldTag
+
+
+def balance(instruction: Instruction):
+    opcode = instruction.opcode_lookup(True)
+    instruction.constrain_equal(opcode, Opcode.BALANCE)
+
+    address = instruction.rlc_to_fq(instruction.stack_pop(), N_BYTES_ACCOUNT_ADDRESS)
+
+    tx_id = instruction.call_context_lookup(CallContextFieldTag.TxId)
+    is_warm = instruction.add_account_to_access_list(tx_id, address, instruction.reversion_info())
+
+    balance = instruction.account_read(address, AccountFieldTag.Balance)
+    instruction.constrain_equal(instruction.stack_push(), balance)
+
+    instruction.step_state_transition_in_same_context(
+        opcode,
+        rw_counter=Transition.delta(7),
+        program_counter=Transition.delta(1),
+        stack_pointer=Transition.same(),
+        dynamic_gas_cost=instruction.select(is_warm, FQ(0), FQ(EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS)),
+    )
diff --git a/tests/evm/test_balance.py b/tests/evm/test_balance.py
new file mode 100644
index 000000000..da7f1fd56
--- /dev/null
+++ b/tests/evm/test_balance.py
@@ -0,0 +1,112 @@
+import pytest
+
+from zkevm_specs.evm import (
+    AccountFieldTag,
+    Block,
+    Bytecode,
+    CallContextFieldTag,
+    ExecutionState,
+    RWDictionary,
+    StepState,
+    Tables,
+    verify_steps,
+)
+from zkevm_specs.util import (
+    EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS,
+    GAS_COST_WARM_ACCESS,
+    RLC,
+    U160,
+    U256,
+    rand_address,
+    rand_fq,
+    rand_range,
+    rand_word,
+)
+
+TESTING_DATA = [
+    (0x30000, 0, True, True),
+    (0x30000, 0, False, True),
+    (0x30000, 200, True, True),
+    (0x30000, 200, False, True),
+    (
+        rand_address(),
+        rand_word(),
+        rand_range(2) == 0,
+        True,  # persistent call
+    ),
+    (
+        rand_address(),
+        rand_word(),
+        rand_range(2) == 0,
+        False,  # reverted call
+    ),
+]
+
+
+@pytest.mark.parametrize("address, balance, is_warm, is_persistent", TESTING_DATA)
+def test_balance(address: U160, balance: U256, is_warm: bool, is_persistent: bool):
+    randomness = rand_fq()
+
+    tx_id = 1
+    call_id = 1
+
+    rw_counter_end_of_reversion = 0 if is_persistent else 7
+    reversible_write_counter = 0
+
+    rw_table = set(
+        RWDictionary(1)
+        .stack_read(call_id, 1023, RLC(address, randomness))
+        .call_context_read(tx_id, CallContextFieldTag.TxId, tx_id)
+        .call_context_read(
+            tx_id, CallContextFieldTag.RwCounterEndOfReversion, rw_counter_end_of_reversion
+        )
+        .call_context_read(tx_id, CallContextFieldTag.IsPersistent, is_persistent)
+        .tx_access_list_account_write(
+            tx_id,
+            address,
+            True,
+            is_warm,
+            rw_counter_of_reversion=rw_counter_end_of_reversion - reversible_write_counter,
+        )
+        .account_read(address, AccountFieldTag.Balance, RLC(balance, randomness))
+        .stack_write(call_id, 1023, RLC(balance, randomness))
+        .rws
+    )
+
+    bytecode = Bytecode().balance()
+    tables = Tables(
+        block_table=Block(),
+        tx_table=set(),
+        bytecode_table=set(bytecode.table_assignments(randomness)),
+        rw_table=rw_table,
+    )
+
+    bytecode_hash = RLC(bytecode.hash(), randomness)
+    verify_steps(
+        randomness=randomness,
+        tables=tables,
+        steps=[
+            StepState(
+                execution_state=ExecutionState.BALANCE,
+                rw_counter=1,
+                call_id=1,
+                is_root=True,
+                is_create=False,
+                code_hash=bytecode_hash,
+                program_counter=0,
+                stack_pointer=1023,
+                gas_left=GAS_COST_WARM_ACCESS + (not is_warm) * EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS,
+            ),
+            StepState(
+                execution_state=ExecutionState.STOP if is_persistent else ExecutionState.REVERT,
+                rw_counter=8,
+                call_id=1,
+                is_root=True,
+                is_create=False,
+                code_hash=bytecode_hash,
+                program_counter=1,
+                stack_pointer=1023,
+                gas_left=0,
+            ),
+        ],
+    )

From b0ee8b76053070cb6d73ea8b125a1233c811d000 Mon Sep 17 00:00:00 2001
From: Steven Gu <asongala@163.com>
Date: Fri, 12 Aug 2022 10:24:51 +0800
Subject: [PATCH 02/14] Add Markdown doc of opcode `BALANCE`.

---
 specs/opcode/31BALANCE.md | 49 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)
 create mode 100644 specs/opcode/31BALANCE.md

diff --git a/specs/opcode/31BALANCE.md b/specs/opcode/31BALANCE.md
new file mode 100644
index 000000000..e5be39a36
--- /dev/null
+++ b/specs/opcode/31BALANCE.md
@@ -0,0 +1,49 @@
+# BALANCE opcode
+
+## Procedure
+
+The `BALANCE` opcode gets balance of the given account.
+
+## EVM behaviour
+
+The `BALANCE` opcode pops `address` (20 bytes of data) off the stack and pushes
+the balance of the corresponding account onto the stack. If the given account
+doesn't exist, then it will push 0 onto the stack instead.
+
+## Circuit behaviour
+
+1. Construct call context table in rw table.
+2. Do bus-mapping lookup for stack read, call context read and account read
+   operations.
+3. Do bus-mapping lookup for transaction access list write and stack write
+   operations.
+
+## Constraints
+
+1. opId = 0x31
+2. State transition:
+   - gc + 7 (1 stack read, 1 stack write, 3 call context reads, 1 account reads,
+     1 transaction access list write)
+   - stack_pointer + 0 (one pop and one push)
+   - pc + 1
+   - gas:
+     - the accessed `address` is warm: GAS_COST_WARM_ACCESS
+     - the accessed `address` is cold: GAS_COST_ACCOUNT_COLD_ACCESS
+3. Lookups: 7
+   - `address` is popped from the stack.
+   - 3 reads from call context for `tx_id`, `rw_counter_end_of_reversion`, and
+     `is_persistent`.
+   - `address` is added to the transaction access list if not already present.
+   - `balance` is read from the given account. Set to 0 if it doesn't exist.
+   - the `balance` result is pushed onto the stack.
+4. Additional Constraints
+   - value `is_warm` matches the gas cost for this opcode.
+
+## Exceptions
+
+1. stack underflow: if the stack starts empty
+2. out of gas: remaining gas is not enough
+
+## Code
+
+Please refer to `src/zkevm_specs/evm/execution/balance.py`.

From 5071cbec66458f9a6a8b3e93b2c166b37ce6d6b9 Mon Sep 17 00:00:00 2001
From: Steven <asongala@163.com>
Date: Wed, 17 Aug 2022 16:19:00 +0800
Subject: [PATCH 03/14] Update specs/opcode/31BALANCE.md
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Carlos Pérez <37264926+CPerezz@users.noreply.github.com>
---
 specs/opcode/31BALANCE.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/specs/opcode/31BALANCE.md b/specs/opcode/31BALANCE.md
index e5be39a36..7a62ecbc1 100644
--- a/specs/opcode/31BALANCE.md
+++ b/specs/opcode/31BALANCE.md
@@ -35,7 +35,7 @@ doesn't exist, then it will push 0 onto the stack instead.
      `is_persistent`.
    - `address` is added to the transaction access list if not already present.
    - `balance` is read from the given account. Set to 0 if it doesn't exist.
-   - the `balance` result is pushed onto the stack.
+   - The BALANCE result is in the top of the stack.
 4. Additional Constraints
    - value `is_warm` matches the gas cost for this opcode.
 

From 270b26172f5b88c0efb3757160b3388efdd010fe Mon Sep 17 00:00:00 2001
From: Steven Gu <asongala@163.com>
Date: Thu, 18 Aug 2022 17:41:58 +0800
Subject: [PATCH 04/14] Update `bus-mapping lookups` in Markdown.

---
 specs/opcode/31BALANCE.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/specs/opcode/31BALANCE.md b/specs/opcode/31BALANCE.md
index 7a62ecbc1..6174bac9a 100644
--- a/specs/opcode/31BALANCE.md
+++ b/specs/opcode/31BALANCE.md
@@ -29,13 +29,13 @@ doesn't exist, then it will push 0 onto the stack instead.
    - gas:
      - the accessed `address` is warm: GAS_COST_WARM_ACCESS
      - the accessed `address` is cold: GAS_COST_ACCOUNT_COLD_ACCESS
-3. Lookups: 7
-   - `address` is popped from the stack.
+3. Lookups: 7 busmapping lookups
+   - `address` is at top of the stack.
    - 3 reads from call context for `tx_id`, `rw_counter_end_of_reversion`, and
      `is_persistent`.
    - `address` is added to the transaction access list if not already present.
    - `balance` is read from the given account. Set to 0 if it doesn't exist.
-   - The BALANCE result is in the top of the stack.
+   - The BALANCE result is at the new top of the stack.
 4. Additional Constraints
    - value `is_warm` matches the gas cost for this opcode.
 

From 931bb884d587aa890c3568ecbcec57c891f9914f Mon Sep 17 00:00:00 2001
From: Steven Gu <asongala@163.com>
Date: Fri, 19 Aug 2022 09:29:30 +0800
Subject: [PATCH 05/14] Fix to get `nonce`, `balance` and `code_hash` to check
 account existence.

---
 specs/opcode/31BALANCE.md                    | 12 +++++++----
 src/zkevm_specs/evm/execution/balance.py     | 21 +++++++++++++++++---
 src/zkevm_specs/evm/execution/extcodehash.py |  3 +++
 3 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/specs/opcode/31BALANCE.md b/specs/opcode/31BALANCE.md
index 6174bac9a..4a7cf5031 100644
--- a/specs/opcode/31BALANCE.md
+++ b/specs/opcode/31BALANCE.md
@@ -8,7 +8,8 @@ The `BALANCE` opcode gets balance of the given account.
 
 The `BALANCE` opcode pops `address` (20 bytes of data) off the stack and pushes
 the balance of the corresponding account onto the stack. If the given account
-doesn't exist, then it will push 0 onto the stack instead.
+doesn't exist (i.e. has nonce = 0, balance = 0 and no code), then it will push 0
+onto the stack instead.
 
 ## Circuit behaviour
 
@@ -22,19 +23,22 @@ doesn't exist, then it will push 0 onto the stack instead.
 
 1. opId = 0x31
 2. State transition:
-   - gc + 7 (1 stack read, 1 stack write, 3 call context reads, 1 account reads,
+   - gc + 9 (1 stack read, 1 stack write, 3 call context reads, 3 account reads,
      1 transaction access list write)
    - stack_pointer + 0 (one pop and one push)
    - pc + 1
    - gas:
      - the accessed `address` is warm: GAS_COST_WARM_ACCESS
      - the accessed `address` is cold: GAS_COST_ACCOUNT_COLD_ACCESS
-3. Lookups: 7 busmapping lookups
+3. Lookups: 9 busmapping lookups
    - `address` is at top of the stack.
    - 3 reads from call context for `tx_id`, `rw_counter_end_of_reversion`, and
      `is_persistent`.
    - `address` is added to the transaction access list if not already present.
-   - `balance` is read from the given account. Set to 0 if it doesn't exist.
+   - nonce of account is `nonce`.
+   - `nonce` is read from account. It is used to check account existence.
+   - `code_hash` is read from account. It is used to check account existence.
+   - `balance` is read from account. Set to 0 if it doesn't exist.
    - The BALANCE result is at the new top of the stack.
 4. Additional Constraints
    - value `is_warm` matches the gas cost for this opcode.
diff --git a/src/zkevm_specs/evm/execution/balance.py b/src/zkevm_specs/evm/execution/balance.py
index 1f23d893d..abadfd2ba 100644
--- a/src/zkevm_specs/evm/execution/balance.py
+++ b/src/zkevm_specs/evm/execution/balance.py
@@ -1,4 +1,4 @@
-from ...util import EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS, FQ, N_BYTES_ACCOUNT_ADDRESS
+from ...util import EMPTY_CODE_HASH, EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS, FQ, N_BYTES_ACCOUNT_ADDRESS
 from ..instruction import Instruction, Transition
 from ..opcode import Opcode
 from ..table import AccountFieldTag, CallContextFieldTag
@@ -13,12 +13,27 @@ def balance(instruction: Instruction):
     tx_id = instruction.call_context_lookup(CallContextFieldTag.TxId)
     is_warm = instruction.add_account_to_access_list(tx_id, address, instruction.reversion_info())
 
+    nonce = instruction.account_read(address, AccountFieldTag.Nonce)
     balance = instruction.account_read(address, AccountFieldTag.Balance)
-    instruction.constrain_equal(instruction.stack_push(), balance)
+    code_hash = instruction.account_read(address, AccountFieldTag.CodeHash)
+
+    # Get flag of the account existance.
+    # TODO: Need to update after the below issue is fixed -
+    # https://github.com/privacy-scaling-explorations/zkevm-specs/issues/249.
+    is_empty = (
+        instruction.is_zero(nonce)
+        * instruction.is_zero(balance)
+        * instruction.is_equal(code_hash, instruction.rlc_encode(EMPTY_CODE_HASH, 32))
+    )
+
+    instruction.constrain_equal(
+        instruction.select(is_empty, FQ(0), balance.expr()),
+        instruction.stack_push(),
+    )
 
     instruction.step_state_transition_in_same_context(
         opcode,
-        rw_counter=Transition.delta(7),
+        rw_counter=Transition.delta(9),
         program_counter=Transition.delta(1),
         stack_pointer=Transition.same(),
         dynamic_gas_cost=instruction.select(is_warm, FQ(0), FQ(EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS)),
diff --git a/src/zkevm_specs/evm/execution/extcodehash.py b/src/zkevm_specs/evm/execution/extcodehash.py
index 18ea4198d..0086eb8c7 100644
--- a/src/zkevm_specs/evm/execution/extcodehash.py
+++ b/src/zkevm_specs/evm/execution/extcodehash.py
@@ -17,6 +17,9 @@ def extcodehash(instruction: Instruction):
     balance = instruction.account_read(address, AccountFieldTag.Balance)
     code_hash = instruction.account_read(address, AccountFieldTag.CodeHash)
 
+    # Get flag of the account existance.
+    # TODO: Need to update after the below issue is fixed -
+    # https://github.com/privacy-scaling-explorations/zkevm-specs/issues/249.
     is_empty = (
         instruction.is_zero(nonce)
         * instruction.is_zero(balance)

From 4ddc636f394f67f286b02eda9ebd426460bb6bd6 Mon Sep 17 00:00:00 2001
From: Steven Gu <asongala@163.com>
Date: Fri, 19 Aug 2022 09:59:02 +0800
Subject: [PATCH 06/14] Fix failed test cases.

---
 tests/evm/test_balance.py | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/tests/evm/test_balance.py b/tests/evm/test_balance.py
index da7f1fd56..90d90b288 100644
--- a/tests/evm/test_balance.py
+++ b/tests/evm/test_balance.py
@@ -12,45 +12,57 @@
     verify_steps,
 )
 from zkevm_specs.util import (
+    EMPTY_CODE_HASH,
     EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS,
     GAS_COST_WARM_ACCESS,
     RLC,
     U160,
     U256,
+    keccak256,
     rand_address,
+    rand_bytes,
     rand_fq,
     rand_range,
     rand_word,
 )
 
 TESTING_DATA = [
-    (0x30000, 0, True, True),
-    (0x30000, 0, False, True),
-    (0x30000, 200, True, True),
-    (0x30000, 200, False, True),
+    (0x30000, 0, 0, bytes(), True, True),
+    (0x30000, 0, 0, bytes(), False, True),
+    (0x30000, 200, 1, bytes(), True, True),
+    (0x30000, 200, 0, bytes([10, 10]), False, True),
     (
         rand_address(),
         rand_word(),
+        rand_word(),
+        rand_bytes(100),
         rand_range(2) == 0,
         True,  # persistent call
     ),
     (
         rand_address(),
         rand_word(),
+        rand_word(),
+        rand_bytes(100),
         rand_range(2) == 0,
         False,  # reverted call
     ),
 ]
 
 
-@pytest.mark.parametrize("address, balance, is_warm, is_persistent", TESTING_DATA)
-def test_balance(address: U160, balance: U256, is_warm: bool, is_persistent: bool):
+@pytest.mark.parametrize("address, balance, nonce, code, is_warm, is_persistent", TESTING_DATA)
+def test_balance(
+    address: U160, balance: U256, nonce: U256, code: bytes, is_warm: bool, is_persistent: bool
+):
     randomness = rand_fq()
 
+    code_hash = int.from_bytes(keccak256(code), "big")
+    result = 0 if (balance == 0 and nonce == 0 and code_hash == EMPTY_CODE_HASH) else balance
+
     tx_id = 1
     call_id = 1
 
-    rw_counter_end_of_reversion = 0 if is_persistent else 7
+    rw_counter_end_of_reversion = 0 if is_persistent else 9
     reversible_write_counter = 0
 
     rw_table = set(
@@ -68,8 +80,10 @@ def test_balance(address: U160, balance: U256, is_warm: bool, is_persistent: boo
             is_warm,
             rw_counter_of_reversion=rw_counter_end_of_reversion - reversible_write_counter,
         )
+        .account_read(address, AccountFieldTag.Nonce, RLC(nonce, randomness))
         .account_read(address, AccountFieldTag.Balance, RLC(balance, randomness))
-        .stack_write(call_id, 1023, RLC(balance, randomness))
+        .account_read(address, AccountFieldTag.CodeHash, RLC(code_hash, randomness))
+        .stack_write(call_id, 1023, RLC(result, randomness))
         .rws
     )
 
@@ -99,7 +113,7 @@ def test_balance(address: U160, balance: U256, is_warm: bool, is_persistent: boo
             ),
             StepState(
                 execution_state=ExecutionState.STOP if is_persistent else ExecutionState.REVERT,
-                rw_counter=8,
+                rw_counter=10,
                 call_id=1,
                 is_root=True,
                 is_create=False,

From e1571d43c1388492d8c9ed1edc16e0a5e65af8ce Mon Sep 17 00:00:00 2001
From: Steven Gu <asongala@163.com>
Date: Wed, 5 Oct 2022 20:12:37 +0800
Subject: [PATCH 07/14] Fix to check account existence by
 `AccountFieldTag.NonExisting`.

---
 specs/opcode/31BALANCE.md                | 12 +++++-----
 src/zkevm_specs/evm/execution/balance.py | 18 ++++-----------
 tests/evm/test_balance.py                | 28 ++++++++++--------------
 3 files changed, 20 insertions(+), 38 deletions(-)

diff --git a/specs/opcode/31BALANCE.md b/specs/opcode/31BALANCE.md
index 4a7cf5031..cfc3751d3 100644
--- a/specs/opcode/31BALANCE.md
+++ b/specs/opcode/31BALANCE.md
@@ -8,8 +8,8 @@ The `BALANCE` opcode gets balance of the given account.
 
 The `BALANCE` opcode pops `address` (20 bytes of data) off the stack and pushes
 the balance of the corresponding account onto the stack. If the given account
-doesn't exist (i.e. has nonce = 0, balance = 0 and no code), then it will push 0
-onto the stack instead.
+doesn't exist (by checking non existing flag), then it will push 0 onto the
+stack instead.
 
 ## Circuit behaviour
 
@@ -23,21 +23,19 @@ onto the stack instead.
 
 1. opId = 0x31
 2. State transition:
-   - gc + 9 (1 stack read, 1 stack write, 3 call context reads, 3 account reads,
+   - gc + 8 (1 stack read, 1 stack write, 3 call context reads, 2 account reads,
      1 transaction access list write)
    - stack_pointer + 0 (one pop and one push)
    - pc + 1
    - gas:
      - the accessed `address` is warm: GAS_COST_WARM_ACCESS
      - the accessed `address` is cold: GAS_COST_ACCOUNT_COLD_ACCESS
-3. Lookups: 9 busmapping lookups
+3. Lookups: 8 busmapping lookups
    - `address` is at top of the stack.
    - 3 reads from call context for `tx_id`, `rw_counter_end_of_reversion`, and
      `is_persistent`.
    - `address` is added to the transaction access list if not already present.
-   - nonce of account is `nonce`.
-   - `nonce` is read from account. It is used to check account existence.
-   - `code_hash` is read from account. It is used to check account existence.
+   - `is_non_existing` is read to check account existence.
    - `balance` is read from account. Set to 0 if it doesn't exist.
    - The BALANCE result is at the new top of the stack.
 4. Additional Constraints
diff --git a/src/zkevm_specs/evm/execution/balance.py b/src/zkevm_specs/evm/execution/balance.py
index abadfd2ba..d7944f5b3 100644
--- a/src/zkevm_specs/evm/execution/balance.py
+++ b/src/zkevm_specs/evm/execution/balance.py
@@ -1,4 +1,4 @@
-from ...util import EMPTY_CODE_HASH, EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS, FQ, N_BYTES_ACCOUNT_ADDRESS
+from ...util import EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS, FQ, N_BYTES_ACCOUNT_ADDRESS
 from ..instruction import Instruction, Transition
 from ..opcode import Opcode
 from ..table import AccountFieldTag, CallContextFieldTag
@@ -13,27 +13,17 @@ def balance(instruction: Instruction):
     tx_id = instruction.call_context_lookup(CallContextFieldTag.TxId)
     is_warm = instruction.add_account_to_access_list(tx_id, address, instruction.reversion_info())
 
-    nonce = instruction.account_read(address, AccountFieldTag.Nonce)
+    is_non_existing = instruction.account_read(address, AccountFieldTag.NonExisting)
     balance = instruction.account_read(address, AccountFieldTag.Balance)
-    code_hash = instruction.account_read(address, AccountFieldTag.CodeHash)
-
-    # Get flag of the account existance.
-    # TODO: Need to update after the below issue is fixed -
-    # https://github.com/privacy-scaling-explorations/zkevm-specs/issues/249.
-    is_empty = (
-        instruction.is_zero(nonce)
-        * instruction.is_zero(balance)
-        * instruction.is_equal(code_hash, instruction.rlc_encode(EMPTY_CODE_HASH, 32))
-    )
 
     instruction.constrain_equal(
-        instruction.select(is_empty, FQ(0), balance.expr()),
+        instruction.select(is_non_existing.expr(), FQ(0), balance.expr()),
         instruction.stack_push(),
     )
 
     instruction.step_state_transition_in_same_context(
         opcode,
-        rw_counter=Transition.delta(9),
+        rw_counter=Transition.delta(8),
         program_counter=Transition.delta(1),
         stack_pointer=Transition.same(),
         dynamic_gas_cost=instruction.select(is_warm, FQ(0), FQ(EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS)),
diff --git a/tests/evm/test_balance.py b/tests/evm/test_balance.py
index 90d90b288..ea7b5ea95 100644
--- a/tests/evm/test_balance.py
+++ b/tests/evm/test_balance.py
@@ -12,13 +12,11 @@
     verify_steps,
 )
 from zkevm_specs.util import (
-    EMPTY_CODE_HASH,
     EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS,
     GAS_COST_WARM_ACCESS,
     RLC,
     U160,
     U256,
-    keccak256,
     rand_address,
     rand_bytes,
     rand_fq,
@@ -27,37 +25,34 @@
 )
 
 TESTING_DATA = [
-    (0x30000, 0, 0, bytes(), True, True),
-    (0x30000, 0, 0, bytes(), False, True),
-    (0x30000, 200, 1, bytes(), True, True),
-    (0x30000, 200, 0, bytes([10, 10]), False, True),
+    (0x30000, 0, True, True, True),
+    (0x30000, 0, True, False, True),
+    (0x30000, 200, False, True, True),
+    (0x30000, 200, False, False, True),
     (
         rand_address(),
         rand_word(),
-        rand_word(),
-        rand_bytes(100),
+        rand_range(2) == 0,
         rand_range(2) == 0,
         True,  # persistent call
     ),
     (
         rand_address(),
         rand_word(),
-        rand_word(),
-        rand_bytes(100),
+        rand_range(2) == 0,
         rand_range(2) == 0,
         False,  # reverted call
     ),
 ]
 
 
-@pytest.mark.parametrize("address, balance, nonce, code, is_warm, is_persistent", TESTING_DATA)
+@pytest.mark.parametrize("address, balance, is_non_existing, is_warm, is_persistent", TESTING_DATA)
 def test_balance(
-    address: U160, balance: U256, nonce: U256, code: bytes, is_warm: bool, is_persistent: bool
+    address: U160, balance: U256, is_non_existing: bool, is_warm: bool, is_persistent: bool
 ):
     randomness = rand_fq()
 
-    code_hash = int.from_bytes(keccak256(code), "big")
-    result = 0 if (balance == 0 and nonce == 0 and code_hash == EMPTY_CODE_HASH) else balance
+    result = 0 if is_non_existing else balance
 
     tx_id = 1
     call_id = 1
@@ -80,9 +75,8 @@ def test_balance(
             is_warm,
             rw_counter_of_reversion=rw_counter_end_of_reversion - reversible_write_counter,
         )
-        .account_read(address, AccountFieldTag.Nonce, RLC(nonce, randomness))
+        .account_read(address, AccountFieldTag.NonExisting, RLC(is_non_existing, randomness))
         .account_read(address, AccountFieldTag.Balance, RLC(balance, randomness))
-        .account_read(address, AccountFieldTag.CodeHash, RLC(code_hash, randomness))
         .stack_write(call_id, 1023, RLC(result, randomness))
         .rws
     )
@@ -113,7 +107,7 @@ def test_balance(
             ),
             StepState(
                 execution_state=ExecutionState.STOP if is_persistent else ExecutionState.REVERT,
-                rw_counter=10,
+                rw_counter=9,
                 call_id=1,
                 is_root=True,
                 is_create=False,

From 6bf2c5c7a9ce7c5ddc761fe01edadf772b5b9b08 Mon Sep 17 00:00:00 2001
From: Steven Gu <asongala@163.com>
Date: Wed, 5 Oct 2022 20:15:15 +0800
Subject: [PATCH 08/14] Revert `extcodehash` fix.

---
 src/zkevm_specs/evm/execution/extcodehash.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/src/zkevm_specs/evm/execution/extcodehash.py b/src/zkevm_specs/evm/execution/extcodehash.py
index 0086eb8c7..18ea4198d 100644
--- a/src/zkevm_specs/evm/execution/extcodehash.py
+++ b/src/zkevm_specs/evm/execution/extcodehash.py
@@ -17,9 +17,6 @@ def extcodehash(instruction: Instruction):
     balance = instruction.account_read(address, AccountFieldTag.Balance)
     code_hash = instruction.account_read(address, AccountFieldTag.CodeHash)
 
-    # Get flag of the account existance.
-    # TODO: Need to update after the below issue is fixed -
-    # https://github.com/privacy-scaling-explorations/zkevm-specs/issues/249.
     is_empty = (
         instruction.is_zero(nonce)
         * instruction.is_zero(balance)

From 21aa597e4bee0e1c1c98a93042bd81fbbb96d60b Mon Sep 17 00:00:00 2001
From: Steven Gu <asongala@163.com>
Date: Sun, 30 Oct 2022 22:13:19 +0800
Subject: [PATCH 09/14] Fix to pass witness value `exists` as `aux_data`, and
 only do `balance` lookup if `exists == 1`, and `non-existing` lookup
 otherwise.

---
 src/zkevm_specs/evm/execution/balance.py | 12 +++++---
 tests/evm/test_balance.py                | 35 +++++++++++++-----------
 2 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/src/zkevm_specs/evm/execution/balance.py b/src/zkevm_specs/evm/execution/balance.py
index d7944f5b3..aec1b28b7 100644
--- a/src/zkevm_specs/evm/execution/balance.py
+++ b/src/zkevm_specs/evm/execution/balance.py
@@ -1,4 +1,4 @@
-from ...util import EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS, FQ, N_BYTES_ACCOUNT_ADDRESS
+from ...util import EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS, FQ, N_BYTES_ACCOUNT_ADDRESS, RLC
 from ..instruction import Instruction, Transition
 from ..opcode import Opcode
 from ..table import AccountFieldTag, CallContextFieldTag
@@ -13,11 +13,15 @@ def balance(instruction: Instruction):
     tx_id = instruction.call_context_lookup(CallContextFieldTag.TxId)
     is_warm = instruction.add_account_to_access_list(tx_id, address, instruction.reversion_info())
 
-    is_non_existing = instruction.account_read(address, AccountFieldTag.NonExisting)
-    balance = instruction.account_read(address, AccountFieldTag.Balance)
+    # Load account `exists` value from auxilary witness data.
+    exists = instruction.curr.aux_data
 
+    if exists == 0:  # 1 - exists == 1
+        instruction.account_read(address, AccountFieldTag.NonExisting)
+
+    balance = instruction.account_read(address, AccountFieldTag.Balance) if exists == 1 else RLC(0)
     instruction.constrain_equal(
-        instruction.select(is_non_existing.expr(), FQ(0), balance.expr()),
+        instruction.select(exists, balance.expr(), FQ(0)),
         instruction.stack_push(),
     )
 
diff --git a/tests/evm/test_balance.py b/tests/evm/test_balance.py
index ea7b5ea95..7ceaea4c7 100644
--- a/tests/evm/test_balance.py
+++ b/tests/evm/test_balance.py
@@ -25,34 +25,32 @@
 )
 
 TESTING_DATA = [
-    (0x30000, 0, True, True, True),
-    (0x30000, 0, True, False, True),
-    (0x30000, 200, False, True, True),
-    (0x30000, 200, False, False, True),
+    (0x30000, 0, 0, True, True),
+    (0x30000, 0, 0, False, True),
+    (0x30000, 200, 1, True, True),
+    (0x30000, 200, 1, False, True),
     (
         rand_address(),
         rand_word(),
-        rand_range(2) == 0,
+        rand_range(2),
         rand_range(2) == 0,
         True,  # persistent call
     ),
     (
         rand_address(),
         rand_word(),
-        rand_range(2) == 0,
+        rand_range(2),
         rand_range(2) == 0,
         False,  # reverted call
     ),
 ]
 
 
-@pytest.mark.parametrize("address, balance, is_non_existing, is_warm, is_persistent", TESTING_DATA)
-def test_balance(
-    address: U160, balance: U256, is_non_existing: bool, is_warm: bool, is_persistent: bool
-):
+@pytest.mark.parametrize("address, balance, exists, is_warm, is_persistent", TESTING_DATA)
+def test_balance(address: U160, balance: U256, exists: int, is_warm: bool, is_persistent: bool):
     randomness = rand_fq()
 
-    result = 0 if is_non_existing else balance
+    result = balance if exists == 1 else 0
 
     tx_id = 1
     call_id = 1
@@ -60,7 +58,7 @@ def test_balance(
     rw_counter_end_of_reversion = 0 if is_persistent else 9
     reversible_write_counter = 0
 
-    rw_table = set(
+    rw_dictionary = (
         RWDictionary(1)
         .stack_read(call_id, 1023, RLC(address, randomness))
         .call_context_read(tx_id, CallContextFieldTag.TxId, tx_id)
@@ -75,11 +73,15 @@ def test_balance(
             is_warm,
             rw_counter_of_reversion=rw_counter_end_of_reversion - reversible_write_counter,
         )
-        .account_read(address, AccountFieldTag.NonExisting, RLC(is_non_existing, randomness))
-        .account_read(address, AccountFieldTag.Balance, RLC(balance, randomness))
-        .stack_write(call_id, 1023, RLC(result, randomness))
-        .rws
     )
+    if exists == 1:
+        rw_dictionary.account_read(address, AccountFieldTag.Balance, RLC(balance, randomness))
+    else:
+        rw_dictionary.account_read(
+            address, AccountFieldTag.NonExisting, RLC(1 - exists, randomness)
+        )
+
+    rw_table = set(rw_dictionary.stack_write(call_id, 1023, RLC(result, randomness)).rws)
 
     bytecode = Bytecode().balance()
     tables = Tables(
@@ -104,6 +106,7 @@ def test_balance(
                 program_counter=0,
                 stack_pointer=1023,
                 gas_left=GAS_COST_WARM_ACCESS + (not is_warm) * EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS,
+                aux_data=exists,
             ),
             StepState(
                 execution_state=ExecutionState.STOP if is_persistent else ExecutionState.REVERT,

From a84d4ae8137c03b55f7d5266dea02b159ce6b422 Mon Sep 17 00:00:00 2001
From: Steven Gu <asongala@163.com>
Date: Sun, 30 Oct 2022 22:19:29 +0800
Subject: [PATCH 10/14] Fix reversible_write_counter to `10` with a comment.

---
 tests/evm/test_balance.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/evm/test_balance.py b/tests/evm/test_balance.py
index 7ceaea4c7..45e6b9885 100644
--- a/tests/evm/test_balance.py
+++ b/tests/evm/test_balance.py
@@ -55,7 +55,8 @@ def test_balance(address: U160, balance: U256, exists: int, is_warm: bool, is_pe
     tx_id = 1
     call_id = 1
 
-    rw_counter_end_of_reversion = 0 if is_persistent else 9
+    # 9 + 1 reversible operation (the account access list write)
+    rw_counter_end_of_reversion = 0 if is_persistent else 10
     reversible_write_counter = 0
 
     rw_dictionary = (

From eeb57ba812e5f0d875da3dffa03beaa1dcd224be Mon Sep 17 00:00:00 2001
From: Steven Gu <asongala@163.com>
Date: Sun, 30 Oct 2022 22:35:17 +0800
Subject: [PATCH 11/14] Update Markdown doc.

---
 specs/opcode/31BALANCE.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/specs/opcode/31BALANCE.md b/specs/opcode/31BALANCE.md
index cfc3751d3..eae867a19 100644
--- a/specs/opcode/31BALANCE.md
+++ b/specs/opcode/31BALANCE.md
@@ -30,13 +30,13 @@ stack instead.
    - gas:
      - the accessed `address` is warm: GAS_COST_WARM_ACCESS
      - the accessed `address` is cold: GAS_COST_ACCOUNT_COLD_ACCESS
-3. Lookups: 8 busmapping lookups
+3. Lookups: 7 busmapping lookups
    - `address` is at top of the stack.
    - 3 reads from call context for `tx_id`, `rw_counter_end_of_reversion`, and
      `is_persistent`.
    - `address` is added to the transaction access list if not already present.
-   - `is_non_existing` is read to check account existence.
-   - `balance` is read from account. Set to 0 if it doesn't exist.
+   - If witness value `exists == 1`, lookup account `balance`. Otherwise lookup
+     account non-existing proof.
    - The BALANCE result is at the new top of the stack.
 4. Additional Constraints
    - value `is_warm` matches the gas cost for this opcode.

From c3a6026b22ee33f1d8e23af18f63ecf90d553c5e Mon Sep 17 00:00:00 2001
From: Steven <asongala@163.com>
Date: Wed, 2 Nov 2022 22:43:10 +0800
Subject: [PATCH 12/14] Update tests/evm/test_balance.py

Co-authored-by: Eduard S. <eduardsanou@posteo.net>
---
 tests/evm/test_balance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/evm/test_balance.py b/tests/evm/test_balance.py
index 45e6b9885..291c611e6 100644
--- a/tests/evm/test_balance.py
+++ b/tests/evm/test_balance.py
@@ -111,7 +111,7 @@ def test_balance(address: U160, balance: U256, exists: int, is_warm: bool, is_pe
             ),
             StepState(
                 execution_state=ExecutionState.STOP if is_persistent else ExecutionState.REVERT,
-                rw_counter=9,
+                rw_counter=8,
                 call_id=1,
                 is_root=True,
                 is_create=False,

From 2b8f73d09c75fc8bfdb19a35fd660ed7348fc00a Mon Sep 17 00:00:00 2001
From: Steven <asongala@163.com>
Date: Wed, 2 Nov 2022 22:43:22 +0800
Subject: [PATCH 13/14] Update src/zkevm_specs/evm/execution/balance.py

Co-authored-by: Eduard S. <eduardsanou@posteo.net>
---
 src/zkevm_specs/evm/execution/balance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/zkevm_specs/evm/execution/balance.py b/src/zkevm_specs/evm/execution/balance.py
index aec1b28b7..6cfcb75e1 100644
--- a/src/zkevm_specs/evm/execution/balance.py
+++ b/src/zkevm_specs/evm/execution/balance.py
@@ -27,7 +27,7 @@ def balance(instruction: Instruction):
 
     instruction.step_state_transition_in_same_context(
         opcode,
-        rw_counter=Transition.delta(8),
+        rw_counter=Transition.delta(7),
         program_counter=Transition.delta(1),
         stack_pointer=Transition.same(),
         dynamic_gas_cost=instruction.select(is_warm, FQ(0), FQ(EXTRA_GAS_COST_ACCOUNT_COLD_ACCESS)),

From b8be6277d82f8e92acc666617d2441a4184f3128 Mon Sep 17 00:00:00 2001
From: Steven <asongala@163.com>
Date: Wed, 2 Nov 2022 22:43:33 +0800
Subject: [PATCH 14/14] Update specs/opcode/31BALANCE.md

Co-authored-by: Eduard S. <eduardsanou@posteo.net>
---
 specs/opcode/31BALANCE.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/specs/opcode/31BALANCE.md b/specs/opcode/31BALANCE.md
index eae867a19..05311532a 100644
--- a/specs/opcode/31BALANCE.md
+++ b/specs/opcode/31BALANCE.md
@@ -23,7 +23,7 @@ stack instead.
 
 1. opId = 0x31
 2. State transition:
-   - gc + 8 (1 stack read, 1 stack write, 3 call context reads, 2 account reads,
+   - gc + 7 (1 stack read, 1 stack write, 3 call context reads, 1 account reads,
      1 transaction access list write)
    - stack_pointer + 0 (one pop and one push)
    - pc + 1