diff --git a/circuits/ecc/src/field_expression/builder.rs b/circuits/ecc/src/field_expression/builder.rs index ee0325b1b4..d01a3524d0 100644 --- a/circuits/ecc/src/field_expression/builder.rs +++ b/circuits/ecc/src/field_expression/builder.rs @@ -51,6 +51,8 @@ pub struct ExprBuilder { // The equations to compute the newly introduced variables. For trace gen only. pub computes: Vec, + + pub output_indices: Vec, } impl ExprBuilder { @@ -74,6 +76,7 @@ impl ExprBuilder { carry_limbs: vec![], constraints: vec![], computes: vec![], + output_indices: vec![], } } diff --git a/circuits/ecc/src/field_expression/field_variable.rs b/circuits/ecc/src/field_expression/field_variable.rs index 2d407edfa5..4553b01585 100644 --- a/circuits/ecc/src/field_expression/field_variable.rs +++ b/circuits/ecc/src/field_expression/field_variable.rs @@ -69,6 +69,12 @@ impl FieldVariable { builder.num_variables - 1 } + pub fn save_output(&mut self) { + let index = self.save(); + let mut builder = self.builder.borrow_mut(); + builder.output_indices.push(index); + } + pub fn canonical_limb_bits(&self) -> usize { self.builder.borrow().limb_bits } diff --git a/vm/src/intrinsics/ecc/sw/add_ne.rs b/vm/src/intrinsics/ecc/sw/add_ne.rs index 23201371cd..20dc786d17 100644 --- a/vm/src/intrinsics/ecc/sw/add_ne.rs +++ b/vm/src/intrinsics/ecc/sw/add_ne.rs @@ -1,227 +1,44 @@ -use std::{cell::RefCell, rc::Rc, sync::Arc}; +use std::{cell::RefCell, rc::Rc}; use afs_primitives::{ - bigint::check_carry_mod_to_zero::CheckCarryModToZeroSubAir, - var_range::{VariableRangeCheckerBus, VariableRangeCheckerChip}, - SubAir, TraceSubRowGenerator, + bigint::check_carry_mod_to_zero::CheckCarryModToZeroSubAir, var_range::VariableRangeCheckerBus, }; -use afs_stark_backend::{interaction::InteractionBuilder, rap::BaseAirWithPublicValues}; -use ax_ecc_primitives::field_expression::{ExprBuilder, FieldExpr, FieldExprCols}; +use ax_ecc_primitives::field_expression::{ExprBuilder, FieldExpr}; use num_bigint_dig::BigUint; -use p3_air::BaseAir; -use p3_field::{AbstractField, Field, PrimeField32}; -use super::super::{EcPoint, FIELD_ELEMENT_BITS}; -use crate::{ - arch::{ - instructions::EccOpcode, AdapterAirContext, AdapterRuntimeContext, DynAdapterInterface, - DynArray, MinimalInstruction, Result, VmAdapterInterface, VmCoreAir, VmCoreChip, - }, - system::program::Instruction, - utils::{biguint_to_limbs_vec, limbs_to_biguint}, -}; - -#[derive(Clone)] -pub struct SwEcAddNeCoreAir { - pub expr: FieldExpr, - pub offset: usize, -} - -impl SwEcAddNeCoreAir { - pub fn new( - modulus: BigUint, // The coordinate field. - num_limbs: usize, - limb_bits: usize, - range_bus: VariableRangeCheckerBus, - offset: usize, - ) -> Self { - assert!(modulus.bits() <= num_limbs * limb_bits); - let subair = CheckCarryModToZeroSubAir::new( - modulus.clone(), - limb_bits, - range_bus.index, - range_bus.range_max_bits, - FIELD_ELEMENT_BITS, - ); - let builder = ExprBuilder::new(modulus, limb_bits, num_limbs, range_bus.range_max_bits); - let builder = Rc::new(RefCell::new(builder)); - - let x1 = ExprBuilder::new_input(builder.clone()); - let y1 = ExprBuilder::new_input(builder.clone()); - let x2 = ExprBuilder::new_input(builder.clone()); - let y2 = ExprBuilder::new_input(builder.clone()); - let mut lambda = (y2 - y1.clone()) / (x2.clone() - x1.clone()); - let mut x3 = lambda.square() - x1.clone() - x2; - x3.save(); - let mut y3 = lambda * (x1 - x3.clone()) - y1; - y3.save(); - - let builder = builder.borrow().clone(); - let expr = FieldExpr { - builder, - check_carry_mod_to_zero: subair, - range_bus, - }; - Self { expr, offset } - } -} - -impl BaseAir for SwEcAddNeCoreAir { - fn width(&self) -> usize { - BaseAir::::width(&self.expr) - } -} - -impl BaseAirWithPublicValues for SwEcAddNeCoreAir {} - -impl VmCoreAir for SwEcAddNeCoreAir -where - I: VmAdapterInterface, - AdapterAirContext: - From>>, -{ - fn eval( - &self, - builder: &mut AB, - local: &[AB::Var], - _from_pc: AB::Var, - ) -> AdapterAirContext { - assert_eq!(local.len(), BaseAir::::width(&self.expr)); - self.expr.eval(builder, local); - - let FieldExprCols { - is_valid, - inputs, - vars, - flags, - .. - } = self.expr.load_vars(local); - assert_eq!(inputs.len(), 4); - assert_eq!(vars.len(), 3); - assert_eq!(flags.len(), 0); - let reads: Vec = inputs.concat().iter().map(|x| (*x).into()).collect(); - let writes: Vec = vars[1..].concat().iter().map(|x| (*x).into()).collect(); - - let expected_opcode = EccOpcode::EC_ADD_NE as usize; - let expected_opcode = AB::Expr::from_canonical_usize(expected_opcode); - - let instruction = MinimalInstruction { - is_valid: is_valid.into(), - opcode: expected_opcode + AB::Expr::from_canonical_usize(self.offset), - }; - - let ctx: AdapterAirContext<_, DynAdapterInterface<_>> = AdapterAirContext { - to_pc: None, - reads: reads.into(), - writes: writes.into(), - instruction: instruction.into(), - }; - ctx.into() - } -} - -pub struct SwEcAddNeCoreChip { - pub air: SwEcAddNeCoreAir, - pub range_checker: Arc, -} - -impl SwEcAddNeCoreChip { - pub fn new( - modulus: BigUint, - num_limbs: usize, - limb_bits: usize, - range_checker: Arc, - offset: usize, - ) -> Self { - let air = SwEcAddNeCoreAir::new(modulus, num_limbs, limb_bits, range_checker.bus(), offset); - Self { air, range_checker } - } -} - -pub struct SwEcAddNeCoreRecord { - pub p1: EcPoint, - pub p2: EcPoint, -} - -impl VmCoreChip for SwEcAddNeCoreChip -where - I: VmAdapterInterface, - I::Reads: Into>, - AdapterRuntimeContext: From>>, -{ - type Record = SwEcAddNeCoreRecord; - type Air = SwEcAddNeCoreAir; - - fn execute_instruction( - &self, - _instruction: &Instruction, - _from_pc: u32, - reads: I::Reads, - ) -> Result<(AdapterRuntimeContext, Self::Record)> { - // Input: 2 EcPoint, so total 4 field elements. - let field_element_limbs = self.air.expr.canonical_num_limbs(); - let limb_bits = self.air.expr.canonical_limb_bits(); - let data: DynArray<_> = reads.into(); - let data = data.0; - assert_eq!(data.len(), 4 * field_element_limbs); - let data_u32: Vec = data.iter().map(|x| x.as_canonical_u32()).collect(); - - let x1 = limbs_to_biguint(&data_u32[..field_element_limbs], limb_bits); - let y1 = limbs_to_biguint( - &data_u32[field_element_limbs..2 * field_element_limbs], - limb_bits, - ); - let x2 = limbs_to_biguint( - &data_u32[2 * field_element_limbs..3 * field_element_limbs], - limb_bits, - ); - let y2 = limbs_to_biguint( - &data_u32[3 * field_element_limbs..4 * field_element_limbs], - limb_bits, - ); - - let vars = self - .air - .expr - .execute(vec![x1.clone(), y1.clone(), x2.clone(), y2.clone()], vec![]); - assert_eq!(vars.len(), 3); // lambda, x3, y3 - let x3 = vars[1].clone(); - let y3 = vars[2].clone(); - - let x3_limbs = biguint_to_limbs_vec(x3, limb_bits, field_element_limbs); - let y3_limbs = biguint_to_limbs_vec(y3, limb_bits, field_element_limbs); - let writes = [x3_limbs, y3_limbs] - .concat() - .into_iter() - .map(|x| F::from_canonical_u32(x)) - .collect::>(); - let ctx = AdapterRuntimeContext::<_, DynAdapterInterface<_>>::without_pc(writes); - - Ok(( - ctx.into(), - SwEcAddNeCoreRecord { - p1: EcPoint { x: x1, y: y1 }, - p2: EcPoint { x: x2, y: y2 }, - }, - )) - } - - fn get_opcode_name(&self, _opcode: usize) -> String { - "SwEcAddNe".to_string() - } - - fn generate_trace_row(&self, row_slice: &mut [F], record: Self::Record) { - self.air.expr.generate_subrow( - ( - &self.range_checker, - vec![record.p1.x, record.p1.y, record.p2.x, record.p2.y], - vec![], - ), - row_slice, - ); - } - - fn air(&self) -> &Self::Air { - &self.air +use super::super::FIELD_ELEMENT_BITS; + +pub fn ec_add_ne_expr( + modulus: BigUint, // The coordinate field. + num_limbs: usize, + limb_bits: usize, + range_bus: VariableRangeCheckerBus, +) -> FieldExpr { + assert!(modulus.bits() <= num_limbs * limb_bits); + let subair = CheckCarryModToZeroSubAir::new( + modulus.clone(), + limb_bits, + range_bus.index, + range_bus.range_max_bits, + FIELD_ELEMENT_BITS, + ); + let builder = ExprBuilder::new(modulus, limb_bits, num_limbs, range_bus.range_max_bits); + let builder = Rc::new(RefCell::new(builder)); + + let x1 = ExprBuilder::new_input(builder.clone()); + let y1 = ExprBuilder::new_input(builder.clone()); + let x2 = ExprBuilder::new_input(builder.clone()); + let y2 = ExprBuilder::new_input(builder.clone()); + let mut lambda = (y2 - y1.clone()) / (x2.clone() - x1.clone()); + let mut x3 = lambda.square() - x1.clone() - x2; + x3.save_output(); + let mut y3 = lambda * (x1 - x3.clone()) - y1; + y3.save_output(); + + let builder = builder.borrow().clone(); + FieldExpr { + builder, + check_carry_mod_to_zero: subair, + range_bus, } } diff --git a/vm/src/intrinsics/ecc/sw/double.rs b/vm/src/intrinsics/ecc/sw/double.rs index fe6b615237..888aa52bf3 100644 --- a/vm/src/intrinsics/ecc/sw/double.rs +++ b/vm/src/intrinsics/ecc/sw/double.rs @@ -1,209 +1,42 @@ -use std::{cell::RefCell, rc::Rc, sync::Arc}; +use std::{cell::RefCell, rc::Rc}; use afs_primitives::{ - bigint::check_carry_mod_to_zero::CheckCarryModToZeroSubAir, - var_range::{VariableRangeCheckerBus, VariableRangeCheckerChip}, - SubAir, TraceSubRowGenerator, + bigint::check_carry_mod_to_zero::CheckCarryModToZeroSubAir, var_range::VariableRangeCheckerBus, }; -use afs_stark_backend::{interaction::InteractionBuilder, rap::BaseAirWithPublicValues}; -use ax_ecc_primitives::field_expression::{ExprBuilder, FieldExpr, FieldExprCols}; +use ax_ecc_primitives::field_expression::{ExprBuilder, FieldExpr}; use num_bigint_dig::BigUint; -use p3_air::BaseAir; -use p3_field::{AbstractField, Field, PrimeField32}; -use super::super::{EcPoint, FIELD_ELEMENT_BITS}; -use crate::{ - arch::{ - instructions::EccOpcode, AdapterAirContext, AdapterRuntimeContext, DynAdapterInterface, - DynArray, MinimalInstruction, Result, VmAdapterInterface, VmCoreAir, VmCoreChip, - }, - system::program::Instruction, - utils::{biguint_to_limbs_vec, limbs_to_biguint}, -}; - -#[derive(Clone)] -pub struct SwEcDoubleCoreAir { - pub expr: FieldExpr, - pub offset: usize, -} - -impl SwEcDoubleCoreAir { - pub fn new( - modulus: BigUint, // The coordinate field. - num_limbs: usize, - limb_bits: usize, - range_bus: VariableRangeCheckerBus, - offset: usize, - ) -> Self { - assert!(modulus.bits() <= num_limbs * limb_bits); - let subair = CheckCarryModToZeroSubAir::new( - modulus.clone(), - limb_bits, - range_bus.index, - range_bus.range_max_bits, - FIELD_ELEMENT_BITS, - ); - let builder = ExprBuilder::new(modulus, limb_bits, num_limbs, range_bus.range_max_bits); - let builder = Rc::new(RefCell::new(builder)); - - let mut x1 = ExprBuilder::new_input(builder.clone()); - let mut y1 = ExprBuilder::new_input(builder.clone()); - let mut lambda = x1.square().int_mul(3) / (y1.int_mul(2)); - let mut x3 = lambda.square() - x1.int_mul(2); - x3.save(); - let mut y3 = lambda * (x1 - x3.clone()) - y1; - y3.save(); - - let builder = builder.borrow().clone(); - let expr = FieldExpr { - builder, - check_carry_mod_to_zero: subair, - range_bus, - }; - Self { expr, offset } - } -} - -impl BaseAir for SwEcDoubleCoreAir { - fn width(&self) -> usize { - BaseAir::::width(&self.expr) - } -} - -impl BaseAirWithPublicValues for SwEcDoubleCoreAir {} - -impl VmCoreAir for SwEcDoubleCoreAir -where - I: VmAdapterInterface, - AdapterAirContext: - From>>, -{ - fn eval( - &self, - builder: &mut AB, - local: &[AB::Var], - _from_pc: AB::Var, - ) -> AdapterAirContext { - assert_eq!(local.len(), BaseAir::::width(&self.expr)); - self.expr.eval(builder, local); - - let FieldExprCols { - is_valid, - inputs, - vars, - flags, - .. - } = self.expr.load_vars(local); - assert_eq!(inputs.len(), 2); - assert_eq!(vars.len(), 3); - assert_eq!(flags.len(), 0); - let reads: Vec = inputs.concat().iter().map(|x| (*x).into()).collect(); - let writes: Vec = vars[1..].concat().iter().map(|x| (*x).into()).collect(); - - let expected_opcode = EccOpcode::EC_DOUBLE as usize; - let expected_opcode = AB::Expr::from_canonical_usize(expected_opcode); - - let instruction = MinimalInstruction { - is_valid: is_valid.into(), - opcode: expected_opcode + AB::Expr::from_canonical_usize(self.offset), - }; - - let ctx: AdapterAirContext<_, DynAdapterInterface<_>> = AdapterAirContext { - to_pc: None, - reads: reads.into(), - writes: writes.into(), - instruction: instruction.into(), - }; - ctx.into() - } -} - -pub struct SwEcDoubleCoreChip { - pub air: SwEcDoubleCoreAir, - pub range_checker: Arc, -} - -impl SwEcDoubleCoreChip { - pub fn new( - modulus: BigUint, - num_limbs: usize, - limb_bits: usize, - range_checker: Arc, - offset: usize, - ) -> Self { - let air = - SwEcDoubleCoreAir::new(modulus, num_limbs, limb_bits, range_checker.bus(), offset); - Self { air, range_checker } - } -} - -pub struct SwEcDoubleCoreRecord { - pub p1: EcPoint, -} - -impl VmCoreChip for SwEcDoubleCoreChip -where - I: VmAdapterInterface, - I::Reads: Into>, - AdapterRuntimeContext: From>>, -{ - type Record = SwEcDoubleCoreRecord; - type Air = SwEcDoubleCoreAir; - - fn execute_instruction( - &self, - _instruction: &Instruction, - _from_pc: u32, - reads: I::Reads, - ) -> Result<(AdapterRuntimeContext, Self::Record)> { - // Input: EcPoint, so total 2 field elements. - let field_element_limbs = self.air.expr.canonical_num_limbs(); - let limb_bits = self.air.expr.canonical_limb_bits(); - let data: DynArray<_> = reads.into(); - let data = data.0; - assert_eq!(data.len(), 2 * field_element_limbs); - let data_u32: Vec = data.iter().map(|x| x.as_canonical_u32()).collect(); - - let x1 = limbs_to_biguint(&data_u32[..field_element_limbs], limb_bits); - let y1 = limbs_to_biguint( - &data_u32[field_element_limbs..2 * field_element_limbs], - limb_bits, - ); - - let vars = self.air.expr.execute(vec![x1.clone(), y1.clone()], vec![]); - assert_eq!(vars.len(), 3); // lambda, x3, y3 - let x3 = vars[1].clone(); - let y3 = vars[2].clone(); - - let x3_limbs = biguint_to_limbs_vec(x3, limb_bits, field_element_limbs); - let y3_limbs = biguint_to_limbs_vec(y3, limb_bits, field_element_limbs); - let writes = [x3_limbs, y3_limbs] - .concat() - .into_iter() - .map(|x| F::from_canonical_u32(x)) - .collect::>(); - let ctx = AdapterRuntimeContext::<_, DynAdapterInterface<_>>::without_pc(writes); - - Ok(( - ctx.into(), - SwEcDoubleCoreRecord { - p1: EcPoint { x: x1, y: y1 }, - }, - )) - } - - fn get_opcode_name(&self, _opcode: usize) -> String { - "SwEcDouble".to_string() - } - - fn generate_trace_row(&self, row_slice: &mut [F], record: Self::Record) { - self.air.expr.generate_subrow( - (&self.range_checker, vec![record.p1.x, record.p1.y], vec![]), - row_slice, - ); - } - - fn air(&self) -> &Self::Air { - &self.air +use super::super::FIELD_ELEMENT_BITS; + +pub fn ec_double_expr( + modulus: BigUint, // The coordinate field. + num_limbs: usize, + limb_bits: usize, + range_bus: VariableRangeCheckerBus, +) -> FieldExpr { + assert!(modulus.bits() <= num_limbs * limb_bits); + let subair = CheckCarryModToZeroSubAir::new( + modulus.clone(), + limb_bits, + range_bus.index, + range_bus.range_max_bits, + FIELD_ELEMENT_BITS, + ); + let builder = ExprBuilder::new(modulus, limb_bits, num_limbs, range_bus.range_max_bits); + let builder = Rc::new(RefCell::new(builder)); + + let mut x1 = ExprBuilder::new_input(builder.clone()); + let mut y1 = ExprBuilder::new_input(builder.clone()); + let mut lambda = x1.square().int_mul(3) / (y1.int_mul(2)); + let mut x3 = lambda.square() - x1.int_mul(2); + x3.save_output(); + let mut y3 = lambda * (x1 - x3.clone()) - y1; + y3.save_output(); + + let builder = builder.borrow().clone(); + FieldExpr { + builder, + check_carry_mod_to_zero: subair, + range_bus, } } diff --git a/vm/src/intrinsics/ecc/sw/tests.rs b/vm/src/intrinsics/ecc/sw/tests.rs index 20105f42a2..e29d1b0f9d 100644 --- a/vm/src/intrinsics/ecc/sw/tests.rs +++ b/vm/src/intrinsics/ecc/sw/tests.rs @@ -7,10 +7,10 @@ use num_traits::FromPrimitive; use p3_baby_bear::BabyBear; use p3_field::AbstractField; -use super::SwEcAddNeCoreChip; +use super::{ec_add_ne_expr, ec_double_expr}; use crate::{ arch::{instructions::EccOpcode, testing::VmChipTestBuilder, VmChipWrapper}, - intrinsics::ecc::sw::SwEcDoubleCoreChip, + intrinsics::field_expression::FieldExpressionCoreChip, rv32im::adapters::{Rv32VecHeapAdapterChip, RV32_REGISTER_NUM_LIMBS}, system::program::Instruction, utils::biguint_to_limbs, @@ -24,12 +24,17 @@ type F = BabyBear; fn test_add_ne() { let mut tester: VmChipTestBuilder = VmChipTestBuilder::default(); let modulus = secp256k1_coord_prime(); - let core = SwEcAddNeCoreChip::new( - modulus.clone(), + let expr = ec_add_ne_expr( + modulus, NUM_LIMBS, LIMB_BITS, - tester.memory_controller().borrow().range_checker.clone(), + tester.memory_controller().borrow().range_checker.bus(), + ); + let core = FieldExpressionCoreChip::new( + expr, EccOpcode::default_offset(), + tester.memory_controller().borrow().range_checker.clone(), + "EcAddNe", ); let adapter = Rv32VecHeapAdapterChip::::new( tester.execution_bus(), @@ -104,12 +109,17 @@ fn test_add_ne() { fn test_double() { let mut tester: VmChipTestBuilder = VmChipTestBuilder::default(); let modulus = secp256k1_coord_prime(); - let core = SwEcDoubleCoreChip::new( - modulus.clone(), + let expr = ec_double_expr( + modulus, NUM_LIMBS, LIMB_BITS, - tester.memory_controller().borrow().range_checker.clone(), + tester.memory_controller().borrow().range_checker.bus(), + ); + let core = FieldExpressionCoreChip::new( + expr, EccOpcode::default_offset(), + tester.memory_controller().borrow().range_checker.clone(), + "EcDouble", ); let adapter = Rv32VecHeapAdapterChip::::new( tester.execution_bus(), diff --git a/vm/src/intrinsics/field_expression/mod.rs b/vm/src/intrinsics/field_expression/mod.rs new file mode 100644 index 0000000000..9a8f4da04b --- /dev/null +++ b/vm/src/intrinsics/field_expression/mod.rs @@ -0,0 +1,206 @@ +use std::sync::Arc; + +use afs_primitives::{var_range::VariableRangeCheckerChip, SubAir, TraceSubRowGenerator}; +use afs_stark_backend::{interaction::InteractionBuilder, rap::BaseAirWithPublicValues}; +use ax_ecc_primitives::field_expression::{FieldExpr, FieldExprCols}; +use itertools::Itertools; +use num_bigint_dig::BigUint; +use p3_air::BaseAir; +use p3_field::{AbstractField, Field, PrimeField32}; + +use crate::{ + arch::{ + AdapterAirContext, AdapterRuntimeContext, DynAdapterInterface, DynArray, + MinimalInstruction, Result, VmAdapterInterface, VmCoreAir, VmCoreChip, + }, + system::program::Instruction, + utils::{biguint_to_limbs_vec, limbs_to_biguint}, +}; +#[derive(Clone)] +pub struct FieldExpressionCoreAir { + pub expr: FieldExpr, + pub offset: usize, +} + +impl FieldExpressionCoreAir { + pub fn new(expr: FieldExpr, offset: usize) -> Self { + Self { expr, offset } + } + + pub fn num_inputs(&self) -> usize { + self.expr.builder.num_input + } + + pub fn num_vars(&self) -> usize { + self.expr.builder.num_variables + } + + pub fn num_flags(&self) -> usize { + self.expr.builder.num_flags + } + + pub fn output_indices(&self) -> &[usize] { + &self.expr.builder.output_indices + } +} + +impl BaseAir for FieldExpressionCoreAir { + fn width(&self) -> usize { + BaseAir::::width(&self.expr) + } +} + +impl BaseAirWithPublicValues for FieldExpressionCoreAir {} + +impl VmCoreAir for FieldExpressionCoreAir +where + I: VmAdapterInterface, + AdapterAirContext: + From>>, +{ + fn eval( + &self, + builder: &mut AB, + local: &[AB::Var], + _from_pc: AB::Var, + ) -> AdapterAirContext { + assert_eq!(local.len(), BaseAir::::width(&self.expr)); + self.expr.eval(builder, local); + let FieldExprCols { + is_valid, + inputs, + vars, + flags, + .. + } = self.expr.load_vars(local); + assert_eq!(inputs.len(), self.num_inputs()); + assert_eq!(vars.len(), self.num_vars()); + assert_eq!(flags.len(), self.num_flags()); + let reads: Vec = inputs.concat().iter().map(|x| (*x).into()).collect(); + let writes: Vec = self + .output_indices() + .iter() + .map(|&i| vars[i].clone()) + .collect::>() + .concat() + .iter() + .map(|x| (*x).into()) + .collect(); + + // TODO: flags -> opcode + let expected_opcode = AB::Expr::from_canonical_usize(self.offset); + + let instruction = MinimalInstruction { + is_valid: is_valid.into(), + opcode: expected_opcode, + }; + + let ctx: AdapterAirContext<_, DynAdapterInterface<_>> = AdapterAirContext { + to_pc: None, + reads: reads.into(), + writes: writes.into(), + instruction: instruction.into(), + }; + ctx.into() + } +} + +pub struct FieldExpressionRecord { + pub inputs: Vec, +} + +pub struct FieldExpressionCoreChip { + pub air: FieldExpressionCoreAir, + pub range_checker: Arc, + + pub name: String, +} + +impl FieldExpressionCoreChip { + pub fn new( + expr: FieldExpr, + offset: usize, + range_checker: Arc, + name: &str, + ) -> Self { + let air = FieldExpressionCoreAir { expr, offset }; + Self { + air, + range_checker, + name: name.to_string(), + } + } +} + +impl VmCoreChip for FieldExpressionCoreChip +where + I: VmAdapterInterface, + I::Reads: Into>, + AdapterRuntimeContext: From>>, +{ + type Record = FieldExpressionRecord; + type Air = FieldExpressionCoreAir; + + fn execute_instruction( + &self, + instruction: &Instruction, + _from_pc: u32, + reads: I::Reads, + ) -> Result<(AdapterRuntimeContext, Self::Record)> { + let field_element_limbs = self.air.expr.canonical_num_limbs(); + let limb_bits = self.air.expr.canonical_limb_bits(); + let data: DynArray<_> = reads.into(); + let data = data.0; + assert_eq!(data.len(), self.air.num_inputs() * field_element_limbs); + let data_u32: Vec = data.iter().map(|x| x.as_canonical_u32()).collect(); + + let mut inputs = vec![]; + for i in 0..self.air.num_inputs() { + let start = i * field_element_limbs; + let end = start + field_element_limbs; + let limb_slice = &data_u32[start..end]; + let input = limbs_to_biguint(limb_slice, limb_bits); + inputs.push(input); + } + + // TODO: local_opcode_index -> flags + let Instruction { opcode, .. } = instruction.clone(); + let _local_opcode_index = opcode - self.air.offset; + let flags = vec![]; + assert_eq!(flags.len(), self.air.num_flags()); + + let vars = self.air.expr.execute(inputs.clone(), flags); + assert_eq!(vars.len(), self.air.num_vars()); + + let outputs: Vec = self + .air + .output_indices() + .iter() + .map(|&i| vars[i].clone()) + .collect(); + let writes: Vec = outputs + .iter() + .map(|x| biguint_to_limbs_vec(x.clone(), limb_bits, field_element_limbs)) + .concat() + .into_iter() + .map(|x| F::from_canonical_u32(x)) + .collect(); + + let ctx = AdapterRuntimeContext::<_, DynAdapterInterface<_>>::without_pc(writes); + Ok((ctx.into(), FieldExpressionRecord { inputs })) + } + + fn get_opcode_name(&self, _opcode: usize) -> String { + self.name.clone() + } + + fn generate_trace_row(&self, row_slice: &mut [F], record: Self::Record) { + self.air + .expr + .generate_subrow((&self.range_checker, record.inputs, vec![]), row_slice); + } + + fn air(&self) -> &Self::Air { + &self.air + } +} diff --git a/vm/src/intrinsics/mod.rs b/vm/src/intrinsics/mod.rs index 1545af12c4..ab4efcb3f1 100644 --- a/vm/src/intrinsics/mod.rs +++ b/vm/src/intrinsics/mod.rs @@ -1,3 +1,4 @@ pub mod ecc; +pub mod field_expression; pub mod hashes; pub mod modular; diff --git a/vm/src/kernels/ecc/mod.rs b/vm/src/kernels/ecc/mod.rs index 0c83ec205b..784f187fa8 100644 --- a/vm/src/kernels/ecc/mod.rs +++ b/vm/src/kernels/ecc/mod.rs @@ -1,20 +1,171 @@ -use super::adapters::native_vec_heap_adapter::{NativeVecHeapAdapterAir, NativeVecHeapAdapterChip}; +use std::sync::Arc; + +use afs_primitives::bigint::utils::secp256k1_coord_prime; +use afs_stark_backend::{ + config::{StarkGenericConfig, Val}, + prover::types::AirProofInput, + rap::AnyRap, + Chip, ChipUsageGetter, +}; +use p3_field::PrimeField32; + +use super::adapters::native_vec_heap_adapter::NativeVecHeapAdapterChip; use crate::{ - arch::{VmAirWrapper, VmChipWrapper}, - intrinsics::ecc::sw::{ - SwEcAddNeCoreAir, SwEcAddNeCoreChip, SwEcDoubleCoreAir, SwEcDoubleCoreChip, + arch::{ExecutionState, InstructionExecutor, Result, VmChipWrapper}, + intrinsics::{ + ecc::sw::{ec_add_ne_expr, ec_double_expr}, + field_expression::FieldExpressionCoreChip, }, + system::{memory::MemoryControllerRef, program::Instruction}, }; -pub type KernelEcAddNeAir = - VmAirWrapper, SwEcAddNeCoreAir>; -pub type KernelEcAddNeChip = - VmChipWrapper, SwEcAddNeCoreChip>; - -pub type KernelEcDoubleAir = - VmAirWrapper, SwEcDoubleCoreAir>; -pub type KernelEcDoubleChip = VmChipWrapper< - F, - NativeVecHeapAdapterChip, - SwEcDoubleCoreChip, ->; +pub struct KernelEcAddNeChip( + VmChipWrapper< + F, + NativeVecHeapAdapterChip, + FieldExpressionCoreChip, + >, +); + +impl KernelEcAddNeChip { + pub fn new( + adapter: NativeVecHeapAdapterChip, + memory_controller: MemoryControllerRef, + limb_bits: usize, + offset: usize, + ) -> Self { + let expr = ec_add_ne_expr( + secp256k1_coord_prime(), + NUM_LIMBS, + limb_bits, + memory_controller.borrow().range_checker.bus(), + ); + let core = FieldExpressionCoreChip::new( + expr, + offset, + memory_controller.borrow().range_checker.clone(), + "EcAddNe", + ); + Self(VmChipWrapper::new(adapter, core, memory_controller)) + } +} + +pub struct KernelEcDoubleChip( + VmChipWrapper< + F, + NativeVecHeapAdapterChip, + FieldExpressionCoreChip, + >, +); + +impl KernelEcDoubleChip { + pub fn new( + adapter: NativeVecHeapAdapterChip, + memory_controller: MemoryControllerRef, + limb_bits: usize, + offset: usize, + ) -> Self { + let expr = ec_double_expr( + secp256k1_coord_prime(), + NUM_LIMBS, + limb_bits, + memory_controller.borrow().range_checker.bus(), + ); + let core = FieldExpressionCoreChip::new( + expr, + offset, + memory_controller.borrow().range_checker.clone(), + "EcDouble", + ); + Self(VmChipWrapper::new(adapter, core, memory_controller)) + } +} + +// TODO: below can be generated by the macro. +impl InstructionExecutor + for KernelEcAddNeChip +{ + fn execute( + &mut self, + instruction: Instruction, + from_state: ExecutionState, + ) -> Result> { + self.0.execute(instruction, from_state) + } + + fn get_opcode_name(&self, opcode: usize) -> String { + self.0.get_opcode_name(opcode) + } +} + +impl ChipUsageGetter for KernelEcAddNeChip { + fn air_name(&self) -> String { + self.0.air_name() + } + + fn current_trace_height(&self) -> usize { + self.0.current_trace_height() + } + + fn trace_width(&self) -> usize { + self.0.trace_width() + } +} + +impl Chip for KernelEcAddNeChip, NUM_LIMBS> +where + SC: StarkGenericConfig, + Val: PrimeField32, +{ + fn air(&self) -> Arc> { + self.0.air() + } + + fn generate_air_proof_input(self) -> AirProofInput { + self.0.generate_air_proof_input() + } +} + +impl InstructionExecutor + for KernelEcDoubleChip +{ + fn execute( + &mut self, + instruction: Instruction, + from_state: ExecutionState, + ) -> Result> { + self.0.execute(instruction, from_state) + } + + fn get_opcode_name(&self, opcode: usize) -> String { + self.0.get_opcode_name(opcode) + } +} + +impl ChipUsageGetter for KernelEcDoubleChip { + fn air_name(&self) -> String { + self.0.air_name() + } + + fn current_trace_height(&self) -> usize { + self.0.current_trace_height() + } + + fn trace_width(&self) -> usize { + self.0.trace_width() + } +} + +impl Chip for KernelEcDoubleChip, NUM_LIMBS> +where + SC: StarkGenericConfig, + Val: PrimeField32, +{ + fn air(&self) -> Arc> { + self.0.air() + } + + fn generate_air_proof_input(self) -> AirProofInput { + self.0.generate_air_proof_input() + } +} diff --git a/vm/src/system/vm/chip_set.rs b/vm/src/system/vm/chip_set.rs index 747b8af51d..f5bf3a4865 100644 --- a/vm/src/system/vm/chip_set.rs +++ b/vm/src/system/vm/chip_set.rs @@ -8,7 +8,6 @@ use std::{ }; use afs_primitives::{ - bigint::utils::secp256k1_coord_prime, range_tuple::{RangeTupleCheckerBus, RangeTupleCheckerChip}, var_range::{VariableRangeCheckerBus, VariableRangeCheckerChip}, xor::XorLookupChip, @@ -34,7 +33,6 @@ use crate::{ arch::{AxVmChip, AxVmInstructionExecutor, ExecutionBus, ExecutorName}, common::nop::NopChip, intrinsics::{ - ecc::sw::{SwEcAddNeCoreChip, SwEcDoubleCoreChip}, hashes::{keccak::hasher::KeccakVmChip, poseidon2::Poseidon2Chip}, modular::{ ModularAddSubChip, ModularAddSubCoreChip, ModularMulDivChip, ModularMulDivCoreChip, @@ -676,14 +674,9 @@ impl VmConfig { program_bus, memory_controller.clone(), ), - SwEcAddNeCoreChip::new( - secp256k1_coord_prime(), - 32, - 8, - memory_controller.borrow().range_checker.clone(), - offset, - ), memory_controller.clone(), + 8, + offset, ))); for opcode in range { executors.insert(opcode, chip.clone().into()); @@ -697,14 +690,9 @@ impl VmConfig { program_bus, memory_controller.clone(), ), - SwEcDoubleCoreChip::new( - secp256k1_coord_prime(), - 32, - 8, - memory_controller.borrow().range_checker.clone(), - offset, - ), memory_controller.clone(), + 8, + offset, ))); for opcode in range { executors.insert(opcode, chip.clone().into());