feat: Add maybe_async_trait

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Changelog
 
 ## 0.10.0 (2024-06-11) - `utils/maybe-async` crate only
+- Added `maybe-async-trait` procedural macro.
 - [BREAKING] Refactored `maybe-async` macro into simpler `maybe-async` and `maybe-await` macros.
 
 ## 0.9.1 (2024-06-24) - `utils/core` crate only

air/src/air/context.rs

@@ -309,8 +309,7 @@ impl<B: StarkField> AirContext<B> {
 
         // we use the identity: ceil(a/b) = (a + b - 1)/b
         let num_constraint_col =
-            (highest_constraint_degree - transition_divisior_degree + trace_length - 1)
-                / trace_length;
+            (highest_constraint_degree - transition_divisior_degree).div_ceil(trace_length);
 
         cmp::max(num_constraint_col, 1)
     }

air/src/proof/ood_frame.rs

@@ -229,6 +229,8 @@ impl Deserializable for OodFrame {
 // OOD FRAME TRACE STATES
 // ================================================================================================
 
+/// Trace evaluation frame at the out-of-domain point.
+///
 /// Stores the trace evaluations at `z` and `gz`, where `z` is a random Field element in
 /// `current_row` and `next_row`, respectively. If the Air contains a Lagrange kernel auxiliary
 /// column, then that column interpolated polynomial will be evaluated at `z`, `gz`, `g^2 z`, ...

crypto/src/hash/mds/mds_f64_12x12.rs

@@ -3,6 +3,20 @@
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
 
+//! This module contains helper functions as well as constants used to perform a 12x12 vector-matrix
+//! multiplication. The special form of our MDS matrix i.e. being circulant, allows us to reduce
+//! the vector-matrix multiplication to a Hadamard product of two vectors in "frequency domain".
+//! This follows from the simple fact that every circulant matrix has the columns of the discrete
+//! Fourier transform matrix as orthogonal eigenvectors.
+//! The implementation also avoids the use of 3-point FFTs, and 3-point iFFTs, and substitutes that
+//! with explicit expressions. It also avoids, due to the form of our matrix in the frequency domain,
+//! divisions by 2 and repeated modular reductions. This is because of our explicit choice of
+//! an MDS matrix that has small powers of 2 entries in frequency domain.
+//! The following implementation has benefited greatly from the discussions and insights of
+//! Hamish Ivey-Law and Jacqueline Nabaglo of Polygon Zero and is based on Nabaglo's implementation
+//! in [Plonky2](https://github.com/mir-protocol/plonky2).
+//! The circulant matrix is identified by its first row: [7, 23, 8, 26, 13, 10, 9, 7, 6, 22, 21, 8].
+
 // FFT-BASED MDS MULTIPLICATION HELPER FUNCTIONS
 // ================================================================================================
 
@@ -12,20 +26,6 @@ use math::{
     FieldElement,
 };
 
-/// This module contains helper functions as well as constants used to perform a 12x12 vector-matrix
-/// multiplication. The special form of our MDS matrix i.e. being circulant, allows us to reduce
-/// the vector-matrix multiplication to a Hadamard product of two vectors in "frequency domain".
-/// This follows from the simple fact that every circulant matrix has the columns of the discrete
-/// Fourier transform matrix as orthogonal eigenvectors.
-/// The implementation also avoids the use of 3-point FFTs, and 3-point iFFTs, and substitutes that
-/// with explicit expressions. It also avoids, due to the form of our matrix in the frequency domain,
-/// divisions by 2 and repeated modular reductions. This is because of our explicit choice of
-/// an MDS matrix that has small powers of 2 entries in frequency domain.
-/// The following implementation has benefited greatly from the discussions and insights of
-/// Hamish Ivey-Law and Jacqueline Nabaglo of Polygon Zero and is based on Nabaglo's implementation
-/// in [Plonky2](https://github.com/mir-protocol/plonky2).
-/// The circulant matrix is identified by its first row: [7, 23, 8, 26, 13, 10, 9, 7, 6, 22, 21, 8].
-
 // MDS matrix in frequency domain.
 // More precisely, this is the output of the three 4-point (real) FFTs of the first column of
 // the MDS matrix i.e. just before the multiplication with the appropriate twiddle factors

crypto/src/hash/mds/mds_f64_8x8.rs

@@ -3,6 +3,19 @@
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
 
+//! This module contains helper functions as well as constants used to perform a 8x8 vector-matrix
+//! multiplication. The special form of our MDS matrix i.e. being circulant, allows us to reduce
+//! the vector-matrix multiplication to a Hadamard product of two vectors in "frequency domain".
+//! This follows from the simple fact that every circulant matrix has the columns of the discrete
+//! Fourier transform matrix as orthogonal eigenvectors.
+//! The implementation also avoids the use of internal 2-point FFTs, and 2-point iFFTs, and substitutes
+//! them with explicit expressions. It also avoids, due to the form of our matrix in the frequency domain,
+//! divisions by 2 and repeated modular reductions. This is because of our explicit choice of
+//! an MDS matrix that has small powers of 2 entries in frequency domain.
+//! The following implementation has benefited greatly from the discussions and insights of
+//! Hamish Ivey-Law and Jacqueline Nabaglo of Polygon Zero is based on Nabaglo's implementation
+//! in [Plonky2](https://github.com/mir-protocol/plonky2).
+
 // FFT-BASED MDS MULTIPLICATION HELPER FUNCTIONS
 // ================================================================================================
 
@@ -12,20 +25,7 @@ use math::{
     FieldElement,
 };
 
-/// This module contains helper functions as well as constants used to perform a 8x8 vector-matrix
-/// multiplication. The special form of our MDS matrix i.e. being circulant, allows us to reduce
-/// the vector-matrix multiplication to a Hadamard product of two vectors in "frequency domain".
-/// This follows from the simple fact that every circulant matrix has the columns of the discrete
-/// Fourier transform matrix as orthogonal eigenvectors.
-/// The implementation also avoids the use of internal 2-point FFTs, and 2-point iFFTs, and substitutes
-/// them with explicit expressions. It also avoids, due to the form of our matrix in the frequency domain,
-/// divisions by 2 and repeated modular reductions. This is because of our explicit choice of
-/// an MDS matrix that has small powers of 2 entries in frequency domain.
-/// The following implementation has benefited greatly from the discussions and insights of
-/// Hamish Ivey-Law and Jacqueline Nabaglo of Polygon Zero is based on Nabaglo's implementation
-/// in [Plonky2](https://github.com/mir-protocol/plonky2).
 /// The circulant matrix is identified by its first row: [23, 8, 13, 10, 7, 6, 21, 8].
-
 // MDS matrix in frequency domain.
 // More precisely, this is the output of the two 4-point (real) FFTs of the first column of
 // the MDS matrix i.e. just before the multiplication with the appropriate twiddle factors

crypto/src/merkle/concurrent.rs

@@ -18,6 +18,8 @@ pub const MIN_CONCURRENT_LEAVES: usize = 1024;
 // PUBLIC FUNCTIONS
 // ================================================================================================
 
+/// Returns internal nodes of a Merkle tree constructed from the provided leaves.
+///
 /// Builds all internal nodes of the Merkle using all available threads and stores the
 /// results in a single vector such that root of the tree is at position 1, nodes immediately
 /// under the root is at positions 2 and 3 etc.

examples/src/utils/rescue.rs

@@ -21,6 +21,8 @@ pub const RATE_WIDTH: usize = 4;
 /// Two elements (32-bytes) are returned as digest.
 const DIGEST_SIZE: usize = 2;
 
+/// Number of rounds in a single permutation of the hash function.
+///
 /// The number of rounds is set to 7 to provide 128-bit security level with 40% security margin;
 /// computed using algorithm 7 from <https://eprint.iacr.org/2020/1143.pdf>
 /// security margin here differs from Rescue Prime specification which suggests 50% security

math/src/field/f64/mod.rs

@@ -5,6 +5,7 @@
 
 //! An implementation of a 64-bit STARK-friendly prime field with modulus $2^{64} - 2^{32} + 1$
 //! using Montgomery representation.
+//!
 //! Our implementation follows <https://eprint.iacr.org/2022/274.pdf> and is constant-time.
 //!
 //! This field supports very fast modular arithmetic and has a number of other attractive

utils/maybe_async/README.md

@@ -70,6 +70,58 @@ async fn world() -> String {
 }
 ```
 
+## maybe_async_trait
+
+The `maybe_async_trait` macro can be applied to traits, and it will conditionally add the `async` keyword to trait methods annotated with `#[maybe_async]`, depending on the async feature being enabled. It also applies `#[async_trait::async_trait(?Send)]` to the trait or impl block when the async feature is on.
+
+For example:
+
+```rust
+// Adding `maybe_async_trait` to a trait definition
+#[maybe_async_trait]
+trait ExampleTrait {
+    #[maybe_async]
+    fn hello_world(&self);
+
+    fn get_hello(&self) -> String;
+}
+
+// Adding `maybe_async_trait` to an implementation of the trait
+#[maybe_async_trait]
+impl ExampleTrait for MyStruct {
+    #[maybe_async]
+    fn hello_world(&self) {
+        // ...
+    }
+
+    fn get_hello(&self) -> String {
+        // ...
+    }
+}
+```
+
+When `async` is set, it gets transformed into:
+
+```rust
+#[async_trait::async_trait(?Send)]
+trait ExampleTrait {
+    async fn hello_world(&self);
+
+    fn get_hello(&self) -> String;
+}
+
+#[async_trait::async_trait(?Send)]
+impl ExampleTrait for MyStruct {
+    async fn hello_world(&self) {
+        // ...
+    }
+
+    fn get_hello(&self) -> String {
+        // ...
+    }
+}
+```
+
 ## License
 
 This project is [MIT licensed](../../LICENSE).

utils/maybe_async/src/lib.rs

@@ -5,7 +5,7 @@
 
 use proc_macro::TokenStream;
 use quote::quote;
-use syn::{parse_macro_input, Expr, ItemFn, TraitItemFn};
+use syn::{parse_macro_input, Expr, ImplItem, ItemFn, ItemImpl, ItemTrait, TraitItem, TraitItemFn};
 
 /// Parses a function (regular or trait) and conditionally adds the `async` keyword depending on
 /// the `async` feature flag being enabled.
@@ -67,6 +67,129 @@ pub fn maybe_async(_attr: TokenStream, input: TokenStream) -> TokenStream {
     }
 }
 
