From 0a254fabf671f31b7ed9e168ee47d82b7a9523b1 Mon Sep 17 00:00:00 2001 From: Chen Kai <281165273grape@gmail.com> Date: Fri, 2 Jan 2026 20:21:17 +0800 Subject: [PATCH 1/3] feat: allow injecting Poseidon2 permutations (16/24) --- src/lib.rs | 114 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 114 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index 595276b..4cfc43f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -33,6 +33,60 @@ static POSEIDON2_24: OnceLock> = OnceLock::new(); /// A lazily-initialized, thread-safe cache for the Poseidon2 permutation with a width of 16. static POSEIDON2_16: OnceLock> = OnceLock::new(); +/// Errors returned when initializing a custom Poseidon2 permutation. +#[derive(Debug, thiserror::Error)] +pub enum Poseidon2InitError { + #[error("Poseidon2 permutation for width {width} was already initialized")] + AlreadyInitialized { width: usize }, +} + +/// Initialize the width-24 Poseidon2 permutation used by this crate. +/// +/// This must be called before the first use of the permutation (i.e. before any code paths that +/// compute message/tweak hashes). If not called, the default Plonky3 permutation is used. +pub fn init_poseidon2_24(perm: Poseidon2KoalaBear<24>) -> Result<(), Poseidon2InitError> { + POSEIDON2_24 + .set(perm) + .map_err(|_| Poseidon2InitError::AlreadyInitialized { width: 24 }) +} + +/// Initialize the width-24 Poseidon2 permutation using a constructor. +/// +/// The constructor will only be called if the permutation has not been initialized yet. +pub fn init_poseidon2_24_with(builder: B) -> Result<(), Poseidon2InitError> +where + B: FnOnce() -> Poseidon2KoalaBear<24>, +{ + if POSEIDON2_24.get().is_some() { + return Err(Poseidon2InitError::AlreadyInitialized { width: 24 }); + } + init_poseidon2_24(builder()) +} + +/// Initialize the width-16 Poseidon2 permutation used by this crate. +/// +/// This must be called before the first use of the permutation. If not called, the default +/// Plonky3 permutation is used. +pub fn init_poseidon2_16(perm: Poseidon2KoalaBear<16>) -> Result<(), Poseidon2InitError> { + POSEIDON2_16 + .set(perm) + .map_err(|_| Poseidon2InitError::AlreadyInitialized { width: 16 }) +} + +/// Initialize the width-16 Poseidon2 permutation using a constructor. +/// +/// The constructor will only be called if the permutation has not been initialized yet. +pub fn init_poseidon2_16_with(builder: B) -> Result<(), Poseidon2InitError> +where + B: FnOnce() -> Poseidon2KoalaBear<16>, +{ + if POSEIDON2_16.get().is_some() { + return Err(Poseidon2InitError::AlreadyInitialized { width: 16 }); + } + + init_poseidon2_16(builder()) +} + /// Poseidon2 permutation (width 24) pub(crate) fn poseidon2_24() -> Poseidon2KoalaBear<24> { POSEIDON2_24 @@ -46,3 +100,63 @@ pub(crate) fn poseidon2_16() -> Poseidon2KoalaBear<16> { .get_or_init(default_koalabear_poseidon2_16) .clone() } + +#[cfg(test)] +mod poseidon2_init_tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use p3_koala_bear::{default_koalabear_poseidon2_16, default_koalabear_poseidon2_24}; + + use crate::{ + Poseidon2InitError, init_poseidon2_16, init_poseidon2_16_with, init_poseidon2_24, + init_poseidon2_24_with, poseidon2_16, poseidon2_24, + }; + + #[test] + fn init_poseidon2_24_returns_already_initialized_and_does_not_call_builder() { + // Ensure the OnceLock is initialized (possibly by other tests too). + let _ = poseidon2_24(); + + let calls = AtomicUsize::new(0); + let res = init_poseidon2_24_with(|| { + calls.fetch_add(1, Ordering::SeqCst); + default_koalabear_poseidon2_24() + }); + + assert!(matches!( + res, + Err(Poseidon2InitError::AlreadyInitialized { width: 24 }) + )); + assert_eq!(calls.load(Ordering::SeqCst), 0); + + let res = init_poseidon2_24(default_koalabear_poseidon2_24()); + assert!(matches!( + res, + Err(Poseidon2InitError::AlreadyInitialized { width: 24 }) + )); + } + + #[test] + fn init_poseidon2_16_returns_already_initialized_and_does_not_call_builder() { + // Ensure the OnceLock is initialized (possibly by other tests too). + let _ = poseidon2_16(); + + let calls = AtomicUsize::new(0); + let res = init_poseidon2_16_with(|| { + calls.fetch_add(1, Ordering::SeqCst); + default_koalabear_poseidon2_16() + }); + + assert!(matches!( + res, + Err(Poseidon2InitError::AlreadyInitialized { width: 16 }) + )); + assert_eq!(calls.load(Ordering::SeqCst), 0); + + let res = init_poseidon2_16(default_koalabear_poseidon2_16()); + assert!(matches!( + res, + Err(Poseidon2InitError::AlreadyInitialized { width: 16 }) + )); + } +} From 9a70b997e44bc44e63824eee602542cc6bcf3684 Mon Sep 17 00:00:00 2001 From: Chen Kai <281165273grape@gmail.com> Date: Fri, 2 Jan 2026 20:27:42 +0800 Subject: [PATCH 2/3] docs: document Poseidon2 init APIs --- src/lib.rs | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index 4cfc43f..f5e07bb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -34,6 +34,11 @@ static POSEIDON2_24: OnceLock> = OnceLock::new(); static POSEIDON2_16: OnceLock> = OnceLock::new(); /// Errors returned when initializing a custom Poseidon2 permutation. +/// +/// This crate caches Poseidon2 permutations (width 24 and 16) using `OnceLock`. +/// Initialization is therefore a one-time operation: attempting to initialize after the +/// permutation has already been set (or lazily created by first use) returns +/// `Poseidon2InitError::AlreadyInitialized`. #[derive(Debug, thiserror::Error)] pub enum Poseidon2InitError { #[error("Poseidon2 permutation for width {width} was already initialized")] @@ -44,6 +49,18 @@ pub enum Poseidon2InitError { /// /// This must be called before the first use of the permutation (i.e. before any code paths that /// compute message/tweak hashes). If not called, the default Plonky3 permutation is used. +/// +/// # Example +/// ```no_run +/// use leansig::init_poseidon2_24_with; +/// +/// // Build your custom Poseidon2(24) permutation here. +/// // For example, from a spec-aligned constant set. +/// init_poseidon2_24_with(|| { +/// // ... construct Poseidon2KoalaBear<24> ... +/// unimplemented!() +/// }).unwrap(); +/// ``` pub fn init_poseidon2_24(perm: Poseidon2KoalaBear<24>) -> Result<(), Poseidon2InitError> { POSEIDON2_24 .set(perm) @@ -53,6 +70,10 @@ pub fn init_poseidon2_24(perm: Poseidon2KoalaBear<24>) -> Result<(), Poseidon2In /// Initialize the width-24 Poseidon2 permutation using a constructor. /// /// The constructor will only be called if the permutation has not been initialized yet. +/// +/// # Errors +/// Returns `Poseidon2InitError::AlreadyInitialized { width: 24 }` if the permutation was already +/// initialized (including via lazy initialization from first use). pub fn init_poseidon2_24_with(builder: B) -> Result<(), Poseidon2InitError> where B: FnOnce() -> Poseidon2KoalaBear<24>, @@ -67,6 +88,16 @@ where /// /// This must be called before the first use of the permutation. If not called, the default /// Plonky3 permutation is used. +/// +/// # Example +/// ```no_run +/// use leansig::init_poseidon2_16_with; +/// +/// init_poseidon2_16_with(|| { +/// // ... construct Poseidon2KoalaBear<16> ... +/// unimplemented!() +/// }).unwrap(); +/// ``` pub fn init_poseidon2_16(perm: Poseidon2KoalaBear<16>) -> Result<(), Poseidon2InitError> { POSEIDON2_16 .set(perm) @@ -76,6 +107,10 @@ pub fn init_poseidon2_16(perm: Poseidon2KoalaBear<16>) -> Result<(), Poseidon2In /// Initialize the width-16 Poseidon2 permutation using a constructor. /// /// The constructor will only be called if the permutation has not been initialized yet. +/// +/// # Errors +/// Returns `Poseidon2InitError::AlreadyInitialized { width: 16 }` if the permutation was already +/// initialized (including via lazy initialization from first use). pub fn init_poseidon2_16_with(builder: B) -> Result<(), Poseidon2InitError> where B: FnOnce() -> Poseidon2KoalaBear<16>, From 27a91f3aef3077939b2ff0cf7c9c8454beb154dc Mon Sep 17 00:00:00 2001 From: Chen Kai <281165273grape@gmail.com> Date: Fri, 2 Jan 2026 20:49:08 +0800 Subject: [PATCH 3/3] refactor: simplify init APIs, clarify single-threaded usage --- src/lib.rs | 33 ++++++++++++++++++--------------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index f5e07bb..7a16ed0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -52,24 +52,26 @@ pub enum Poseidon2InitError { /// /// # Example /// ```no_run -/// use leansig::init_poseidon2_24_with; +/// use leansig::init_poseidon2_24; +/// use p3_koala_bear::Poseidon2KoalaBear; /// /// // Build your custom Poseidon2(24) permutation here. /// // For example, from a spec-aligned constant set. -/// init_poseidon2_24_with(|| { -/// // ... construct Poseidon2KoalaBear<24> ... -/// unimplemented!() -/// }).unwrap(); +/// let perm: Poseidon2KoalaBear<24> = unimplemented!(); +/// init_poseidon2_24(perm).unwrap(); /// ``` +/// +/// For a builder-style API, see [`init_poseidon2_24_with`]. pub fn init_poseidon2_24(perm: Poseidon2KoalaBear<24>) -> Result<(), Poseidon2InitError> { POSEIDON2_24 .set(perm) .map_err(|_| Poseidon2InitError::AlreadyInitialized { width: 24 }) } -/// Initialize the width-24 Poseidon2 permutation using a constructor. +/// Initialize the width-24 Poseidon2 permutation using a builder. /// -/// The constructor will only be called if the permutation has not been initialized yet. +/// The builder is only called if the permutation has not been initialized yet. +/// This function is intended for single-threaded initialization at program startup. /// /// # Errors /// Returns `Poseidon2InitError::AlreadyInitialized { width: 24 }` if the permutation was already @@ -91,22 +93,24 @@ where /// /// # Example /// ```no_run -/// use leansig::init_poseidon2_16_with; +/// use leansig::init_poseidon2_16; +/// use p3_koala_bear::Poseidon2KoalaBear; /// -/// init_poseidon2_16_with(|| { -/// // ... construct Poseidon2KoalaBear<16> ... -/// unimplemented!() -/// }).unwrap(); +/// let perm: Poseidon2KoalaBear<16> = unimplemented!(); +/// init_poseidon2_16(perm).unwrap(); /// ``` +/// +/// For a builder-style API, see [`init_poseidon2_16_with`]. pub fn init_poseidon2_16(perm: Poseidon2KoalaBear<16>) -> Result<(), Poseidon2InitError> { POSEIDON2_16 .set(perm) .map_err(|_| Poseidon2InitError::AlreadyInitialized { width: 16 }) } -/// Initialize the width-16 Poseidon2 permutation using a constructor. +/// Initialize the width-16 Poseidon2 permutation using a builder. /// -/// The constructor will only be called if the permutation has not been initialized yet. +/// The builder is only called if the permutation has not been initialized yet. +/// This function is intended for single-threaded initialization at program startup. /// /// # Errors /// Returns `Poseidon2InitError::AlreadyInitialized { width: 16 }` if the permutation was already @@ -118,7 +122,6 @@ where if POSEIDON2_16.get().is_some() { return Err(Poseidon2InitError::AlreadyInitialized { width: 16 }); } - init_poseidon2_16(builder()) }