intruction: better encoder

tcoratger · tcoratger · commit dc6822359842 · 2025-08-18T21:48:44.000+02:00
diff --git a/crates/leanVm/src/air/constant.rs b/crates/leanVm/src/air/constant.rs
@@ -41,8 +41,17 @@ pub(crate) const COL_INDEX_DEREF: usize = 8;
 pub(crate) const COL_INDEX_JUZ: usize = 9;
 /// An auxiliary operand used for multi-purpose logic, like specifying the `res` type in DEREF.
 pub(crate) const COL_INDEX_AUX: usize = 10;
+/// The opcode flag that activates the Poseidon16 precompile's constraints.
+pub(crate) const COL_INDEX_POSEIDON_16: usize = 11;
+/// The opcode flag that activates the Poseidon24 precompile's constraints.
+pub(crate) const COL_INDEX_POSEIDON_24: usize = 12;
+/// The opcode flag that activates the DotProduct precompile's constraints.
+pub(crate) const COL_INDEX_DOT_PRODUCT: usize = 13;
+/// The opcode flag that activates the MultilinearEval precompile's constraints.
+pub(crate) const COL_INDEX_MULTILINEAR_EVAL: usize = 14;
+
 /// The total number of columns representing a single decoded instruction from the bytecode.
-pub(crate) const N_BYTECODE_COLUMNS: usize = 11;
+pub(crate) const N_BYTECODE_COLUMNS: usize = 15;
 
 // Execution trace columns (committed by the prover)
 
diff --git a/crates/leanVm/src/bytecode/instruction/computation.rs b/crates/leanVm/src/bytecode/instruction/computation.rs
@@ -16,7 +16,7 @@ pub struct ComputationInstruction {
     /// The first operand of the computation.
     pub arg_a: MemOrConstant,
     /// The second operand of the computation.
-    pub arg_b: MemOrFp,
+    pub arg_c: MemOrFp,
     /// The memory location or constant that the result must be equal to.
     pub res: MemOrConstant,
 }