+/// Conditionally add `async` keyword to functions.
+///
+/// Parses a trait or an `impl` block and conditionally adds the `async` keyword to methods that
+/// are annotated with `#[maybe_async]`, depending on the `async` feature flag being enabled.
+/// Additionally, if applied to a trait definition or impl block, it will add
+/// `#[async_trait::async_trait(?Send)]` to the it.
+///
+/// For example, given the following trait definition:
+/// ```ignore
+/// #[maybe_async_trait]
+/// trait ExampleTrait {
+///     #[maybe_async]
+///     fn hello_world(&self);
+///
+///     fn get_hello(&self) -> String;
+/// }
+/// ```
+///
+/// And the following implementation:
+/// ```ignore
+/// #[maybe_async_trait]
+/// impl ExampleTrait for MyStruct {
+///     #[maybe_async]
+///     fn hello_world(&self) {
+///         // ...
+///     }
+///
+///     fn get_hello(&self) -> String {
+///         // ...
+///     }
+/// }
+/// ```
+///
+/// When the `async` feature is enabled, this will be transformed into:
+/// ```ignore
+/// #[async_trait::async_trait(?Send)]
+/// trait ExampleTrait {
+///     async fn hello_world(&self);
+///
+///     fn get_hello(&self) -> String;
+/// }
+///
+/// #[async_trait::async_trait(?Send)]
+/// impl ExampleTrait for MyStruct {
+///     async fn hello_world(&self) {
+///         // ...
+///     }
+///
+///     fn get_hello(&self) -> String {
+///         // ...
+///     }
+/// }
+/// ```
+///
+/// When the `async` feature is disabled, the code remains unchanged, and neither the `async`
+/// keyword nor the `#[async_trait::async_trait(?Send)]` attribute is applied.
+#[proc_macro_attribute]
+pub fn maybe_async_trait(_attr: TokenStream, input: TokenStream) -> TokenStream {
+    // Try parsing the input as a trait definition
+    if let Ok(mut trait_item) = syn::parse::<ItemTrait>(input.clone()) {
+        let output = if cfg!(feature = "async") {
+            for item in &mut trait_item.items {
+                if let TraitItem::Fn(method) = item {
+                    // Remove the #[maybe_async] and make method async
+                    method.attrs.retain(|attr| {
+                        if attr.path().is_ident("maybe_async") {
+                            method.sig.asyncness = Some(syn::token::Async::default());
+                            false
+                        } else {
+                            true
+                        }
+                    });
+                }
+            }
+
+            quote! {
+                #[async_trait::async_trait(?Send)]
+                #trait_item
+            }
+        } else {
+            quote! {
+                #trait_item
+            }
+        };
+
+        return output.into();
+    }
+    // Check if it is an Impl block
+    else if let Ok(mut impl_item) = syn::parse::<ItemImpl>(input.clone()) {
+        let output = if cfg!(feature = "async") {
+            for item in &mut impl_item.items {
+                if let ImplItem::Fn(method) = item {
+                    // Remove #[maybe_async] and make method async
+                    method.attrs.retain(|attr| {
+                        if attr.path().is_ident("maybe_async") {
+                            method.sig.asyncness = Some(syn::token::Async::default());
+                            false // Remove the attribute
+                        } else {
+                            true // Keep other attributes
+                        }
+                    });
+                }
+            }
+            quote! {
+                #[async_trait::async_trait(?Send)]
+                #impl_item
+            }
+        } else {
+            quote! {
+                #[cfg(not(feature = "async"))]
+                #impl_item
+            }
+        };
+
+        return output.into();
+    }
+
+    // If input is neither a trait nor an impl block, emit a compile-time error
+    quote! {
+        compile_error!("`maybe_async_trait` can only be applied to trait definitions and trait impl blocks");
+    }.into()
+}
+
 /// Parses an expression and conditionally adds the `.await` keyword at the end of it depending on
 /// the `async` feature flag being enabled.
 ///