[refactor] primitives (#621)

* feat: rework SubAir trait

* refactor: done with AssertLessThanAir

* refactor: var_range

* refactor: range

* chore: delete assert_sorted

* wip: is less than

* feat: better documentation of `SubAir` concept

* refactor: is_less_than

* chore: delete `sum`

we can easily restore later

* chore: docs

* refactor: xor

* refactor: is_zero

* refactor: is_zero

* feat: is_equal_array following BEQ air

* chore: cleanup, just bigint left

* chore: fix derive

* chore: doc comments

* refactor: finish primitives

* wip

* feat: IsLtArrayAir (SubAir)

* fix tests

* subair: call me by my name

* fix: put back tests

* fix: static verifier must have traces sorted

* chore: combine tests to reduce compile time

* chore: split out toolchain tests for faster CI

* chore: display the diff

* fix: toolchain tests weren't using runs-on

* todo: turn off CI diff check for ELF

.github/workflows/recursion.yml

@@ -39,5 +39,4 @@ jobs:
       - name: Run recursion crate tests
         working-directory: lib/recursion
         run: |
-          cargo nextest run --cargo-profile=fast --features parallel
-          cargo nextest run --cargo-profile=fast --features parallel,static-verifier -- halo2::tests
+          cargo nextest run --features parallel,static-verifier

.github/workflows/toolchain-edsl.yml

@@ -0,0 +1,37 @@
+name: eDSL Compiler Tests
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+    branches: ["main"]
+    paths:
+      - "vm/**"
+      - "toolchain/compiler/**"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: true
+
+env:
+  CARGO_TERM_COLOR: always
+  AXIOM_FAST_TEST: "1"
+
+jobs:
+  tests:
+    runs-on:
+      - runs-on=${{ github.run_id }}
+      - runner=64cpu-linux-arm64
+
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@nightly
+      - uses: Swatinem/rust-cache@v2
+        with:
+          cache-on-failure: true
+      - uses: taiki-e/install-action@nextest
+
+      - name: Run compiler tests
+        working-directory: toolchain/compiler
+        run: |
+          cargo nextest run --cargo-profile=fast --features parallel

.github/workflows/toolchain.yml

@@ -7,7 +7,7 @@ on:
     branches: ["main"]
     paths:
       - "vm/**"
-      - "toolchain/**"
+      - "toolchain/riscv/**"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
@@ -19,7 +19,9 @@ env:
 
 jobs:
   tests:
-    runs-on: ubuntu-arm64-16core-64ram
+    runs-on:
+      - runs-on=${{ github.run_id }}
+      - runner=64cpu-linux-arm64
 
     steps:
       - uses: actions/checkout@v4
@@ -34,15 +36,10 @@ jobs:
         run: |
           rustup component add rust-src --toolchain nightly-aarch64-unknown-linux-gnu
           cargo +nightly build --release
-          ls target
-          cp target/riscv32im-risc0-zkvm-elf/release/axvm-fibonacci-program ../../../transpiler/data/rv32im-fibonacci-program-elf-release
+          # TODO: aarch and linux builds have different ELFs
+          # diff target/riscv32im-risc0-zkvm-elf/release/axvm-fibonacci-program ../../../transpiler/data/rv32im-fibonacci-program-elf-release
 
       - name: Run transpiler tests
         working-directory: toolchain/riscv/transpiler
         run: |
           cargo nextest run --cargo-profile=fast
-
-      - name: Run compiler tests
-        working-directory: toolchain/compiler
-        run: |
-          cargo nextest run --cargo-profile=fast --features parallel

Cargo.toml

@@ -71,6 +71,7 @@ afs-compiler = { path = "toolchain/compiler", default-features = false }
 stark-vm = { path = "vm", default-features = false }
 afs-primitives = { path = "circuits/primitives", default-features = false }
 ax-ecc-primitives = { path = "circuits/ecc", default-features = false }
+ax-ecc-lib = { path = "lib/ecc", default-features = false }
 afs-derive = { path = "derive", default-features = false }
 ax-sdk = { path = "sdk", default-features = false }
 afs-transpiler = { path = "toolchain/riscv/transpiler", default-features = false }

circuits/ecc/Cargo.toml

@@ -25,6 +25,7 @@ halo2curves-axiom = { git = "https://github.com/axiom-crypto/halo2curves.git" }
 afs-primitives = { path = "../primitives" }
 hex-literal = { workspace = true }
 sha3 = "0.10.8"
+lazy_static.workspace = true
 
 [features]
 default = []

circuits/ecc/src/field_expression/builder.rs

@@ -1,4 +1,4 @@
-use std::{cell::RefCell, ops::Deref, rc::Rc, sync::Arc};
+use std::{cell::RefCell, ops::Deref, rc::Rc};
 
 use afs_primitives::{
     bigint::{
@@ -7,18 +7,17 @@ use afs_primitives::{
         utils::*,
         OverflowInt,
     },
-    sub_chip::{AirConfig, LocalTraceInstructions, SubAir},
-    var_range::{bus::VariableRangeCheckerBus, VariableRangeCheckerChip},
+    var_range::{VariableRangeCheckerBus, VariableRangeCheckerChip},
+    SubAir, TraceSubRowGenerator,
 };
 use afs_stark_backend::{
     interaction::InteractionBuilder,
     rap::{BaseAirWithPublicValues, PartitionedBaseAir},
 };
 use num_bigint_dig::{BigInt, BigUint, Sign};
 use num_traits::Zero;
-use p3_air::{Air, AirBuilder, BaseAir};
+use p3_air::{AirBuilder, BaseAir};
 use p3_field::{Field, PrimeField64};
-use p3_matrix::Matrix;
 
 use super::{FieldVariable, SymbolicExpr};
 
@@ -150,28 +149,24 @@ impl<F: Field> BaseAir<F> for FieldExpr {
     }
 }
 
-impl<AB: InteractionBuilder> Air<AB> for FieldExpr {
-    fn eval(&self, builder: &mut AB) {
-        let main = builder.main();
-        let local = main.row_slice(0);
-        let local = local.to_vec();
-        SubAir::eval(self, builder, local, ());
-    }
-}
-
 impl<AB: InteractionBuilder> SubAir<AB> for FieldExpr {
-    type IoView = Vec<AB::Var>;
-    type AuxView = ();
-
-    fn eval(&self, builder: &mut AB, io: Vec<AB::Var>, _aux: ()) {
+    /// The sub-row slice owned by the expression builder.
+    type AirContext<'a> = &'a [AB::Var]
+    where AB: 'a, AB::Var: 'a, AB::Expr: 'a;
+
+    fn eval<'a>(&'a self, builder: &'a mut AB, local: &'a [AB::Var])
+    where
+        AB::Var: 'a,
+        AB::Expr: 'a,
+    {
         let FieldExprCols {
             is_valid,
             inputs,
             vars,
             q_limbs,
             carry_limbs,
             flags,
-        } = self.load_vars(&io);
+        } = self.load_vars(local);
         let inputs = load_overflow::<AB>(inputs, self.limb_bits);
         let vars = load_overflow::<AB>(vars, self.limb_bits);
 
@@ -180,14 +175,16 @@ impl<AB: InteractionBuilder> SubAir<AB> for FieldExpr {
         }
         for i in 0..self.constraints.len() {
             let expr = self.constraints[i].evaluate_overflow_expr::<AB>(&inputs, &vars, &flags);
-            self.check_carry_mod_to_zero.constrain_carry_mod_to_zero(
+            self.check_carry_mod_to_zero.eval(
                 builder,
-                expr,
-                CheckCarryModToZeroCols {
-                    carries: carry_limbs[i].clone(),
-                    quotient: q_limbs[i].clone(),
-                },
-                is_valid,
+                (
+                    expr,
+                    CheckCarryModToZeroCols {
+                        carries: carry_limbs[i].clone(),
+                        quotient: q_limbs[i].clone(),
+                    },
+                    is_valid,
+                ),
             )
         }
 
@@ -217,16 +214,15 @@ pub struct FieldExprCols<T> {
     pub flags: Vec<T>,
 }
 
-impl AirConfig for FieldExpr {
-    // No column struct.
-    type Cols<T> = Vec<T>;
-}
-
-impl<F: PrimeField64> LocalTraceInstructions<F> for FieldExpr {
-    type LocalInput = (Vec<BigUint>, Arc<VariableRangeCheckerChip>, Vec<bool>);
+impl<F: PrimeField64> TraceSubRowGenerator<F> for FieldExpr {
+    type TraceContext<'a> = (&'a VariableRangeCheckerChip, Vec<BigUint>, Vec<bool>);
+    type ColsMut<'a> = &'a mut [F];
 
-    fn generate_trace_row(&self, local_input: Self::LocalInput) -> Self::Cols<F> {
-        let (inputs, range_checker, flags) = local_input;
+    fn generate_subrow<'a>(
+        &'a self,
+        (range_checker, inputs, flags): (&'a VariableRangeCheckerChip, Vec<BigUint>, Vec<bool>),
+        sub_row: &'a mut [F],
+    ) {
         assert_eq!(inputs.len(), self.num_input);
         // Remove this if this is no longer the case in the future.
         assert_eq!(self.num_variables, self.constraints.len());
@@ -309,15 +305,18 @@ impl<F: PrimeField64> LocalTraceInstructions<F> for FieldExpr {
             .map(|x| vec_isize_to_f::<F>(x.limbs.clone()))
             .collect::<Vec<_>>();
 
-        [
-            vec![F::one()],
-            input_limbs.concat(),
-            vars_limbs.concat(),
-            all_q.concat(),
-            all_carry.concat(),
-            flags.iter().map(|x| F::from_bool(*x)).collect::<Vec<_>>(),
-        ]
-        .concat()
+        // TODO: avoid all these copies and directly allocate
+        sub_row.copy_from_slice(
+            &[
+                vec![F::one()],
+                input_limbs.concat(),
+                vars_limbs.concat(),
+                all_q.concat(),
+                all_carry.concat(),
+                flags.iter().map(|x| F::from_bool(*x)).collect::<Vec<_>>(),
+            ]
+            .concat(),
+        );
     }
 }