feat: sha256 air WIP

Cargo.toml

@@ -12,6 +12,7 @@ members = [
     "vm",
     "vm/bin",
     "poseidon2-air",
+    "sha256-air",
 ]
 resolver = "2"
 

chips/src/bits/bit_decompose/air.rs

@@ -0,0 +1,56 @@
+use super::columns::BitDecomposeCols;
+use crate::{
+    sub_chip::{AirConfig, SubAir},
+    utils::Word32,
+};
+use p3_air::{Air, AirBuilder, BaseAir};
+use p3_field::{AbstractField, Field};
+use p3_matrix::Matrix;
+use std::borrow::Borrow;
+
+pub struct BitDecomposeAir<const N: usize> {
+    // TOOD: do we need this?
+    bus_index: usize,
+}
+
+impl<const N: usize> BitDecomposeAir<N> {
+    pub fn new(bus_index: usize) -> Self {
+        Self { bus_index }
+    }
+}
+
+impl<const N: usize> AirConfig for BitDecomposeAir<N> {
+    type Cols<T> = BitDecomposeCols<N, T>;
+}
+
+impl<F: Field, const N: usize> BaseAir<F> for BitDecomposeAir<N> {
+    fn width(&self) -> usize {
+        N + 1
+    }
+}
+
+impl<AB: AirBuilder, const N: usize> Air<AB> for BitDecomposeAir<N> {
+    fn eval(&self, builder: &mut AB) {
+        let main = builder.main();
+        let local = main.row_slice(0);
+        let local: &[AB::Var] = (*local).borrow();
+        let cols = BitDecomposeCols::<N, AB::Var>::from_slice(local);
+
+        SubAir::eval(self, builder, cols, ());
+    }
+}
+
+impl<const N: usize, AB: AirBuilder> SubAir<AB> for BitDecomposeAir<N> {
+    type IoView = BitDecomposeCols<N, AB::Var>;
+    type AuxView = ();
+
+    fn eval(&self, builder: &mut AB, io: Self::IoView, _aux: Self::AuxView) {
+        let mut from_bits = AB::Expr::zero();
+        for (i, &bit) in io.x_bits.iter().enumerate() {
+            from_bits += bit * AB::Expr::from_canonical_u32(1 << i);
+        }
+        // How to make this a function of Word32?
+        let from_x = io.x.0[0] * AB::Expr::from_canonical_u32(1 << 16) + io.x.0[1];
+        builder.assert_eq(from_bits, from_x);
+    }
+}

chips/src/bits/bit_decompose/columns.rs

@@ -0,0 +1,14 @@
+use crate::utils::Word32;
+
+pub struct BitDecomposeCols<const N: usize, T> {
+    pub x: Word32<T>,
+    pub x_bits: Vec<T>,
+}
+
+impl<const N: usize, T: Clone> BitDecomposeCols<N, T> {
+    pub fn from_slice(slice: &[T]) -> Self {
+        let x = Word32([slice[0].clone(), slice[1].clone()]);
+        let x_bits = slice[2..].to_vec();
+        Self { x, x_bits }
+    }
+}

chips/src/bits/bit_decompose/mod.rs

@@ -0,0 +1,2 @@
+pub mod air;
+pub mod columns;

chips/src/lib.rs

@@ -1,4 +1,7 @@
 pub mod assert_sorted;
+pub mod bits {
+    pub mod bit_decompose;
+}
 pub mod common;
 pub mod execution_air;
 pub mod group_by;

chips/src/utils.rs

@@ -1,5 +1,8 @@
+use afs_derive::AlignedBorrow;
+use p3_air::Air;
 use p3_air::{AirBuilder, VirtualPairCol};
 use p3_field::Field;