@@ -38,7 +38,7 @@ impl ComputationInstruction {
         memory_manager: &mut MemoryManager,
     ) -> Result<(), VirtualMachineError> {
         let val_a = run_context.value_from_mem_or_constant(&self.arg_a, memory_manager);
-        let val_b = run_context.value_from_mem_or_fp(&self.arg_b, memory_manager);
+        let val_b = run_context.value_from_mem_or_fp(&self.arg_c, memory_manager);
         let val_res = run_context.value_from_mem_or_constant(&self.res, memory_manager);
 
         match (val_a, val_b, val_res) {
diff --git a/crates/leanVm/src/bytecode/instruction/mod.rs b/crates/leanVm/src/bytecode/instruction/mod.rs
@@ -8,12 +8,13 @@ use p3_symmetric::Permutation;
 use poseidon16::Poseidon2_16Instruction;
 use poseidon24::Poseidon2_24Instruction;
 
-use super::operand::{MemOrConstant, MemOrFp, MemOrFpOrConstant};
+use super::operand::MemOrFpOrConstant;
 use crate::{
     air::constant::{
-        COL_INDEX_ADD, COL_INDEX_AUX, COL_INDEX_DEREF, COL_INDEX_FLAG_A, COL_INDEX_FLAG_B,
-        COL_INDEX_FLAG_C, COL_INDEX_JUZ, COL_INDEX_MUL, COL_INDEX_OPERAND_A, COL_INDEX_OPERAND_B,
-        COL_INDEX_OPERAND_C, N_INSTRUCTION_COLUMNS,
+        COL_INDEX_ADD, COL_INDEX_AUX, COL_INDEX_DEREF, COL_INDEX_DOT_PRODUCT, COL_INDEX_FLAG_A,
+        COL_INDEX_FLAG_B, COL_INDEX_FLAG_C, COL_INDEX_JUZ, COL_INDEX_MUL,
+        COL_INDEX_MULTILINEAR_EVAL, COL_INDEX_OPERAND_A, COL_INDEX_OPERAND_B, COL_INDEX_OPERAND_C,
+        COL_INDEX_POSEIDON_16, COL_INDEX_POSEIDON_24, N_INSTRUCTION_COLUMNS,
     },
     bytecode::operation::Operation,
     constant::{DIMENSION, F},
@@ -161,26 +162,18 @@ impl Instruction {
             Self::Computation(ComputationInstruction {
                 operation,
                 arg_a,
-                arg_b,
+                arg_c,
                 res,
             }) => {
                 // Set the appropriate opcode flag for ADD or MUL.
                 match operation {
                     Operation::Add => repr[COL_INDEX_ADD] = F::ONE,
                     Operation::Mul => repr[COL_INDEX_MUL] = F::ONE,
                 }
-                // Convert operands to their field and flag representations.
-                let (op_a, flag_a) = op_to_field(arg_a);
-                let (op_b, flag_b) = op_to_field_fp(arg_b);
-                let (op_c, flag_c) = op_to_field(res);
-
-                // Populate the representation vector.
-                repr[COL_INDEX_OPERAND_A] = op_a;
-                repr[COL_INDEX_FLAG_A] = flag_a;
-                repr[COL_INDEX_OPERAND_B] = op_b;
-                repr[COL_INDEX_FLAG_B] = flag_b;
-                repr[COL_INDEX_OPERAND_C] = op_c;
-                repr[COL_INDEX_FLAG_C] = flag_c;
+                // Convert operands to their field and flag representations using the helper methods.
+                (repr[COL_INDEX_OPERAND_A], repr[COL_INDEX_FLAG_A]) = arg_a.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_B], repr[COL_INDEX_FLAG_B]) = res.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_C], repr[COL_INDEX_FLAG_C]) = arg_c.to_field_and_flag();
             }
             Self::Deref(DerefInstruction {
                 shift_0,
@@ -221,37 +214,58 @@ impl Instruction {
                 // Set the JUZ opcode flag.
                 repr[COL_INDEX_JUZ] = F::ONE;
                 // Convert operands to their field and flag representations.
-                let (op_a, flag_a) = op_to_field(condition);
-                let (op_b, flag_b) = op_to_field_fp(updated_fp);
-                let (op_c, flag_c) = op_to_field(dest);
-
-                // Populate the representation vector.
-                repr[COL_INDEX_OPERAND_A] = op_a;
-                repr[COL_INDEX_FLAG_A] = flag_a;
-                repr[COL_INDEX_OPERAND_B] = op_b;
-                repr[COL_INDEX_FLAG_B] = flag_b;
-                repr[COL_INDEX_OPERAND_C] = op_c;
-                repr[COL_INDEX_FLAG_C] = flag_c;
+                (repr[COL_INDEX_OPERAND_A], repr[COL_INDEX_FLAG_A]) = condition.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_B], repr[COL_INDEX_FLAG_B]) =
+                    updated_fp.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_C], repr[COL_INDEX_FLAG_C]) = dest.to_field_and_flag();
+            }
+            Self::Poseidon2_16(Poseidon2_16Instruction { arg_a, arg_b, res }) => {
+                // Set the Poseidon16 opcode flag.
+                repr[COL_INDEX_POSEIDON_16] = F::ONE;
+                // Convert operands to their field and flag representations.
+                (repr[COL_INDEX_OPERAND_A], repr[COL_INDEX_FLAG_A]) = arg_a.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_B], repr[COL_INDEX_FLAG_B]) = arg_b.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_C], repr[COL_INDEX_FLAG_C]) = res.to_field_and_flag();
+            }
+            Self::Poseidon2_24(Poseidon2_24Instruction { arg_a, arg_b, res }) => {
+                // Set the Poseidon24 opcode flag.
+                repr[COL_INDEX_POSEIDON_24] = F::ONE;
+                // Convert operands to their field and flag representations.
+                (repr[COL_INDEX_OPERAND_A], repr[COL_INDEX_FLAG_A]) = arg_a.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_B], repr[COL_INDEX_FLAG_B]) = arg_b.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_C], repr[COL_INDEX_FLAG_C]) = res.to_field_and_flag();
+            }
+            Self::DotProduct(DotProductInstruction {
+                arg0,
+                arg1,
+                res,
+                size,
+            }) => {
+                // Set the DotProduct opcode flag.
+                repr[COL_INDEX_DOT_PRODUCT] = F::ONE;
+                // Convert operands to their field and flag representations.
+                (repr[COL_INDEX_OPERAND_A], repr[COL_INDEX_FLAG_A]) = arg0.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_B], repr[COL_INDEX_FLAG_B]) = arg1.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_C], repr[COL_INDEX_FLAG_C]) = res.to_field_and_flag();
+                // Use the AUX column to store the size of the vectors.
+                repr[COL_INDEX_AUX] = F::from_usize(*size);
+            }
+            Self::MultilinearEval(MultilinearEvalInstruction {
+                coeffs,
+                point,
+                res,
+                n_vars,
+            }) => {
+                // Set the MultilinearEval opcode flag.
+                repr[COL_INDEX_MULTILINEAR_EVAL] = F::ONE;
+                // Convert operands to their field and flag representations.
+                (repr[COL_INDEX_OPERAND_A], repr[COL_INDEX_FLAG_A]) = coeffs.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_B], repr[COL_INDEX_FLAG_B]) = point.to_field_and_flag();
+                (repr[COL_INDEX_OPERAND_C], repr[COL_INDEX_FLAG_C]) = res.to_field_and_flag();
+                // Use the AUX column to store the number of variables.
+                repr[COL_INDEX_AUX] = F::from_usize(*n_vars);
             }
