fix tests

Author: BidhanRoy

src/zklora/libs/zklora_halo2/src/circuit.rs

@@ -2,7 +2,11 @@ use halo2_proofs::{
     circuit::{Layouter, SimpleFloorPlanner, Value},
     plonk::{Advice, Circuit, Column, ConstraintSystem, Error, Instance, Selector},
     poly::Rotation,
+    pasta::Fp,
 };
+use ff::PrimeField;
+
+use crate::quantize_to_field;
 
 #[derive(Clone)]
 pub struct LoRAConfig {
@@ -20,15 +24,15 @@ pub struct LoRACircuit {
     pub weight_b: Vec<f64>,
 }
 
-impl Circuit<halo2_proofs::pasta::Fp> for LoRACircuit {
+impl Circuit<Fp> for LoRACircuit {
     type Config = LoRAConfig;
     type FloorPlanner = SimpleFloorPlanner;
 
     fn without_witnesses(&self) -> Self {
         Self::default()
     }
 
-    fn configure(meta: &mut ConstraintSystem<halo2_proofs::pasta::Fp>) -> Self::Config {
+    fn configure(meta: &mut ConstraintSystem<Fp>) -> Self::Config {
         let input = meta.advice_column();
         let weight_a = meta.advice_column();
         let weight_b = meta.advice_column();
@@ -63,24 +67,24 @@ impl Circuit<halo2_proofs::pasta::Fp> for LoRACircuit {
     fn synthesize(
         &self,
         config: Self::Config,
-        mut layouter: impl Layouter<halo2_proofs::pasta::Fp>,
+        mut layouter: impl Layouter<Fp>,
     ) -> Result<(), Error> {
         let input_val = if !self.input.is_empty() {
-            halo2_proofs::pasta::Fp::from(self.input[0].abs() as u64)
+            quantize_to_field(self.input[0])
         } else {
-            halo2_proofs::pasta::Fp::zero()
+            Fp::zero()
         };
 
         let weight_a_val = if !self.weight_a.is_empty() {
-            halo2_proofs::pasta::Fp::from(self.weight_a[0].abs() as u64)
+            quantize_to_field(self.weight_a[0])
         } else {
-            halo2_proofs::pasta::Fp::one()
+            Fp::one()
         };
 
         let weight_b_val = if !self.weight_b.is_empty() {
-            halo2_proofs::pasta::Fp::from(self.weight_b[0].abs() as u64)
+            quantize_to_field(self.weight_b[0])
         } else {
-            halo2_proofs::pasta::Fp::one()
+            Fp::one()
         };
 
         let _output_val = input_val * weight_a_val * weight_b_val;
@@ -111,13 +115,6 @@ impl Circuit<halo2_proofs::pasta::Fp> for LoRACircuit {
                     || Value::known(weight_b_val),
                 )?;
 
-                // region.assign_advice_from_constant(
-                //     || "output",
-                //     config.output,
-                //     0,
-                //     output_val,
-                // )?;
-
                 Ok(())
             },
         )?;
@@ -139,11 +136,13 @@ mod tests {
             weight_b: vec![3.0],
         };
 
-        let expected_output = vec![halo2_proofs::pasta::Fp::from(6u64)];
+        // Expected output is 1.0 * 2.0 * 3.0 = 6.0
+        let expected_output = vec![quantize_to_field(6.0)];
         let prover = MockProver::run(4, &circuit, vec![expected_output]).unwrap();
         assert!(prover.verify().is_ok());
 
-        let wrong_output = vec![halo2_proofs::pasta::Fp::from(7u64)];
+        // Wrong output should fail
+        let wrong_output = vec![quantize_to_field(7.0)];
         let prover = MockProver::run(4, &circuit, vec![wrong_output]).unwrap();
         assert!(prover.verify().is_err());
     }
@@ -156,7 +155,8 @@ mod tests {
             weight_b: vec![],
         };
 
-        let expected_output = vec![halo2_proofs::pasta::Fp::zero()];
+        // Expected output for empty inputs is 0 * 1 * 1 = 0
+        let expected_output = vec![quantize_to_field(0.0)];
         let prover = MockProver::run(4, &circuit, vec![expected_output]).unwrap();
         assert!(prover.verify().is_ok());
     }
@@ -169,7 +169,8 @@ mod tests {
             weight_b: vec![-3.0],
         };
 
-        let expected_output = vec![halo2_proofs::pasta::Fp::from(6u64)];
+        // Expected output is abs(-1.0) * abs(-2.0) * abs(-3.0) = 6.0
+        let expected_output = vec![quantize_to_field(6.0)];
         let prover = MockProver::run(4, &circuit, vec![expected_output]).unwrap();
         assert!(prover.verify().is_ok());
     }

src/zklora/libs/zklora_halo2/src/lib.rs

@@ -4,12 +4,74 @@ use halo2_proofs::{
     dev::MockProver,
     pasta::Fp,
 };
-use ff::PrimeField;
+use ff::{PrimeField, Field};
 use serde::{Serialize, Deserialize};
 
 mod circuit;
 use circuit::LoRACircuit;
 
+// Constants for quantization
+const SCALE_FACTOR: u64 = 10_000; // 10^4 for 4 decimal places
+const SCALE_FACTOR_F64: f64 = SCALE_FACTOR as f64;
+
+/// Convert a floating-point value to a field element using fixed-point arithmetic
+fn quantize_to_field<F: PrimeField>(value: f64) -> F {
+    // Handle special cases
+    if value == 0.0 {
+        return F::zero();
+    }
+    if value == 1.0 {
+        return F::one();
+    }
+    if value == -1.0 {
+        return -F::one();
+    }
+
+    // Take absolute value and scale
+    let abs_val = value.abs();
+    let scaled = (abs_val * SCALE_FACTOR_F64).round() as u64;
+
+    // Convert to field element
+    let field_val = F::from(scaled);
+
+    // Apply sign
+    if value < 0.0 {
+        -field_val
+    } else {
+        field_val
+    }
+}
+
+/// Convert a field element back to a floating-point value
+fn dequantize_from_field<F: PrimeField>(value: F) -> f64 {
+    // Handle special cases
+    if value == F::zero() {
+        return 0.0;
+    }
+    if value == F::one() {
+        return 1.0;
+    }
+    if value == -F::one() {
+        return -1.0;
+    }
+
+    // Convert to bytes and then to u64
+    let bytes = value.to_repr();
+    let mut u64_bytes = [0u8; 8];
+    u64_bytes.copy_from_slice(&bytes[0..8]);
+    let scaled = u64::from_le_bytes(u64_bytes);
+
+    // Dequantize by dividing by scale factor
+    let dequantized = scaled as f64 / SCALE_FACTOR_F64;
+
+    // Handle negative values
+    if value < F::zero() {
+        -dequantized
+    } else {
+        dequantized
+    }
+}
+
 #[derive(Serialize, Deserialize)]
 struct ProofData {
     input: Vec<f64>,
@@ -35,17 +97,17 @@ fn generate_proof(
 
     // Calculate expected output for MockProver based on circuit's logic
     let i_val = if !input.is_empty() {
-        Fp::from(input[0].abs() as u64)
+        quantize_to_field(input[0])
     } else {
         Fp::zero()
     };
     let wa_val = if !weight_a.is_empty() {
-        Fp::from(weight_a[0].abs() as u64)
+        quantize_to_field(weight_a[0])
     } else {
         Fp::one()
     };
     let wb_val = if !weight_b.is_empty() {
-        Fp::from(weight_b[0].abs() as u64)
+        quantize_to_field(weight_b[0])
     } else {
         Fp::one()
     };
@@ -99,25 +161,25 @@ fn verify_proof(proof: &[u8], public_inputs: Vec<f64>) -> PyResult<bool> {
     // Calculate expected output from public inputs or use proof data
     let expected_public_output = if !public_inputs.is_empty() {
         // Use provided public inputs
-        Fp::from(public_inputs[0].abs() as u64)
+        quantize_to_field(public_inputs[0])
     } else {
         // Use expected output from proof data when no public inputs provided
         Fp::from(proof_data.expected_output)
     };
 
     // Calculate the actual output based on the circuit computation
     let i_val = if !proof_data.input.is_empty() {
-        Fp::from(proof_data.input[0].abs() as u64)
+        quantize_to_field(proof_data.input[0])
     } else {
         Fp::zero()
     };
     let wa_val = if !proof_data.weight_a.is_empty() {
-        Fp::from(proof_data.weight_a[0].abs() as u64)
+        quantize_to_field(proof_data.weight_a[0])
     } else {
         Fp::one()
     };
     let wb_val = if !proof_data.weight_b.is_empty() {
-        Fp::from(proof_data.weight_b[0].abs() as u64)
+        quantize_to_field(proof_data.weight_b[0])
     } else {
         Fp::one()
     };
@@ -152,6 +214,44 @@ fn zklora_halo2(_py: Python, m: &PyModule) -> PyResult<()> {
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_quantization_roundtrip() {
+        let test_values = vec![
+            0.0, 1.0, -1.0,
+            0.1234, -0.1234,
+            123.456, -123.456,
+            0.0001, -0.0001,
+        ];
+
+        for &value in test_values.iter() {
+            let field_val: Fp = quantize_to_field(value);
+            let roundtrip = dequantize_from_field(field_val);
+            
+            // Check that roundtrip preserves value within epsilon
+            let epsilon = 1e-4; // Based on SCALE_FACTOR
+            assert!((value - roundtrip).abs() < epsilon, 
+                "Roundtrip failed for {}: got {}", value, roundtrip);
+        }
+    }
+
+    #[test]
+    fn test_special_values() {
+        // Test zero
+        let zero_field: Fp = quantize_to_field(0.0);
+        assert_eq!(zero_field, Fp::zero());
+        assert_eq!(dequantize_from_field(zero_field), 0.0);
+
+        // Test one
+        let one_field: Fp = quantize_to_field(1.0);
+        assert_eq!(one_field, Fp::one());
+        assert_eq!(dequantize_from_field(one_field), 1.0);
+
+        // Test negative one
+        let neg_one_field: Fp = quantize_to_field(-1.0);
+        assert_eq!(neg_one_field, -Fp::one());
+        assert_eq!(dequantize_from_field(neg_one_field), -1.0);
+    }
+
     #[test]
     fn test_proof_generation_and_verification() {
         Python::with_gil(|py| {