[ETCM-912] Implement gas calculation changes in BALANCE

Author: lukasz-golebiewski

src/main/scala/io/iohk/ethereum/vm/OpCode.scala

@@ -386,12 +386,18 @@ case object SHA3 extends OpCode(0x20, 2, 1, _.G_sha3) {
 
 case object ADDRESS extends ConstOp(0x30)(_.env.ownerAddr.toUInt256)
 
-case object BALANCE extends OpCode(0x31, 1, 1, _.G_balance) with ConstGas {
+case object BALANCE extends OpCode(0x31, 1, 1, _.G_balance) with AddrAccessGas with ConstGas {
   protected def exec[W <: WorldStateProxy[W, S], S <: Storage[S]](state: ProgramState[W, S]): ProgramState[W, S] = {
     val (accountAddress, stack1) = state.stack.pop
-    val accountBalance = state.world.getBalance(Address(accountAddress))
+    val addr = Address(accountAddress)
+    val accountBalance = state.world.getBalance(addr)
     val stack2 = stack1.push(accountBalance)
-    state.withStack(stack2).step()
+    state.withStack(stack2).addAccessedAddress(addr).step()
+  }
+
+  protected def address[W <: WorldStateProxy[W, S], S <: Storage[S]](state: ProgramState[W, S]): Address = {
+    val (accountAddress, _) = state.stack.pop
+    Address(accountAddress)
   }
 }
 

src/test/scala/io/iohk/ethereum/vm/OpCodeFunSpec.scala

@@ -156,13 +156,14 @@ class OpCodeFunSpec extends AnyFunSuite with OpCodeTesting with Matchers with Sc
 
   test(BALANCE) { op =>
     forAll(getProgramStateGen(), getUInt256Gen()) { (stateIn, accountBalance) =>
+      val (addr, _) = stateIn.stack.pop
       val stateOut = executeOp(op, stateIn)
       withStackVerification(op, stateIn, stateOut) {
         val (_, stack1) = stateIn.stack.pop
-        stateOut shouldEqual stateIn.withStack(stack1.push(UInt256.Zero)).step()
+        stateOut shouldEqual stateIn.addAccessedAddress(Address(addr)).withStack(stack1.push(UInt256.Zero)).step()
       }
 
-      val (addr, stack1) = stateIn.stack.pop
+      val (_, stack1) = stateIn.stack.pop
 
       val account = Account(balance = accountBalance)
       val world1 = stateIn.world.saveAccount(Address(addr.mod(UInt256(BigInt(2).pow(160)))), account)
@@ -172,7 +173,7 @@ class OpCodeFunSpec extends AnyFunSuite with OpCodeTesting with Matchers with Sc
 
       withStackVerification(op, stateInWithAccount, stateOutWithAccount) {
         val stack2 = stack1.push(accountBalance)
-        stateOutWithAccount shouldEqual stateInWithAccount.withStack(stack2).step()
+        stateOutWithAccount shouldEqual stateInWithAccount.addAccessedAddress(Address(addr)).withStack(stack2).step()
       }
     }
   }

src/test/scala/io/iohk/ethereum/vm/OpCodeGasSpecPostEip2929.scala

@@ -17,7 +17,7 @@ class OpCodeGasSpecPostEip2929 extends AnyFunSuite with OpCodeTesting with Match
 
   import config.feeSchedule._
 
-  test(EXTCODESIZE, EXTCODEHASH) { op =>
+  test(EXTCODESIZE, EXTCODEHASH, BALANCE) { op =>
     val stateGen = getProgramStateGen(
       evmConfig = config,
       stackGen = getStackGen(elems = 1),