feat: poseidon2 hash

Cargo.toml

@@ -30,26 +30,26 @@ debug = true
 debug-assertions = true
 
 [workspace.dependencies]
-p3-air = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-field = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-commit = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-matrix = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
+p3-air = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-field = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-commit = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-matrix = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
 p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3.git", features = [
   "nightly-features",
-], rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-util = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-challenger = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-dft = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-fri = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-keccak = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-keccak-air = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-blake3 = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-merkle-tree = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-poseidon2 = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-symmetric = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-uni-stark = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-maybe-rayon = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
-p3-bn254-fr = { git = "https://github.com/Plonky3/Plonky3.git", rev = "42d52e8608c1d12d337cfc8bfd692777ef13532f" }
+], rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-util = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-challenger = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-dft = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-fri = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-keccak = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-keccak-air = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-blake3 = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-merkle-tree = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-poseidon2 = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-symmetric = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-uni-stark = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-maybe-rayon = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
+p3-bn254-fr = { git = "https://github.com/Plonky3/Plonky3.git", rev = "b447924a508a780b997d451f9864e593efce0843" }
 
 # For local development. 
 

core/src/air/builder.rs

@@ -52,6 +52,11 @@ pub trait BaseAirBuilder: AirBuilder + MessageBuilder<AirInteraction<Self::Expr>
         }
     }
 
+    /// Asserts that an iterator of expressions are all zero.
+    fn assert_all_zero<I: Into<Self::Expr>>(&mut self, iter: impl IntoIterator<Item = I>) {
+        iter.into_iter().for_each(|expr| self.assert_zero(expr));
+    }
+
     /// Will return `a` if `condition` is 1, else `b`.  This assumes that `condition` is already
     /// checked to be a boolean.
     #[inline]

derive/src/lib.rs

@@ -32,6 +32,7 @@ use syn::Data;
 use syn::DeriveInput;
 use syn::GenericParam;
 use syn::ItemFn;