+use p3_field::PrimeField32;
 
 // TODO: Ideally upstream PrimeField implements From<T>
 pub trait FieldFrom<T> {
@@ -45,3 +48,40 @@ pub fn or<AB: AirBuilder>(a: AB::Expr, b: AB::Expr) -> AB::Expr {
 pub fn implies<AB: AirBuilder>(a: AB::Expr, b: AB::Expr) -> AB::Expr {
     or::<AB>(AB::Expr::one() - a, b)
 }
+
+#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
+pub struct Word32<T>(pub [T; 2]);
+impl<AB, T> Word32<T>
+where
+    AB: AirBuilder,
+    T: AB::Var,
+{
+    pub fn get_value(&self) -> AB::Expr {
+        let upper = AB::Expr::from(self.0[0]);
+        let lower = AB::Expr::from(self.0[1]);
+        lower + (upper * AB::Expr::from_canonical_u64(1 << 16))
+    }
+}
+
+// impl<T> Word32<T> {
+//     pub fn get_value<AB: AirBuilder<Var = T>>(&self) -> AB::Expr {
+//         let lower = AB::Expr::from(self.0[1]);
+//         let upper = AB::Expr::from(self.0[0]);
+//         lower + (upper * AB::Expr::from_canonical_u64(1 << 16))
+//         // AB::Expr::from_canonical_u16(0)
+//     }
+// }
+
+// impl<AB: AirBuilder> Word32<AB::Var> {
+//     fn concat_u8(x: u8, y: u8) -> u16 {
+//         ((x as u16) << 8) | (y as u16)
+//     }
+//     // TODO: verify from and to 32. probably wrong now. Just make it compile first.
+//     pub fn from_u32(x: u32) -> Self {
+//         let bytes = x.to_le_bytes(); // TODO: should this be le or be?
+//         Word32([
+//             T::from_canonical_u16(Self::concat_u8(bytes[0], bytes[1])),
+//             T::from_canonical_u16(Self::concat_u8(bytes[2], bytes[3])),
+//         ])
+//     }
+// }

sha256-air/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "sha256-air"
+version.workspace = true
+authors.workspace = true
+edition.workspace = true
+
+[dependencies]
+p3-air = { workspace = true }
+p3-field = { workspace = true }
+p3-matrix = { workspace = true }
+
+afs-chips = { path = "../chips" }
+afs-derive = { path = "../derive" }

sha256-air/src/air.rs

@@ -0,0 +1,131 @@
+use crate::columns::{Sha256AuxCols, Sha256IoCols, Word};
+
+use super::columns::Sha256Cols;
+use afs_chips::bits::bit_decompose::air::BitDecomposeAir;
+use afs_chips::bits::bit_decompose::columns::BitDecomposeCols;
+use afs_chips::sub_chip::{AirConfig, SubAir};
+use p3_air::{Air, AirBuilder, BaseAir, FilteredAirBuilder};
+use p3_field::{AbstractField, Field};
+use p3_matrix::Matrix;
+use std::borrow::Borrow;
+
+// TODO: update
+pub const NUM_COLUMNS: usize = 100;
+
+pub const H: [u32; 8] = [
+    0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
+];
+
+pub struct Sha256Air {
+    bus_index: usize,
+}
+
+impl AirConfig for Sha256Air {
+    type Cols<T> = Sha256Cols<T>;
+}
+
+impl<F: Field> BaseAir<F> for Sha256Air {
+    fn width(&self) -> usize {
+        NUM_COLUMNS
+    }
+}
+
+impl<AB: AirBuilder> Air<AB> for Sha256Air {
+    fn eval(&self, builder: &mut AB) {
+        let main = builder.main();
+        let (local, next) = (main.row_slice(0), main.row_slice(1));
+        let local: &Sha256Cols<AB::Var> = (*local).borrow();
+        let next: &Sha256Cols<AB::Var> = (*next).borrow();
+
+        SubAir::<AB>::eval(self, builder, local.io, (local.aux, next.aux));
+    }
+}
+
+impl<AB: AirBuilder> SubAir<AB> for Sha256Air {
+    type IoView = Sha256IoCols<AB::Var>;
+    // local, next
+    type AuxView = (Sha256AuxCols<AB::Var>, Sha256AuxCols<AB::Var>);
+
+    fn eval(&self, builder: &mut AB, io: Self::IoView, aux: Self::AuxView) {
+        self.eval_init(builder, &aux.0);
+        self.eval_compression(builder, &io, &aux.0, &aux.1);
+    }
+}
+
+impl Sha256Air {
+    fn eval_io<AB: AirBuilder>(
+        &self,
+        builder: &mut AB,
+        io: &Sha256IoCols<AB::Var>,
+        aux: &Sha256AuxCols<AB::Var>,
+    ) {
+        todo!()
+    }
+
+    fn eval_init<AB: AirBuilder>(&self, builder: &mut AB, aux: &Sha256AuxCols<AB::Var>) {
+        self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h0, H[0]);
+        self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h1, H[1]);
+        self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h2, H[2]);
+        self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h3, H[3]);
+        self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h4, H[4]);
+        self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h5, H[5]);
+        self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h6, H[6]);
+        self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h7, H[7]);
+    }
+
+    fn eval_compression<AB: AirBuilder>(
+        &self,
+        builder: &mut AB,
+        local_io: &Sha256IoCols<AB::Var>,
+        local_aux: &Sha256AuxCols<AB::Var>,
+        next_aux: &Sha256AuxCols<AB::Var>,
+    ) {
+        // Working variables set to hashes at the start of each block.
+        let mut block_start = builder.when(local_aux.is_block_start);
+        self.assert_words_equal(&mut block_start, &local_aux.a, &next_aux.h0);
+        self.assert_words_equal(&mut block_start, &local_aux.b, &next_aux.h1);
+        self.assert_words_equal(&mut block_start, &local_aux.c, &next_aux.h2);
+        self.assert_words_equal(&mut block_start, &local_aux.d, &next_aux.h3);
+        self.assert_words_equal(&mut block_start, &local_aux.e, &next_aux.h4);
+        self.assert_words_equal(&mut block_start, &local_aux.f, &next_aux.h5);
+        self.assert_words_equal(&mut block_start, &local_aux.g, &next_aux.h6);
+        self.assert_words_equal(&mut block_start, &local_aux.h, &next_aux.h7);
+
+        let mut is_main_loop = builder.when_ne(local_aux.row_idx, AB::F::zero());
+        let bit_decompose_air = BitDecomposeAir::<32>::new(self.bus_index);
+
+        let x = std::iter::once(local_aux.e).chain(local_aux.e_bits);
+        SubAir::eval(
+            &bit_decompose_air,
+            builder,
+            BitDecomposeCols::from_slice(std::iter::once(local_aux.e).chain(local_aux.e_bits)),
+            (),
+        );
+
+        // Update the working variables for next round.
+
+        // Update hashes at the end of each block.
+    }
+
+    fn assert_word_equal_u32<AB: AirBuilder>(
+        &self,
+        builder: &mut FilteredAirBuilder<'_, AB>,
+        word: &Word<AB::Var>,
+        x: u32,
+    ) {
+        let high = AB::F::from_canonical_u16((x >> 16) as u16);
+        let low = AB::F::from_canonical_u16((x & 0xffff) as u16);
+        builder.assert_eq(word.0[0], high);
+        builder.assert_eq(word.0[1], low);
+    }
+
+    fn assert_words_equal<AB: AirBuilder>(
+        &self,
+        builder: &mut FilteredAirBuilder<'_, AB>,
+        w1: &Word<AB::Var>,
+        w2: &Word<AB::Var>,
+    ) {
+        builder.assert_eq(w1.0[0], w2.0[0]);
+        builder.assert_eq(w1.0[1], w2.0[1]);
+    }
+}

