From 01c2608e7faf942cc288e0b5681b793fa64f539e Mon Sep 17 00:00:00 2001
From: Mason Liang <mason@scroll.tech>
Date: Fri, 4 Feb 2022 17:30:03 -0500
Subject: [PATCH] Add spec and test for selfbalance opcode

---
 specs/opcode/47SELFBALANCE.md                | 32 ++++++++
 src/zkevm_specs/evm/execution/__init__.py    |  2 +
 src/zkevm_specs/evm/execution/selfbalance.py | 19 +++++
 tests/evm/test_selfbalance.py                | 81 ++++++++++++++++++++
 4 files changed, 134 insertions(+)
 create mode 100644 specs/opcode/47SELFBALANCE.md
 create mode 100644 src/zkevm_specs/evm/execution/selfbalance.py
 create mode 100644 tests/evm/test_selfbalance.py

diff --git a/specs/opcode/47SELFBALANCE.md b/specs/opcode/47SELFBALANCE.md
new file mode 100644
index 000000000..8df0c7b4e
--- /dev/null
+++ b/specs/opcode/47SELFBALANCE.md
@@ -0,0 +1,32 @@
+# Selfbalance op code
+
+## Procedure
+
+The `SELFBALANCE` opcode pushes the balance (32 bytes of data) of the currently executing address onto the stack.
+
+## Circuit behaviour
+
+1. Construct call context table in rw table
+2. Do busmapping lookup for stack write operation
+
+## Constraints
+
+1. opId = 0x47
+2. State transition:
+   - gc + 3 (1 stack write, 1 account balance read, and 1 callee address read)
+   - stack_pointer - 1
+   - pc + 1
+   - gas + 5
+3. Lookups: 3
+   - `callee_address` is the callee address of the call context
+   - `self_balance` is the balance of `callee_address`
+   - `self_balance` is on the top of stack
+
+## Exceptions
+
+1. stack overflow: stack is full, stack pointer = 0
+2. out of gas: remaining gas is not enough
+
+## Code
+
+Please refer to `src/zkevm_specs/evm/execution/selfbalance.py`.
diff --git a/src/zkevm_specs/evm/execution/__init__.py b/src/zkevm_specs/evm/execution/__init__.py
index 7c158b3df..a4bf3645c 100644
--- a/src/zkevm_specs/evm/execution/__init__.py
+++ b/src/zkevm_specs/evm/execution/__init__.py
@@ -18,6 +18,7 @@
 from .jumpi import *
 from .push import *
 from .slt_sgt import *
+from .selfbalance import *
 
 
 EXECUTION_STATE_IMPL: Dict[ExecutionState, Callable] = {
@@ -35,4 +36,5 @@
     ExecutionState.PUSH: push,
     ExecutionState.SCMP: scmp,
     ExecutionState.GAS: gas,
+    ExecutionState.SELFBALANCE: selfbalance,
 }
diff --git a/src/zkevm_specs/evm/execution/selfbalance.py b/src/zkevm_specs/evm/execution/selfbalance.py
new file mode 100644
index 000000000..b69448757
--- /dev/null
+++ b/src/zkevm_specs/evm/execution/selfbalance.py
@@ -0,0 +1,19 @@
+from ..instruction import Instruction, Transition
+from ..table import AccountFieldTag, CallContextFieldTag
+from ..opcode import Opcode
+
+
+def selfbalance(instruction: Instruction):
+    opcode = instruction.opcode_lookup(True)
+    instruction.constrain_equal(opcode, Opcode.SELFBALANCE)
+
+    callee_address = instruction.call_context_lookup(CallContextFieldTag.CalleeAddress)
+    balance = instruction.account_read(callee_address, AccountFieldTag.Balance)
+    instruction.constrain_equal(instruction.stack_push(), balance)
+
+    instruction.step_state_transition_in_same_context(
+        opcode,
+        rw_counter=Transition.delta(3),
+        program_counter=Transition.delta(1),
+        stack_pointer=Transition.delta(-1),
+    )
diff --git a/tests/evm/test_selfbalance.py b/tests/evm/test_selfbalance.py
new file mode 100644
index 000000000..2894dd4b8
--- /dev/null
+++ b/tests/evm/test_selfbalance.py
@@ -0,0 +1,81 @@
+import pytest
+
+from zkevm_specs.evm import (
+    ExecutionState,
+    StepState,
+    verify_steps,
+    Tables,
+    RWTableTag,
+    RW,
+    Block,
+    Bytecode,
+    CallContextFieldTag,
+    AccountFieldTag,
+)
+from zkevm_specs.util import rand_address, rand_word, rand_fp, RLC, U256, U160
+
+TESTING_DATA = [(0, 0), (0, 10), (rand_address(), rand_word())]
+
+
+@pytest.mark.parametrize("callee_address, balance", TESTING_DATA)
+def test_selfbalance(callee_address: U160, balance: U256):
+    randomness = rand_fp()
+
+    bytecode = Bytecode().selfbalance()
+    bytecode_hash = RLC(bytecode.hash(), randomness)
+
+    rlc_balance = RLC(balance, randomness)
+
+    tables = Tables(
+        block_table=Block(),
+        tx_table=set(),
+        bytecode_table=set(bytecode.table_assignments(randomness)),
+        rw_table=set(
+            [
+                (9, RW.Read, RWTableTag.CallContext, 1, CallContextFieldTag.CalleeAddress, 0, callee_address, 0, 0, 0),
+                (
+                    10,
+                    RW.Read,
+                    RWTableTag.Account,
+                    callee_address,
+                    AccountFieldTag.Balance,
+                    0,
+                    rlc_balance,
+                    rlc_balance,
+                    0,
+                    0,
+                    0,
+                ),
+                (11, RW.Write, RWTableTag.Stack, 1, 1023, 0, rlc_balance, 0, 0, 0),
+            ]
+        ),
+    )
+
+    verify_steps(
+        randomness=randomness,
+        tables=tables,
+        steps=[
+            StepState(
+                execution_state=ExecutionState.SELFBALANCE,
+                rw_counter=9,
+                call_id=1,
+                is_root=True,
+                is_create=False,
+                code_source=bytecode_hash,
+                program_counter=0,
+                stack_pointer=1024,
+                gas_left=5,
+            ),
+            StepState(
+                execution_state=ExecutionState.STOP,
+                rw_counter=12,
+                call_id=1,
+                is_root=True,
+                is_create=False,
+                code_source=bytecode_hash,
+                program_counter=1,
+                stack_pointer=1023,
+                gas_left=0,
+            ),
+        ],
+    )