+use syn::WherePredicate;
 
 #[proc_macro_derive(AlignedBorrow)]
 pub fn aligned_borrow_derive(input: TokenStream) -> TokenStream {
@@ -94,7 +95,13 @@ pub fn aligned_borrow_derive(input: TokenStream) -> TokenStream {
 
 #[proc_macro_derive(
     MachineAir,
-    attributes(sp1_core_path, execution_record_path, program_path, builder_path)
+    attributes(
+        sp1_core_path,
+        execution_record_path,
+        program_path,
+        builder_path,
+        eval_trait_bound
+    )
 )]
 pub fn machine_air_derive(input: TokenStream) -> TokenStream {
     let ast: syn::DeriveInput = syn::parse(input).unwrap();
@@ -105,6 +112,7 @@ pub fn machine_air_derive(input: TokenStream) -> TokenStream {
     let execution_record_path = find_execution_record_path(&ast.attrs);
     let program_path = find_program_path(&ast.attrs);
     let builder_path = find_builder_path(&ast.attrs);
+    let eval_trait_bound = find_eval_trait_bound(&ast.attrs);
     let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
 
     match &ast.data {
@@ -257,6 +265,13 @@ pub fn machine_air_derive(input: TokenStream) -> TokenStream {
 
             let (air_impl_generics, _, _) = new_generics.split_for_impl();
 
+            let mut new_generics = generics.clone();
+            let where_clause = new_generics.make_where_clause();
+            if eval_trait_bound.is_some() {
+                let predicate: WherePredicate = syn::parse_str(&eval_trait_bound.unwrap()).unwrap();
+                where_clause.predicates.push(predicate);
+            }
+
             let air = quote! {
                 impl #air_impl_generics p3_air::Air<AB> for #name #ty_generics #where_clause {
                     fn eval(&self, builder: &mut AB) {
@@ -360,3 +375,17 @@ fn find_builder_path(attrs: &[syn::Attribute]) -> syn::Path {
     }
     parse_quote!(crate::air::SP1AirBuilder<F = F>)
 }
+
+fn find_eval_trait_bound(attrs: &[syn::Attribute]) -> Option<String> {
+    for attr in attrs {
+        if attr.path.is_ident("eval_trait_bound") {
+            if let Ok(syn::Meta::NameValue(meta)) = attr.parse_meta() {
+                if let syn::Lit::Str(lit_str) = &meta.lit {
+                    return Some(lit_str.value());
+                }
+            }
+        }
+    }
+
+    None
+}

prover/src/lib.rs

@@ -74,7 +74,7 @@ pub type OuterSC = BabyBearPoseidon2Outer;
 
 const REDUCE_DEGREE: usize = 3;
 const COMPRESS_DEGREE: usize = 9;
-const WRAP_DEGREE: usize = 9;
+const WRAP_DEGREE: usize = 17;
 
 pub type ReduceAir<F> = RecursionAir<F, REDUCE_DEGREE>;
 pub type CompressAir<F> = RecursionAir<F, COMPRESS_DEGREE>;

recursion/circuit/src/stark.rs

@@ -23,7 +23,7 @@ use sp1_recursion_compiler::ir::{Usize, Witness};
 use sp1_recursion_compiler::prelude::SymbolicVar;
 use sp1_recursion_core::air::{RecursionPublicValues, NUM_PV_ELMS_TO_HASH};
 use sp1_recursion_core::stark::config::{outer_fri_config, BabyBearPoseidon2Outer};
-use sp1_recursion_core::stark::RecursionAirSkinnyDeg9;
+use sp1_recursion_core::stark::RecursionAirWideDeg17;
 use sp1_recursion_program::commit::PolynomialSpaceVariable;
 use sp1_recursion_program::stark::RecursiveVerifierConstraintFolder;
 use sp1_recursion_program::types::QuotientDataValues;
@@ -244,7 +244,7 @@ pub fn build_wrap_circuit(
     template_proof: ShardProof<OuterSC>,
 ) -> Vec<Constraint> {
     let outer_config = OuterSC::new();
-    let outer_machine = RecursionAirSkinnyDeg9::<OuterF>::wrap_machine(outer_config);
+    let outer_machine = RecursionAirWideDeg17::<OuterF>::wrap_machine(outer_config);
 
     let mut builder = Builder::<OuterConfig>::default();
     let mut challenger = MultiField32ChallengerVariable::new(&mut builder);

recursion/compiler/src/asm/compiler.rs

@@ -517,7 +517,32 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
                         _ => unimplemented!(),
                     }
                 }
-
+                DslIr::Poseidon2AbsorbBabyBear(p2_hash_num, input) => match input {
+                    Array::Dyn(input, input_size) => {
+                        if let Usize::Var(input_size) = input_size {
+                            self.push(
+                                AsmInstruction::Poseidon2Absorb(
+                                    p2_hash_num.fp(),
+                                    input.fp(),
+                                    input_size.fp(),
+                                ),
+                                trace,
+                            );
+                        } else {
+                            unimplemented!();
+                        }
+                    }
+                    _ => unimplemented!(),
+                },
+                DslIr::Poseidon2FinalizeBabyBear(p2_hash_num, output) => match output {
+                    Array::Dyn(output, _) => {
+                        self.push(
+                            AsmInstruction::Poseidon2Finalize(p2_hash_num.fp(), output.fp()),
+                            trace,
+                        );
+                    }
+                    _ => unimplemented!(),
+                },
                 DslIr::Commit(val, index) => {
                     self.push(AsmInstruction::Commit(val.fp(), index.fp()), trace);
                 }

recursion/compiler/src/asm/instruction.rs

@@ -147,8 +147,16 @@ pub enum AsmInstruction<F, EF> {
 
     /// Perform a permutation of the Poseidon2 hash function on the array specified by the ptr.
     Poseidon2Permute(i32, i32),
+
+    /// Perform a Poseidon2 compress.
     Poseidon2Compress(i32, i32, i32),
 
+    /// Performs a Posedion2 absorb.
+    Poseidon2Absorb(i32, i32, i32),
+
+    /// Performs a Poseidon2 finalize.
+    Poseidon2Finalize(i32, i32),
+
     /// Print a variable.
     PrintV(i32),
 
@@ -846,6 +854,28 @@ impl<F: PrimeField32, EF: ExtensionField<F>> AsmInstruction<F, EF> {
                 false,
                 "".to_string(),
             ),
+            AsmInstruction::Poseidon2Absorb(hash_num, input_ptr, input_len) => Instruction::new(
+                Opcode::Poseidon2Absorb,
+                i32_f(hash_num),
+                i32_f_arr(input_ptr),
+                i32_f_arr(input_len),
+                F::zero(),
+                F::zero(),
+                false,
+                false,
+                "".to_string(),
+            ),
+            AsmInstruction::Poseidon2Finalize(hash_num, output_ptr) => Instruction::new(
+                Opcode::Poseidon2Finalize,
+                i32_f(hash_num),
+                i32_f_arr(output_ptr),
+                f_u32(F::zero()),
+                F::zero(),
+                F::zero(),
+                false,
+                false,
+                "".to_string(),
+            ),
             AsmInstruction::Commit(val, index) => Instruction::new(
                 Opcode::Commit,
                 i32_f(val),
@@ -1144,6 +1174,16 @@ impl<F: PrimeField32, EF: ExtensionField<F>> AsmInstruction<F, EF> {
                     result, src1, src2
                 )
             }
+            AsmInstruction::Poseidon2Absorb(hash_num, input_ptr, input_len) => {
+                write!(
+                    f,
+                    "poseidon2_absorb ({})fp, {})fp, ({})fp",
+                    hash_num, input_ptr, input_len,
+                )
+            }
+            AsmInstruction::Poseidon2Finalize(hash_num, output_ptr) => {
+                write!(f, "poseidon2_finalize ({})fp, {})fp", hash_num, output_ptr,)
+            }
             AsmInstruction::Commit(val, index) => {
                 write!(f, "commit ({})fp ({})fp", val, index)
             }

recursion/compiler/src/ir/builder.rs

@@ -90,7 +90,7 @@ impl<T> IntoIterator for TracedVec<T> {
 /// A builder for the DSL.
 ///
 /// Can compile to both assembly and a set of constraints.
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone)]
 pub struct Builder<C: Config> {
     pub(crate) felt_count: u32,
     pub(crate) ext_count: u32,
@@ -100,17 +100,43 @@ pub struct Builder<C: Config> {
     pub(crate) witness_var_count: u32,
     pub(crate) witness_felt_count: u32,
     pub(crate) witness_ext_count: u32,
+    pub(crate) p2_hash_num: Var<C::N>,
     pub(crate) debug: bool,
     pub(crate) is_sub_builder: bool,
 }
 
+impl<C: Config> Default for Builder<C> {
+    fn default() -> Self {
+        // We need to create a temporary placeholder for the p2_hash_num variable.
+        let placeholder_p2_hash_num = Var::new(0);
+
+        let mut new_builder = Self {
+            felt_count: 0,
+            ext_count: 0,
+            var_count: 0,
+            witness_var_count: 0,
+            witness_felt_count: 0,
+            witness_ext_count: 0,
+            operations: Default::default(),
+            nb_public_values: None,
+            p2_hash_num: placeholder_p2_hash_num,
+            debug: false,
+            is_sub_builder: false,
+        };
+
+        new_builder.p2_hash_num = new_builder.uninit();
+        new_builder
+    }
+}
+
 impl<C: Config> Builder<C> {
     /// Creates a new builder with a given number of counts for each type.
     pub fn new_sub_builder(
         var_count: u32,
         felt_count: u32,
         ext_count: u32,
         nb_public_values: Option<Var<C::N>>,
+        p2_hash_num: Var<C::N>,
         debug: bool,
     ) -> Self {
         Self {
@@ -124,6 +150,7 @@ impl<C: Config> Builder<C> {
             witness_ext_count: 0,
             operations: Default::default(),
             nb_public_values,
+            p2_hash_num,
             debug,
             is_sub_builder: true,
         }
@@ -517,9 +544,12 @@ impl<'a, C: Config> IfBuilder<'a, C> {
             self.builder.felt_count,
             self.builder.ext_count,
             self.builder.nb_public_values,
+            self.builder.p2_hash_num,
             self.builder.debug,
         );
         f(&mut f_builder);
+        self.builder.p2_hash_num = f_builder.p2_hash_num;
+
         let then_instructions = f_builder.operations;
 
         // Dispatch instructions to the correct conditional block.
@@ -565,21 +595,27 @@ impl<'a, C: Config> IfBuilder<'a, C> {
             self.builder.felt_count,
             self.builder.ext_count,
             self.builder.nb_public_values,
+            self.builder.p2_hash_num,
             self.builder.debug,
         );
 
         // Execute the `then` and `else_then` blocks and collect the instructions.
         then_f(&mut then_builder);
+        self.builder.p2_hash_num = then_builder.p2_hash_num;
+
         let then_instructions = then_builder.operations;
 
         let mut else_builder = Builder::<C>::new_sub_builder(
             self.builder.var_count,
             self.builder.felt_count,
             self.builder.ext_count,
             self.builder.nb_public_values,
+            self.builder.p2_hash_num,
             self.builder.debug,
         );
         else_f(&mut else_builder);
+        self.builder.p2_hash_num = else_builder.p2_hash_num;
+
         let else_instructions = else_builder.operations;
 
         // Dispatch instructions to the correct conditional block.
@@ -711,10 +747,12 @@ impl<'a, C: Config> RangeBuilder<'a, C> {
             self.builder.felt_count,
             self.builder.ext_count,
             self.builder.nb_public_values,
+            self.builder.p2_hash_num,
             self.builder.debug,
         );
 
         f(loop_variable, &mut loop_body_builder);
+        self.builder.p2_hash_num = loop_body_builder.p2_hash_num;
 
         let loop_instructions = loop_body_builder.operations;
 

recursion/compiler/src/ir/instructions.rs

@@ -198,6 +198,10 @@ pub enum DslIr<C: Config> {
         Array<C, Felt<C::F>>,
         Array<C, Felt<C::F>>,
     ),
+    /// Absorb an array of baby bear elements for a specified hash instance.
+    Poseidon2AbsorbBabyBear(Var<C::N>, Array<C, Felt<C::F>>),
+    /// Finalize and return the hash digest of a specified hash instance.
+    Poseidon2FinalizeBabyBear(Var<C::N>, Array<C, Felt<C::F>>),
     /// Permutes an array of Bn254 elements using Poseidon2 (output = p2_permute(array)). Should only
     /// be used when target is a gnark circuit.
     CircuitPoseidon2Permute([Var<C::N>; 3]),