refactor(tests): Proposed patch for bloatnet SSTORE tests

tests/benchmark/test_bloatnet.py

@@ -19,64 +19,218 @@
 from ethereum_test_tools.vm.opcode import Opcodes as Op
 
 
-@pytest.mark.valid_from("Prague")
-@pytest.mark.parametrize("final_storage_value", [0x02 << 248, 0x02])
-def test_bloatnet(
-    blockchain_test: BlockchainTestFiller, pre: Alloc, fork: Fork, final_storage_value: int
+@pytest.mark.valid_from("Osaka")
+def test_bloatnet_sstore_cold(
+    blockchain_test: BlockchainTestFiller, pre: Alloc, fork: Fork, gas_benchmark_value: int
 ):
     """
-    A test that calls a contract with many SSTOREs.
-
-    The first block will have many SSTORES that go from 0 -> 1
-    and the 2nd block will have many SSTORES that go from 1 -> 2
+    Benchmark test that maximizes cold SSTORE operations (0 -> 1) by filling
+    a block with multiple transactions, with each one containing a contract
+    that performs a set of cold SSTOREs until the block gas limit is reached.
     """
-    # Get gas costs for the current fork
     gas_costs = fork.gas_costs()
-
-    # this is only used for computing the intinsic gas
-    data = final_storage_value.to_bytes(32, "big").rstrip(b"\x00")
-
-    storage = Storage()
-
-    # Initial gas for PUSH0 + CALLDATALOAD + POP (at the end)
-    totalgas = gas_costs.G_BASE * 2 + gas_costs.G_VERY_LOW
-    totalgas = totalgas + fork.transaction_intrinsic_cost_calculator()(calldata=data)
-    gas_increment = gas_costs.G_VERY_LOW * 2 + gas_costs.G_STORAGE_SET + gas_costs.G_COLD_SLOAD
-    sstore_code = Op.PUSH0 + Op.CALLDATALOAD
-    storage_slot: int = 0
-    while totalgas + gas_increment < Environment().gas_limit:
-        totalgas += gas_increment
-        sstore_code = sstore_code + Op.SSTORE(storage_slot, Op.DUP1)
-        storage[storage_slot] = final_storage_value
-        storage_slot += 1
-
-    sstore_code = sstore_code + Op.POP  # Drop last value on the stack
-
-    sender = pre.fund_eoa()
-    print(sender)
-    contract_address = pre.deploy_contract(
-        code=sstore_code,
-        storage=Storage(),
+    intrinsic_gas_calc = fork.transaction_intrinsic_cost_calculator()
+
+    # TODO: We should maybe not use `None` for tx limit cap as this leads to typing
+    #  issues. Maybe some really high value or the block gas limit?
+    tx_gas_cap = fork.transaction_gas_limit_cap() or gas_benchmark_value
+
+    storage_value = 1
+    calldata = storage_value.to_bytes(32, "big")
+
+    total_sstores = 0
+    total_block_gas_used = 0
+    all_txs = []
+
+    expected_storage_state = {}
+
+    while total_block_gas_used < gas_benchmark_value:
+        remaining_block_gas = gas_benchmark_value - total_block_gas_used
+        tx_gas_limit = min(remaining_block_gas, tx_gas_cap)
+
+        intrinsic_gas_with_data_floor = intrinsic_gas_calc(calldata=calldata)
+        if tx_gas_limit <= intrinsic_gas_with_data_floor:
+            break
+
+        opcode_gas_budget = tx_gas_limit - intrinsic_gas_with_data_floor
+
+        tx_contract_code = Op.PUSH0 + Op.CALLDATALOAD
+        tx_opcode_gas = (
+            gas_costs.G_BASE  # PUSH0
+            + gas_costs.G_VERY_LOW  # CALLDATALOAD
+        )
+        tx_sstores_count = 0
+
+        current_slot = total_sstores
+
+        pop_gas = gas_costs.G_BASE
+
+        while True:
+            sstore_per_op_cost = (
+                gas_costs.G_VERY_LOW * 2  # PUSH + DUP1
+                + gas_costs.G_COLD_SLOAD
+                + gas_costs.G_STORAGE_SET  # SSTORE
+            )
+
+            if tx_opcode_gas + sstore_per_op_cost + pop_gas > opcode_gas_budget:
+                break
+
+            tx_opcode_gas += sstore_per_op_cost
+            tx_contract_code += Op.SSTORE(current_slot, Op.DUP1)
+            tx_sstores_count += 1
+            current_slot += 1
+
+        # If no SSTOREs could be added, we've filled the block
+        if tx_sstores_count == 0:
+            break
+
+        # Add a POP to clean up the stack at the end
+        tx_contract_code += Op.POP
+        tx_opcode_gas += pop_gas
+
+        contract_address = pre.deploy_contract(code=tx_contract_code)
+        tx = Transaction(
+            to=contract_address,
+            gas_limit=tx_gas_limit,
+            data=calldata,
+            sender=pre.fund_eoa(),
+        )
+        all_txs.append(tx)
+
+        actual_intrinsic_consumed = intrinsic_gas_calc(
+            calldata=calldata,
+            # The actual gas consumed uses the standard intrinsic cost
+            # (prior execution), not the floor cost used for validation
+            return_cost_deducted_prior_execution=True,
+        )
+
+        tx_gas_used = actual_intrinsic_consumed + tx_opcode_gas
+        total_block_gas_used += tx_gas_used
+
+        total_sstores += tx_sstores_count
+
+        # update expected storage state for each contract
+        expected_storage_state[contract_address] = Account(
+            storage=Storage(
+                {
+                    HashInt(slot): HashInt(storage_value)
+                    for slot in range(current_slot - tx_sstores_count, current_slot)
+                }
+            )
+        )
+
+    blockchain_test(
+        pre=pre,
+        blocks=[Block(txs=all_txs)],
+        post=expected_storage_state,
+        expected_benchmark_gas_used=total_block_gas_used,
     )
 
-    tx_0_1 = Transaction(
-        to=contract_address,
-        gas_limit=Environment().gas_limit,
-        data=(final_storage_value // 2).to_bytes(32, "big").rstrip(b"\x00"),
-        value=0,
-        sender=sender,
-    )
-    tx_1_2 = Transaction(
-        to=contract_address,
-        gas_limit=Environment().gas_limit,
-        data=final_storage_value.to_bytes(32, "big").rstrip(b"\x00"),
-        value=0,
-        sender=sender,
-    )
 
-    post = {contract_address: Account(storage=storage)}
+@pytest.mark.valid_from("Osaka")
+def test_bloatnet_sstore_warm(
+    blockchain_test: BlockchainTestFiller, pre: Alloc, fork: Fork, gas_benchmark_value: int
+):
+    """
+    Benchmark test that maximizes warm SSTORE operations (1 -> 2).
 
-    blockchain_test(pre=pre, blocks=[Block(txs=[tx_0_1]), Block(txs=[tx_1_2])], post=post)
+    It fills a block with multiple transactions, each containing a contract
+    that performs a set of warm SSTOREs until the block gas limit is reached.
+    """
+    gas_costs = fork.gas_costs()
+    intrinsic_gas_calc = fork.transaction_intrinsic_cost_calculator()
+    tx_gas_cap = fork.transaction_gas_limit_cap() or gas_benchmark_value
+
+    initial_value = 1
+    new_value = 2
+    calldata = new_value.to_bytes(32, "big")
+
+    total_sstores = 0
+    total_block_gas_used = 0
+    all_txs = []
+    expected_storage_state = {}
+
+    while total_block_gas_used < gas_benchmark_value:
+        remaining_block_gas = gas_benchmark_value - total_block_gas_used
+        tx_gas_limit = min(remaining_block_gas, tx_gas_cap)
+
+        intrinsic_gas_with_data_floor = intrinsic_gas_calc(calldata=calldata)
+        if tx_gas_limit <= intrinsic_gas_with_data_floor:
+            break
+
+        opcode_gas_budget = tx_gas_limit - intrinsic_gas_with_data_floor
+
+        tx_contract_code = Op.PUSH0 + Op.CALLDATALOAD
+        tx_opcode_gas = (
+            gas_costs.G_BASE  # PUSH0
+            + gas_costs.G_VERY_LOW  # CALLDATALOAD
+        )
+        tx_sstores_count = 0
+
+        current_slot = total_sstores
+        pop_gas = gas_costs.G_BASE
+
+        warm_sstore_cost = gas_costs.G_COLD_SLOAD + gas_costs.G_STORAGE_RESET
+        while True:
+            sstore_per_op_cost = (
+                gas_costs.G_VERY_LOW * 2  # PUSH + DUP1
+                + warm_sstore_cost  # SSTORE
+            )
+
+            if tx_opcode_gas + sstore_per_op_cost + pop_gas > opcode_gas_budget:
+                break
+
+            tx_opcode_gas += sstore_per_op_cost
+            tx_contract_code += Op.SSTORE(current_slot, Op.DUP1)
+            tx_sstores_count += 1
+            current_slot += 1
+
+        if tx_sstores_count == 0:
+            break
+
+        tx_contract_code += Op.POP
+        tx_opcode_gas += pop_gas
+
+        # Pre-fill storage with initial values
+        initial_storage = {
+            slot: initial_value for slot in range(total_sstores, total_sstores + tx_sstores_count)
+        }
+
+        contract_address = pre.deploy_contract(
+            code=tx_contract_code,
+            storage=initial_storage,  # type: ignore
+        )
+        tx = Transaction(
+            to=contract_address,
+            gas_limit=tx_gas_limit,
+            data=calldata,
+            sender=pre.fund_eoa(),
+        )
+        all_txs.append(tx)
+
+        actual_intrinsic_consumed = intrinsic_gas_calc(
+            calldata=calldata, return_cost_deducted_prior_execution=True
+        )
+
+        tx_gas_used = actual_intrinsic_consumed + tx_opcode_gas
+        total_block_gas_used += tx_gas_used
+        total_sstores += tx_sstores_count
+
+        expected_storage_state[contract_address] = Account(
+            storage=Storage(
+                {
+                    HashInt(slot): HashInt(new_value)
+                    for slot in range(current_slot - tx_sstores_count, current_slot)
+                }
+            )
+        )
+
+    blockchain_test(
+        pre=pre,
+        blocks=[Block(txs=all_txs)],
+        post=expected_storage_state,
+        expected_benchmark_gas_used=total_block_gas_used,
+    )
 
 
 # Warm reads are very cheap, which means you can really fill a block