generalize ram tag to custom

Author: hero78119

ceno_zkvm/src/instructions/riscv/ecall/keccak.rs

@@ -135,7 +135,7 @@ impl<E: ExtensionField> Instruction<E> for KeccakEcallInstruction<E> {
 
         cb.write_record(
             || "keccak_state_in",
-            RAMType::Undefined,
+            RAMType::Custom,
             keccak_state_record(
                 vm_state.ts.expr(),
                 state_ptr_value.expr_unaligned(),
@@ -145,7 +145,7 @@ impl<E: ExtensionField> Instruction<E> for KeccakEcallInstruction<E> {
         )?;
         cb.read_record(
             || "keccak_state_out",
-            RAMType::Undefined,
+            RAMType::Custom,
             keccak_state_record(
                 vm_state.ts.expr(),
                 state_ptr_value.expr_unaligned(),

ceno_zkvm/src/precompiles/lookup_keccakf.rs

@@ -46,7 +46,7 @@ use crate::{
         utils::{Mask, MaskRepresentation, not8_expr, set_slice_felts_from_u64 as push_instance},
     },
     scheme::utils::gkr_witness,
-    structs::RAMType,
+    structs::{CustomRWTag, RAMType},
 };
 
 pub const ROUNDS: usize = 24;
@@ -103,7 +103,6 @@ pub const AND_LOOKUPS: usize = AND_LOOKUPS_PER_ROUND;
 pub const XOR_LOOKUPS: usize = XOR_LOOKUPS_PER_ROUND;
 pub const RANGE_LOOKUPS: usize = RANGE_LOOKUPS_PER_ROUND;
 pub const STRUCTURAL_WITIN: usize = 6;
-pub const KECCAK_STATE_TAG: u64 = 0x4b4543;
 pub const KECCAK_STATE_PHASE_INPUT: u64 = 0;
 pub const KECCAK_STATE_PHASE_OUTPUT: u64 = 1;
 
@@ -173,7 +172,7 @@ pub fn keccak_state_record<E: ExtensionField>(
     state: impl IntoIterator<Item = Expression<E>>,
 ) -> Vec<Expression<E>> {
     [
-        E::BaseField::from_canonical_u64(KECCAK_STATE_TAG).expr(),
+        CustomRWTag::KeccakState.expr::<E>(),
         cycle,
         state_ptr,
         E::BaseField::from_canonical_u64(phase).expr(),
@@ -557,7 +556,7 @@ impl<E: ExtensionField> ProtocolBuilder<E> for KeccakLayout<E> {
 
         system.read_record(
             || "keccak_state_in",
-            RAMType::Undefined,
+            RAMType::Custom,
             keccak_state_record(
                 layout.cycle.expr(),
                 layout.state_ptr.expr(),
@@ -570,7 +569,7 @@ impl<E: ExtensionField> ProtocolBuilder<E> for KeccakLayout<E> {
         )?;
         system.write_record(
             || "keccak_state_out",
-            RAMType::Undefined,
+            RAMType::Custom,
             keccak_state_record(
                 layout.cycle.expr(),
                 layout.state_ptr.expr(),

ceno_zkvm/src/precompiles/mod.rs

@@ -9,8 +9,8 @@ mod weierstrass;
 
 pub use lookup_keccakf::{
     AND_LOOKUPS, KECCAK_INPUT32_SIZE, KECCAK_OUT_EVAL_SIZE, KECCAK_STATE_PHASE_INPUT,
-    KECCAK_STATE_PHASE_OUTPUT, KECCAK_STATE_TAG, KeccakInstance, KeccakLayout, KeccakParams,
-    KeccakStateInstance, KeccakTrace, KeccakWitInstance, RANGE_LOOKUPS, ROUNDS as KECCAK_ROUNDS,
+    KECCAK_STATE_PHASE_OUTPUT, KeccakInstance, KeccakLayout, KeccakParams, KeccakStateInstance,
+    KeccakTrace, KeccakWitInstance, RANGE_LOOKUPS, ROUNDS as KECCAK_ROUNDS,
     ROUNDS_CEIL_LOG2 as KECCAK_ROUNDS_CEIL_LOG2, XOR_LOOKUPS, keccak_state_record,
     run_lookup_keccakf, setup_gkr_circuit as setup_lookup_keccak_gkr_circuit,
 };

ceno_zkvm/src/scheme/mock_prover.rs

@@ -1498,6 +1498,18 @@ Hints:
             gs
         );
 
+        // part4 custom local buses
+        let (custom_rs, rs_grp_by_anno, custom_ws, ws_grp_by_anno, _) =
+            derive_ram_rws!(RAMType::Custom);
+        find_rw_mismatch!(
+            custom_rs,
+            rs_grp_by_anno,
+            custom_ws,
+            ws_grp_by_anno,
+            RAMType::Custom,
+            gs
+        );
+
         if num_rw_mismatch_errors > 0 {
             panic!("found {} r/w mismatch errors", num_rw_mismatch_errors);
         }

ceno_zkvm/src/structs.rs

@@ -18,7 +18,7 @@ use gkr_iop::{
 };
 use itertools::Itertools;
 use mpcs::{Point, PolynomialCommitmentScheme};
-use multilinear_extensions::Instance;
+use multilinear_extensions::{Expression, Instance, ToExpr};
 use p3::field::FieldAlgebra;
 use poseidon::challenger::{CanObserve, DefaultChallenger, FieldChallenger};
 use rayon::{
@@ -110,6 +110,18 @@ pub type ROMType = LookupTable;
 
 pub type RAMType = gkr_iop::RAMType;
 
+#[derive(Clone, Debug, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[repr(u16)]
+pub enum CustomRWTag {
+    KeccakState = 0,
+}
+
+impl CustomRWTag {
+    pub fn expr<E: ExtensionField>(self) -> Expression<E> {
+        E::BaseField::from_canonical_u16(self as u16).expr()
+    }
+}
+
 pub type PointAndEval<F> = multilinear_extensions::mle::PointAndEval<F>;
 
 #[derive(Clone)]

gkr_iop/src/lib.rs

@@ -106,6 +106,7 @@ pub enum RAMType {
     GlobalState = 0,
     Register,
     Memory,
+    Custom,
     Undefined,
 }