perf(evm): reuse CFG memory proofs for KECCAK expansion

src/compiler/evm_frontend/evm_memory_analysis.h

@@ -876,6 +876,8 @@ class MemoryGuaranteedMinBytesAnalysis {
     switch (Op.Kind) {
     case MemoryOpKind::MLoad:
       return Op.Reads.size() == 1 && getIntervalEnd(Op.Reads[0], End);
+    case MemoryOpKind::Keccak:
+      return Op.Reads.size() == 1 && getIntervalEnd(Op.Reads[0], End);
     case MemoryOpKind::MStore:
     case MemoryOpKind::MStore8:
       return Op.Writes.size() == 1 && getIntervalEnd(Op.Writes[0], End);

src/compiler/evm_frontend/evm_mir_compiler.cpp

@@ -5861,7 +5861,12 @@ EVMMirBuilder::handleKeccak256(Operand OffsetComponents,
     syncGasToMemory();
 #endif
   } else {
-    preExpandMemoryRange(OffsetComponents, LengthComponents);
+    const bool ExpansionCovered = OffsetWasConstU64 && LengthWasConstU64 &&
+                                  tryUseGuaranteedMinBytesExpansionElision(
+                                      true, ConstOffset, ConstLength);
+    if (!ExpansionCovered) {
+      preExpandMemoryRange(OffsetComponents, LengthComponents);
+    }
     MInstruction *Length = extractKnownU64LowOperand(LengthComponents);
     chargeKeccakWordGasIR(Length);
   }
@@ -8639,6 +8644,7 @@ void EVMMirBuilder::noteMemoryOpcodeInBlock(evmc_opcode Opcode, uint64_t PC) {
 }
 
 void EVMMirBuilder::noteHelperOpcodeInBlock(evmc_opcode Opcode, uint64_t PC) {
+  CurrentMemoryOpPC = PC;
   if (!CurBlockMemStats.Active) {
     return;
   }

src/tests/evm_differential_tests.cpp

@@ -73,10 +73,11 @@ struct EVMExecutionResult {
 // config can fall back to the interpreter for a single call, which would make
 // a differential test vacuous; DisableMultipassMultithread compiles before the
 // run so the JIT side is non-vacuous.
-EVMExecutionResult
-runEvmBytecode(const std::string &Label, const std::vector<uint8_t> &Bytecode,
-               common::RunMode Mode, const std::vector<uint8_t> &CallData = {},
-               uint32_t MessageFlags = 0u, bool EnableGasMetering = false) {
+EVMExecutionResult runEvmBytecode(
+    const std::string &Label, const std::vector<uint8_t> &Bytecode,
+    common::RunMode Mode, const std::vector<uint8_t> &CallData = {},
+    uint32_t MessageFlags = 0u, bool EnableGasMetering = false,
+    uint64_t GasLimit = 0xFFFF'FFFF'FFFF - zen::evm::BASIC_EXECUTION_COST) {
   EVMExecutionResult Empty;
 
   RuntimeConfig Config;
@@ -106,7 +107,6 @@ runEvmBytecode(const std::string &Label, const std::vector<uint8_t> &Bytecode,
     return Empty;
   }
 
-  const uint64_t GasLimit = 0xFFFF'FFFF'FFFF - zen::evm::BASIC_EXECUTION_COST;
   auto InstRet = Iso->createEVMInstance(*Mod, GasLimit);
   if (!InstRet) {
     ADD_FAILURE() << "instance create failed: " << Label;
@@ -262,6 +262,68 @@ std::string fixtureTestName(const testing::TestParamInfo<std::string> &Info) {
 
 } // namespace
 
+TEST(EVMKeccakMemoryProofDifferential,
+     CrossBlockProofReusePreservesHashAndGas) {
+  const std::vector<uint8_t> Bytecode = {
+      0x60, 0x01, // PC0: PUSH1 1
+      0x60, 0x80, // PC2: PUSH1 0x80
+      0x52,       // PC4: MSTORE, proves memory through 0xa0
+      0x60, 0x08, // PC5: PUSH1 successor
+      0x56,       // PC7: JUMP
+      0x5b,       // PC8: JUMPDEST
+      0x60, 0x20, // PC9: PUSH1 hash length
+      0x60, 0x80, // PC11: PUSH1 hash offset
+      0x20,       // PC13: KECCAK256
+      0x5f,       // PC14: PUSH0 result offset
+      0x52,       // PC15: MSTORE hash result
+      0x60, 0x20, // PC16: PUSH1 return length
+      0x5f,       // PC18: PUSH0 return offset
+      0xf3,       // PC19: RETURN
+  };
+
+  const auto Output =
+      expectInterpMatchesMultipassWithGas("keccak_cfg_proof", Bytecode, {});
+  EXPECT_EQ(Output,
+            "B10E2D527612073B26EECDFD717E6A320CF44B4AFAC2B0732D9FCBE2B7FA0CF6");
+}
+
+TEST(EVMKeccakMemoryProofDifferential,
+     WordGasOutOfGasBoundaryMatchesInterpreter) {
+  // KECCAK256 is the final operation, so one gas below the successful
+  // execution budget exercises its final dynamic word-gas charge.
+  const std::vector<uint8_t> Bytecode = {
+      0x60, 0x01, // PUSH1 1
+      0x60, 0x80, // PUSH1 0x80
+      0x52,       // MSTORE, expands memory before KECCAK256
+      0x60, 0x20, // PUSH1 hash length
+      0x60, 0x80, // PUSH1 hash offset
+      0x20,       // KECCAK256
+  };
+  constexpr uint64_t ProbeGasLimit = 1'000'000;
+  const auto Reference =
+      runEvmBytecode("keccak_word_gas_reference", Bytecode,
+                     common::RunMode::InterpMode, {}, 0u, true, ProbeGasLimit);
+  ASSERT_EQ(Reference.Status, EVMC_SUCCESS);
+  ASSERT_GE(Reference.GasLeft, 0);
+  const auto RequiredGas =
+      ProbeGasLimit - static_cast<uint64_t>(Reference.GasLeft);
+  ASSERT_GT(RequiredGas, uint64_t{1});
+
+  const auto Interp = runEvmBytecode("keccak_word_gas_oog_interp", Bytecode,
+                                     common::RunMode::InterpMode, {}, 0u, true,
+                                     RequiredGas - 1);
+  const auto Multi = runEvmBytecode("keccak_word_gas_oog_multipass", Bytecode,
+                                    common::RunMode::MultipassMode, {}, 0u,
+                                    true, RequiredGas - 1);
+#ifdef ZEN_ENABLE_JIT
+  EXPECT_TRUE(Multi.JITCompiled);
+#endif
+  EXPECT_EQ(Interp.Status, EVMC_OUT_OF_GAS);
+  EXPECT_EQ(Multi.Status, Interp.Status);
+  EXPECT_EQ(Multi.GasLeft, Interp.GasLeft);
+  EXPECT_EQ(Multi.OutputHex, Interp.OutputHex);
+}
+
 // ===========================================================================
 // Fixture-based differential suite.
 //

src/tests/evm_jit_frontend_tests.cpp

@@ -3,6 +3,7 @@
 
 #include "action/evm_bytecode_visitor.h"
 #include "compiler/evm_frontend/evm_analyzer.h"
+#include "compiler/evm_frontend/evm_imported.h"
 #include "compiler/evm_frontend/evm_memory_analysis.h"
 #include "compiler/evm_frontend/evm_memory_facts.h"
 #include "compiler/evm_frontend/evm_memory_grouping.h"
@@ -396,6 +397,130 @@ TEST(EVMMemoryGuaranteedMinBytesAnalysisTest, LearnsConstMCopyUnionEnd) {
   EXPECT_EQ(Guaranteed.getGuaranteedMinBytesBeforeOp(Facts.Ops[1].Id), 0x40u);
 }
 
+TEST(EVMMemoryGuaranteedMinBytesAnalysisTest, LearnsConstKeccakReadEnd) {
+  const std::vector<uint8_t> Bytecode = {OP_PUSH1, 0x40,         OP_PUSH1,
+                                         0x20,     OP_KECCAK256, OP_POP,
+                                         OP_PUSH1, 0x20,         OP_MLOAD};
+
+  COMPILER::MemoryFacts Facts = collectMemoryFacts(Bytecode);
+  COMPILER::MemoryGuaranteedMinBytesAnalysis Guaranteed(Facts);
+
+  ASSERT_EQ(Facts.Ops.size(), 2u);
+  EXPECT_EQ(Guaranteed.getGuaranteedMinBytesBeforeOp(Facts.Ops[1].Id), 0x60u);
+}
+
+TEST(EVMMemoryGuaranteedMinBytesAnalysisTest,
+     PropagatesConstKeccakGuaranteeToSuccessor) {
+  const std::vector<uint8_t> Bytecode = {
+      OP_PUSH1, 0x40,        OP_PUSH1, 0x20, OP_KECCAK256,
+      OP_POP,   OP_JUMPDEST, OP_PUSH1, 0x20, OP_MLOAD};
+
+  COMPILER::MemoryFacts Facts = collectAnalyzerMemoryFacts(Bytecode);
+  COMPILER::MemoryGuaranteedMinBytesAnalysis Guaranteed(Facts);
+
+  ASSERT_EQ(Facts.Ops.size(), 2u);
+  EXPECT_EQ(Guaranteed.getGuaranteedMinBytesAtEntry(6), 0x60u);
+}
+
+TEST(EVMMemoryGuaranteedMinBytesAnalysisTest, IgnoresZeroLengthKeccak) {
+  const std::vector<uint8_t> Bytecode = {OP_PUSH1, 0x00,         OP_PUSH1,
+                                         0x80,     OP_KECCAK256, OP_POP,
+                                         OP_PUSH1, 0x00,         OP_MLOAD};
+
+  COMPILER::MemoryFacts Facts = collectMemoryFacts(Bytecode);
+  COMPILER::MemoryGuaranteedMinBytesAnalysis Guaranteed(Facts);
+
+  ASSERT_EQ(Facts.Ops.size(), 2u);
+  EXPECT_EQ(Guaranteed.getGuaranteedMinBytesBeforeOp(Facts.Ops[1].Id), 0u);
+}
+
+#ifdef ZEN_ENABLE_EVM_MEMORY_PLAN_FRAMEWORK
+std::optional<size_t>
+countKeccakExpandMemoryCalls(const std::vector<uint8_t> &Bytecode) {
+  COMPILER::EVMFrontendContext Ctx;
+  Ctx.setRevision(EVMC_CANCUN);
+  Ctx.setBytecode(reinterpret_cast<const zen::common::Byte *>(Bytecode.data()),
+                  Bytecode.size());
+
+  COMPILER::MModule Mod(Ctx);
+  std::array<COMPILER::MType *, 1> ParamTypes = {
+      COMPILER::MPointerType::create(Ctx, Ctx.VoidType)};
+  COMPILER::MFunctionType *FuncType = COMPILER::MFunctionType::create(
+      Ctx, Ctx.VoidType, llvm::ArrayRef<COMPILER::MType *>(ParamTypes));
+  Mod.addFuncType(FuncType);
+
+  COMPILER::MFunction Func(Ctx, 0);
+  Func.setFunctionType(FuncType);
+  EVMMirBuilder Builder(Ctx, Func);
+  if (!Builder.compile(&Ctx)) {
+    return std::nullopt;
+  }
+
+  std::string Mir;
+  llvm::raw_string_ostream OS(Mir);
+  Func.print(OS);
+  OS.flush();
+
+  const auto &RuntimeFunctions = COMPILER::getRuntimeFunctionTable();
+  const auto Address =
+      COMPILER::getFunctionAddress(RuntimeFunctions.ExpandMemoryNoGas);
+  const std::string Needle =
+      "target = const.i64 " + std::to_string(Address) + ", ";
+  size_t Count = 0;
+  for (size_t Pos = Mir.find(Needle); Pos != std::string::npos;
+       Pos = Mir.find(Needle, Pos + Needle.size())) {
+    ++Count;
+  }
+  return Count;
+}
+
+TEST(EVMMirBuilderMemoryProofTest,
+     ReusesCrossBlockGuaranteedSizeForKeccakConsumer) {
+  const std::vector<uint8_t> ProducerOnly = {
+      OP_PUSH1, 0x01, OP_PUSH1, 0x80,        OP_MSTORE,
+      OP_PUSH1, 0x08, OP_JUMP,  OP_JUMPDEST, OP_STOP};
+  const std::vector<uint8_t> CoveredKeccak = {
+      OP_PUSH1, 0x01,         OP_PUSH1,    0x80,     OP_MSTORE, OP_PUSH1,
+      0x08,     OP_JUMP,      OP_JUMPDEST, OP_PUSH1, 0x20,      OP_PUSH1,
+      0x80,     OP_KECCAK256, OP_POP,      OP_STOP};
+  const std::vector<uint8_t> UncoveredKeccak = {
+      OP_PUSH1, 0x01,         OP_PUSH1,    0x80,     OP_MSTORE, OP_PUSH1,
+      0x08,     OP_JUMP,      OP_JUMPDEST, OP_PUSH1, 0x20,      OP_PUSH1,
+      0xa0,     OP_KECCAK256, OP_POP,      OP_STOP};
+
+  const std::optional<size_t> ProducerExpansionCalls =
+      countKeccakExpandMemoryCalls(ProducerOnly);
+  const std::optional<size_t> CoveredExpansionCalls =
+      countKeccakExpandMemoryCalls(CoveredKeccak);
+  const std::optional<size_t> UncoveredExpansionCalls =
+      countKeccakExpandMemoryCalls(UncoveredKeccak);
+
+  ASSERT_TRUE(ProducerExpansionCalls.has_value());
+  ASSERT_TRUE(CoveredExpansionCalls.has_value());
+  ASSERT_TRUE(UncoveredExpansionCalls.has_value());
+  EXPECT_EQ(*CoveredExpansionCalls, *ProducerExpansionCalls);
+  EXPECT_EQ(*UncoveredExpansionCalls, *ProducerExpansionCalls + 1);
+}
+
+TEST(EVMMirBuilderMemoryProofTest, ReusesPriorKeccakProofAtExactOpcodePC) {
+  const std::vector<uint8_t> OneKeccak = {OP_PUSH1,     0x20,   OP_PUSH1, 0x80,
+                                          OP_KECCAK256, OP_POP, OP_STOP};
+  const std::vector<uint8_t> TwoKeccaks = {
+      OP_PUSH1,     0x20,     OP_PUSH1, 0x80,     OP_KECCAK256,
+      OP_POP,       OP_PUSH1, 0x20,     OP_PUSH1, 0x80,
+      OP_KECCAK256, OP_POP,   OP_STOP};
+
+  const std::optional<size_t> OneExpansionCalls =
+      countKeccakExpandMemoryCalls(OneKeccak);
+  const std::optional<size_t> TwoExpansionCalls =
+      countKeccakExpandMemoryCalls(TwoKeccaks);
+
+  ASSERT_TRUE(OneExpansionCalls.has_value());
+  ASSERT_TRUE(TwoExpansionCalls.has_value());
+  EXPECT_EQ(*TwoExpansionCalls, *OneExpansionCalls);
+}
+#endif
+
 TEST(EVMMemoryGuaranteedMinBytesAnalysisTest, IgnoresZeroLengthMCopy) {
   const std::vector<uint8_t> Bytecode = {OP_PUSH1, 0x00,    OP_PUSH1, 0x80,
                                          OP_PUSH1, 0xa0,    OP_MCOPY, OP_PUSH1,