[feat] support weiestrass curve, add p256 (#1107)

* execution works

* fix proving

* cleanup

* cleanup

* make a optional

* update book

* restore

* restore

* Add a feature

* Fix pairing tests?

* Wip? Doesn't work

* test passes, without setup constraints

* feat: add setup constraints with a for ecdouble

* tests passgs

* fix

* fix

* fix

---------

Co-authored-by: Alexander Golovanov <[email protected]>
Co-authored-by: Jonathan Wang <[email protected]>

benchmarks/programs/ecrecover/src/main.rs

@@ -11,7 +11,7 @@ use k256::{
 use openvm::io::read_vec;
 #[allow(unused_imports)]
 use openvm_ecc_guest::{
-    algebra::IntMod, ecdsa::VerifyingKey, k256::Secp256k1Coord, weierstrass::WeierstrassPoint,
+    algebra::IntMod, ecdsa::VerifyingKey, k256::Secp256k1Point, weierstrass::WeierstrassPoint,
 };
 #[allow(unused_imports, clippy::single_component_path_imports)]
 use openvm_keccak256_guest; // export native keccak
@@ -26,7 +26,7 @@ openvm_algebra_guest::moduli_setup::moduli_init! {
     "0xFFFFFFFF FFFFFFFF FFFFFFFF FFFFFFFE BAAEDCE6 AF48A03B BFD25E8C D0364141"
 }
 openvm_ecc_guest::sw_setup::sw_init! {
-    Secp256k1Coord,
+    Secp256k1Point,
 }
 
 pub fn main() {

book/src/custom-extensions/ecc.md

@@ -34,18 +34,18 @@ For elliptic curve cryptography, the `openvm-ecc-guest` crate provides macros si
 ```rust
 sw_declare! {
     Bls12_381G1Affine { mod_type = Bls12_381Fp, b = BLS12_381_B },
-    Bn254G1Affine { mod_type = Bn254Fp, b = BN254_B },
+    P256Affine { mod_type = P256Coord, a = P256_A, b = P256_B },
 }
 ```
 
-Each declared curve must specify the `mod_type` (implementing `IntMod`) and a constant `b` for the Weierstrass curve equation \\(y^2 = x^3 + b\\).
-This creates `Bls12_381G1Affine` and `Bn254G1Affine` structs which implement the `Group` and `WeierstrassPoint` traits. The underlying memory layout of the structs uses the memory layout of the `Bls12_381Fp` and `Bn254Fp` structs, respectively.
+Each declared curve must specify the `mod_type` (implementing `IntMod`) and a constant `b` for the Weierstrass curve equation \\(y^2 = x^3 + ax + b\\). `a` is optional and defaults to 0 for short Weierstrass curves.
+This creates `Bls12_381G1Affine` and `P256Affine` structs which implement the `Group` and `WeierstrassPoint` traits. The underlying memory layout of the structs uses the memory layout of the `Bls12_381Fp` and `P256Coord` structs, respectively.
 
 2. **Init**: Called once, it enumerates these curves and allows the compiler to produce optimized instructions:
 
 ```rust
 sw_init! {
-    Bls12_381Fp, Bn254Fp,
+    Bls12_381G1Affine, P256Affine,
 }
 ```
 

crates/circuits/mod-builder/src/README.md

@@ -1,4 +1,35 @@
-### Example usage
+## Notes on flags and setup
+
+Setup opcode is a special op that verifies the modulus is correct.
+There are some chips that don't need it because we hardcode the modulus. E.g. the pairing ones.
+For those chips need setup, setup is derived: `setup = is_valid - sum(all_flags)`. Note that all chips have `is_valid`.
+Therefore when the chip needs setup and only supports one opcode, user won't explicitly create a flag for it
+and we will create a default flag for it on finalizing.
+
+There are two independent properties of the chip built by `ExprBuilder`: whether it needs setup, and whether it supports multiple (2 for now) flags, and hence four types of chips:
+
+| needs_setup | multi_flags | Example |
+|-------------|-------------|---------|
+| true        | true        | modular, Fp2 |
+| true        | false       | EcAdd and EcDouble |
+| false       | true        | Not supported, no such chips |
+| false       | false       | Pairing ones, hardcoded modulus so no setup needed |
+
+
+1. For the first type (modular and Fp2), there are two flags (e.g. `add_flag` and `sub_flag`) and `setup = is_valid - sum(all_flags)`. That is, when doing setup both flags are 0.
+2. For the second type (EcAdd and EcDouble), the chip only supports one operation so technically it doesn't need a flag. But for impelementation simplicity, we still create a dummy flag for it, and it's always 1 unless it's doing setup. And this `setup = is_valid - sum(all_flags)` still holds.
+3. No chip is in the third type right now.
+4. For the fourth type, there is no setup needed and no flags for selecting operations. Only `is_valid` is needed.
+
+### Finalizing
+
+The STARK backend requires the trace height to be a power of 2. Usually we pad the trace with empty (all 0) rows to make the height a power of 2. Since `is_valid` is 0 for padded rows, the constraints including interaction with memory are satisfied.
+However, there are some cases that all-0 row doesn't satisfy the constraints: when the computation involves non-zero constant values:
+
+- Some pairing chips involves adding a small constant value (like `1`). But since pairing chips don't have any flags, we will pad the trace with the last valid row and set `is_valid` to 0.
+- EcDouble for short Weierstrass curve: `y^2 = x^3 + ax + b` where `a != 0`. Last valid row with `is_valid = 0` won't work as in that case `setup = is_valid - sum(all_flags) = 0 - 1 = -1` is not a bool (0/1). So we will pad the trace with the first valid row (which is a setup row) and set `is_valid` to 0.
+
+## Example usage
 
 Ec Double chip:
 

crates/circuits/mod-builder/src/builder.rs

@@ -166,13 +166,16 @@ impl ExprBuilder {
         )
     }
 
+    /// Creates a new constant (compile-time known) FieldVariable from `value` where
+    /// the big integer `value` is decomposed into `num_limbs` limbs of `limb_bits` bits,
+    /// with `num_limbs, limb_bits` specified by the builder config.
     pub fn new_const(builder: Rc<RefCell<ExprBuilder>>, value: BigUint) -> FieldVariable {
         let mut borrowed = builder.borrow_mut();
         let index = borrowed.constants.len();
-        let limbs = big_uint_to_limbs(&value, borrowed.limb_bits);
-        let num_limbs = limbs.len();
-        let range_checker_bits = borrowed.range_checker_bits;
         let limb_bits = borrowed.limb_bits;
+        let num_limbs = borrowed.num_limbs;
+        let limbs = big_uint_to_num_limbs(&value, limb_bits, num_limbs);
+        let range_checker_bits = borrowed.range_checker_bits;
         borrowed.constants.push((value.clone(), limbs));
         drop(borrowed);
 

crates/circuits/mod-builder/src/core_chip.rs

@@ -123,11 +123,8 @@ where
         let writes: Vec<AB::Expr> = self
             .output_indices()
             .iter()
-            .map(|&i| vars[i].clone())
-            .collect::<Vec<_>>()
-            .concat()
-            .iter()
-            .map(|x| (*x).into())
+            .flat_map(|&i| vars[i].clone())
+            .map(Into::into)
             .collect();
 
         let opcode_flags_except_last = self.opcode_flag_idx.iter().map(|&i| flags[i]).collect_vec();
@@ -289,7 +286,8 @@ where
             return;
         }
         // We will copy over the core part of last row to padded rows (all rows after num_records).
-        let adapter_width = trace.width() - <Self::Air as BaseAir<F>>::width(&self.air);
+        let core_width = <Self::Air as BaseAir<F>>::width(&self.air);
+        let adapter_width = trace.width() - core_width;
         let last_row = trace
             .rows()
             .nth(num_records - 1)

crates/circuits/primitives/src/bigint/utils.rs

@@ -94,6 +94,7 @@ pub fn big_uint_to_limbs(x: &BigUint, limb_bits: usize) -> Vec<usize> {
 }
 
 pub fn big_uint_to_num_limbs(x: &BigUint, limb_bits: usize, num_limbs: usize) -> Vec<usize> {
+    // TODO: inefficient, should allocate `num_limbs` ahead of time.
     let limbs = big_uint_to_limbs(x, limb_bits);
     let num_limbs = max(num_limbs, limbs.len());
     limbs

examples/ecc/Cargo.toml

@@ -7,17 +7,12 @@ edition = "2021"
 members = []
 
 [dependencies]
-openvm = { git = "https://github.com/openvm-org/openvm.git" }
-openvm-platform = { git = "https://github.com/openvm-org/openvm.git" }
-openvm-algebra-guest = { git = "https://github.com/openvm-org/openvm.git" }
-openvm-ecc-guest = { git = "https://github.com/openvm-org/openvm.git", features = ["k256"] }
+openvm = { path = "../../crates/toolchain/openvm" }
+openvm-platform = { path = "../../crates/toolchain/platform" }
+openvm-algebra-guest = { path = "../../extensions/algebra/guest" }
+openvm-ecc-guest = { path = "../../extensions/ecc/guest", features = ["k256"] }
 hex-literal = { version = "0.4.1", default-features = false }
 
 [features]
 default = []
-std = [
-    "openvm/std",
-    "openvm-algebra-guest/std",
-    "openvm-ecc-guest/std",
-]
-
+std = ["openvm/std", "openvm-algebra-guest/std", "openvm-ecc-guest/std"]

examples/ecc/src/main.rs

@@ -17,7 +17,7 @@ openvm_algebra_guest::moduli_setup::moduli_init! {
 }
 
 openvm_ecc_guest::sw_setup::sw_init! {
-    Secp256k1Coord,
+    Secp256k1Point,
 }
 // ANCHOR_END: init
 

extensions/ecc/circuit/Cargo.toml

@@ -30,6 +30,8 @@ eyre = { workspace = true }
 num-integer = { workspace = true }
 serde = { workspace = true }
 serde_with = { workspace = true }
+itertools = { workspace = true }
+
 [dev-dependencies]
 openvm-stark-sdk = { workspace = true }
 openvm-mod-circuit-builder = { workspace = true, features = ["test-utils"] }
@@ -38,4 +40,9 @@ openvm-rv32-adapters = { workspace = true, features = ["test-utils"] }
 lazy_static = { workspace = true }
 
 [target.'cfg(not(target_os = "zkvm"))'.dependencies]
-openvm-ecc-guest = { workspace = true, features = ["halo2curves", "k256"] }
+openvm-ecc-guest = { workspace = true, features = [
+    "halo2curves",
+    "k256",
+    "p256",
+] }
+openvm-algebra-guest = { workspace = true }