diff --git a/src/lib.rs b/src/lib.rs index 595276b..7a16ed0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -33,6 +33,98 @@ 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. +/// +/// 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")] + 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. +/// +/// # Example +/// ```no_run +/// 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. +/// 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 builder. +/// +/// 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 +/// initialized (including via lazy initialization from first use). +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. +/// +/// # Example +/// ```no_run +/// use leansig::init_poseidon2_16; +/// use p3_koala_bear::Poseidon2KoalaBear; +/// +/// 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 builder. +/// +/// 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 +/// initialized (including via lazy initialization from first use). +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 +138,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 }) + )); + } +}