sha256-air/src/columns.rs

@@ -0,0 +1,66 @@
+use afs_chips::utils::Word32;
+use afs_derive::AlignedBorrow;
+
+#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
+pub struct Sha256Cols<T> {
+    pub io: Sha256IoCols<T>,
+    pub aux: Sha256AuxCols<T>,
+}
+
+#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
+pub struct Sha256IoCols<T> {
+    // Two field elements to hold one word (32bits) of input
+    pub input: Word32<T>,
+    // 16 field elements to hold one 256bit output
+    pub output: [T; 16],
+}
+
+#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
+pub struct Sha256AuxCols<T> {
+    // The message schedule at the current round.
+    pub w: Word32<T>,
+    // Sliding window of the previous 16 w.
+    pub w_prev: [T; 32],
+    // The round constant.
+    pub k: Word32<T>,
+    // The working variables a~h.
+    pub a: Word32<T>,
+    pub b: Word32<T>,
+    pub c: Word32<T>,
+    pub d: Word32<T>,
+    pub e: Word32<T>,
+    pub f: Word32<T>,
+    pub g: Word32<T>,
+    pub h: Word32<T>,
+    // The bit operation variables to compute w.
+    pub s0: Word32<T>,
+    pub s1: Word32<T>,
+    // The variables to update working variables each round.
+    pub cs1: Word32<T>, // S1 in pseudocode
+    pub ch: Word32<T>,
+    pub temp1: Word32<T>,
+    pub cs0: Word32<T>, // S0 in pseudocode
+    pub maj: Word32<T>,
+    pub temp2: Word32<T>,
+    // TODO: use Vec? Fix for now so I can use aligned borrow and skip implementing from slice.
+    pub e_bits: [T; 32],
+    // Hash values.
+    pub h0: Word32<T>,
+    pub h1: Word32<T>,
+    pub h2: Word32<T>,
+    pub h3: Word32<T>,
+    pub h4: Word32<T>,
+    pub h5: Word32<T>,
+    pub h6: Word32<T>,
+    pub h7: Word32<T>,
+    // Control flow columns.
+    pub idx: T,
+    pub row_idx: T,
+    pub is_block_start: T, // Bool, basically whether row_idx == 0, but equality is hard to condition on.
+}
+
+impl<T> Sha256Cols<T> {
+    pub fn flatten(&self) -> Vec<T> {
+        todo!()
+    }
+}

sha256-air/src/lib.rs

@@ -0,0 +1,2 @@
+pub mod air;
+pub mod columns;