-            // Precompiles do not set flags in the main instruction columns.
-            _ => {}
         }
         repr
     }
 }
-
-/// Helper to convert a `MemOrConstant` operand into its (value, flag) pair.
-fn op_to_field(op: &MemOrConstant) -> (F, F) {
-    match op {
-        MemOrConstant::Constant(c) => (*c, F::ONE),
-        MemOrConstant::MemoryAfterFp { offset } => (F::from_usize(*offset), F::ZERO),
-    }
-}
-
-/// Helper to convert a `MemOrFp` operand into its (value, flag) pair.
-fn op_to_field_fp(op: &MemOrFp) -> (F, F) {
-    match op {
-        MemOrFp::Fp => (F::ZERO, F::ONE), // Represents the `fp` register itself.
-        MemOrFp::MemoryAfterFp { offset } => (F::from_usize(*offset), F::ZERO),
-    }
-}
diff --git a/crates/leanVm/src/bytecode/operand.rs b/crates/leanVm/src/bytecode/operand.rs
@@ -1,3 +1,5 @@
+use p3_field::PrimeCharacteristicRing;
+
 use crate::constant::F;
 
 /// Represents a value that can either be a constant or a value from memory.
@@ -14,6 +16,23 @@ pub enum MemOrConstant {
     },
 }
 
+impl MemOrConstant {
+    /// Converts the operand into its raw field representation for the trace.
+    ///
+    /// Returns a tuple `(operand, flag)` where:
+    /// - `operand`: The field element representing the constant value or memory offset.
+    /// - `flag`: A flag that is `1` for a constant and `0` for a memory access.
+    #[must_use]
+    pub fn to_field_and_flag(&self) -> (F, F) {
+        match self {
+            // If it's a constant, the flag is 1 and the value is the constant itself.
+            Self::Constant(c) => (*c, F::ONE),
+            // If it's a memory location, the flag is 0 and the value is the offset.
+            Self::MemoryAfterFp { offset } => (F::from_usize(*offset), F::ZERO),
+        }
+    }
+}
+
 /// Represents a value that can be a memory location, the `fp` register itself, or a constant.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub enum MemOrFpOrConstant {
@@ -39,3 +58,20 @@ pub enum MemOrFp {
     /// The value of the frame pointer (`fp`) register itself.
     Fp,
 }
+
+impl MemOrFp {
+    /// Converts the operand into its raw field representation for the trace.
+    ///
+    /// Returns a tuple `(operand, flag)` where:
+    /// - `operand`: The field element representing the memory offset (or 0 if `Fp`).
+    /// - `flag`: A flag that is `1` for the `Fp` register and `0` for a memory access.
+    #[must_use]
+    pub fn to_field_and_flag(&self) -> (F, F) {
+        match self {
+            // If it's the frame pointer, the flag is 1 and the operand value is 0.
+            Self::Fp => (F::ZERO, F::ONE),
+            // If it's a memory location, the flag is 0 and the value is the offset.
+            Self::MemoryAfterFp { offset } => (F::from_usize(*offset), F::ZERO),
+        }
+    }
+}
diff --git a/crates/leanVm/src/core.rs b/crates/leanVm/src/core.rs
@@ -473,7 +473,7 @@ mod tests {
         let instruction = Instruction::Computation(ComputationInstruction {
             operation: Operation::Add,
             arg_a: MemOrConstant::Constant(F::ONE),
-            arg_b: MemOrFp::Fp,
+            arg_c: MemOrFp::Fp,
             res: MemOrConstant::MemoryAfterFp { offset: 0 },
         });
         // Execute: Update the PC based on this instruction.
