chore: Lagrange polynomial cache refactor

.gitignore

@@ -1,5 +1,6 @@
 .claude/settings.local.json
 ./CLAUDE.local.md
+.codex
 /target*
 **/temp/*
 /keys

core/experiments/benches/algebra.rs

@@ -1,6 +1,6 @@
 use aes_prng::AesRng;
 use algebra::poly::Poly;
-use algebra::poly::lagrange_interpolation;
+use algebra::poly::lagrange_interpolation_with_polys;
 use algebra::poly::lagrange_polynomials;
 use algebra::sharing::shamir::InputOp;
 use algebra::sharing::shamir::ShamirSharings;
@@ -23,32 +23,27 @@ fn bench_lagrange_poly(c: &mut Criterion) {
         let p_str = format!("n:{num_parties}/t:{threshold}/e:{max_err}");
         assert!(num_parties >= (threshold + 1) + 2 * max_err);
 
+        let mut rng = AesRng::seed_from_u64(0);
+        let secret = LevelOne::from_u128(2345);
+        let sharing = ShamirSharings::share(&mut rng, secret, num_parties, threshold).unwrap();
+        let xs: Vec<_> = sharing
+            .shares
+            .iter()
+            .map(|s| LevelOne::from_u128(s.owner().one_based() as u128))
+            .collect();
+        let ys: Vec<_> = sharing.shares.iter().map(|s| s.value()).collect();
+
+        // Pre-compute Lagrange polynomials for the memoized benchmark
+        let lagrange_polys = lagrange_polynomials(&xs);
+
         group.bench_function(BenchmarkId::new("lagrange_mem", p_str.clone()), |b| {
-            let mut rng = AesRng::seed_from_u64(0);
-            let secret = LevelOne::from_u128(2345);
-            let sharing = ShamirSharings::share(&mut rng, secret, num_parties, threshold).unwrap();
-            let xs: Vec<_> = sharing
-                .shares
-                .iter()
-                .map(|s| LevelOne::from_u128(s.owner().one_based() as u128))
-                .collect();
-            let ys: Vec<_> = sharing.shares.iter().map(|s| s.value()).collect();
             b.iter(|| {
-                let interpolated = lagrange_interpolation(&xs, &ys);
+                let interpolated = lagrange_interpolation_with_polys(&lagrange_polys, &ys);
                 assert_eq!(interpolated.unwrap().eval(&LevelOne::from_u128(0)), secret);
             });
         });
 
         group.bench_function(BenchmarkId::new("lagrange_no_mem", p_str), |b| {
-            let mut rng = AesRng::seed_from_u64(0);
-            let secret = LevelOne::from_u128(2345);
-            let sharing = ShamirSharings::share(&mut rng, secret, num_parties, threshold).unwrap();
-            let xs: Vec<_> = sharing
-                .shares
-                .iter()
-                .map(|s| LevelOne::from_u128(s.owner().one_based() as u128))
-                .collect();
-            let ys: Vec<_> = sharing.shares.iter().map(|s| s.value()).collect();
             b.iter(|| {
                 let ls = lagrange_polynomials(&xs);
                 let mut res = Poly::zero();

core/experiments/src/choreography/server.rs

@@ -1,11 +1,12 @@
 use crate::conf::party::PartyConf;
 use algebra::{
     base_ring::{Z64, Z128},
+    galois_fields::lagrange::init_lagrange_stores,
     galois_rings::common::ResiduePoly,
     structure_traits::{Derive, ErrorCorrect, Invert, Solve, Syndrome},
 };
 use observability::telemetry::make_span;
-use std::sync::Arc;
+use std::{num::NonZero, sync::Arc};
 use threshold_execution::online::preprocessing::{
     PreprocessorFactory, create_memory_factory, create_redis_factory,
 };
@@ -34,6 +35,15 @@ where
     ResiduePoly<Z128, EXTENSION_DEGREE>: Syndrome + ErrorCorrect + Invert + Solve + Derive,
 {
     let my_role: Role = settings.protocol().host().into();
+    if let Some(peers) = settings.protocol().peers() {
+        let num_parties = peers.len() + 1;
+        // The threshold-fhe config currently assumes n = 3t + 1.
+        let threshold = (num_parties - 1) / 3;
+        init_lagrange_stores(
+            NonZero::new(num_parties).expect("num_parties must be non-zero"),
+            threshold,
+        )?;
+    }
 
     let tls_conf = settings
         .certpaths

core/service/src/bin/kms-server.rs

@@ -1,3 +1,4 @@
+use algebra::galois_fields::lagrange::init_lagrange_stores;
 use anyhow::ensure;
 use clap::Parser;
 use futures_util::future::OptionFuture;
@@ -32,7 +33,7 @@ use kms_lib::{
         },
     },
 };
-use std::{net::ToSocketAddrs, sync::Arc, thread};
+use std::{net::ToSocketAddrs, num::NonZero, sync::Arc, thread};
 use thread_handles::init_rayon_thread_pool;
 use threshold_networking::tls::AttestedVerifier;
 use tokio::net::TcpListener;
@@ -367,6 +368,19 @@ async fn main_exec() -> anyhow::Result<()> {
 
     tracing::info!("Starting KMS Server with core config: {:?}", &core_config);
 
+    // NOTE: Cache for GF16 (which we use here -- for now) is fully filled.
+    // The call below is effectively a no-op unless we work in bigger fields.
+    if let Some(threshold_config) = core_config.threshold.as_ref()
+        && let Some(peers) = threshold_config.peers.as_ref()
+        && !peers.is_empty()
+    {
+        init_lagrange_stores(
+            NonZero::new(peers.len())
+                .expect("peers.len() (i.e. number of parties) was just checked to be non-zero"),
+            threshold_config.threshold as usize,
+        )?;
+    }
+
     tracing::info!(
         "Multi-threading values: tokio::num_workers: {}, rayon_num_threads: {}, total_num_cpus: {}",
         tokio::runtime::Handle::current().metrics().num_workers(),

core/threshold-algebra/src/galois_fields/gf128.rs

@@ -1,16 +1,11 @@
-use std::collections::HashMap;
-
-use error_utils::anyhow_error_and_log;
-
-use crate::{galois_fields::LagrangeMap, poly::lagrange_polynomials};
-
 use crate::{
+    galois_fields::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::{LazyLock, RwLock};
+use std::sync::{LazyLock, OnceLock};
 
 g2p!(
     GF128,
@@ -89,32 +84,11 @@ impl RingWithExceptionalSequence for GF128 {
     }
 }
 
-static LAGRANGE_STORE: LazyLock<RwLock<LagrangeMap<GF128>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+pub(crate) static LAGRANGE_STORE: OnceLock<LagrangeMap<GF128>> = OnceLock::new();
 
 impl Field for GF128 {
-    fn memoize_lagrange(points: &[Self]) -> anyhow::Result<Vec<Poly<Self>>> {
-        if let Ok(lock_lagrange_store) = LAGRANGE_STORE.read() {
-            match lock_lagrange_store.get(points) {
-                Some(v) => Ok(v.clone()),
-                None => {
-                    drop(lock_lagrange_store);
-                    if let Ok(mut lock_lagrange_store) = LAGRANGE_STORE.write() {
-                        let lagrange_pols = lagrange_polynomials(points);
-                        lock_lagrange_store.insert(points.to_vec(), lagrange_pols.clone());
-                        Ok(lagrange_pols)
-                    } else {
-                        Err(anyhow_error_and_log(
-                            "Error writing LAGRANGE_STORE".to_string(),
-                        ))
-                    }
-                }
-            }
-        } else {
-            Err(anyhow_error_and_log(
-                "Error reading LAGRANGE_STORE".to_string(),
-            ))
-        }
+    fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
+        LAGRANGE_STORE.get()?.get(points).map(|v| v.as_slice())
     }
 }
 

core/threshold-algebra/src/galois_fields/gf16.rs

@@ -1,14 +1,11 @@
-use crate::{galois_fields::LagrangeMap, poly::lagrange_polynomials};
-use error_utils::anyhow_error_and_log;
-use std::collections::HashMap;
-
 use crate::{
+    galois_fields::{LagrangeMap, lagrange::build_lagrange_map},
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::{LazyLock, RwLock};
+use std::{num::NonZero, sync::LazyLock};
 
 g2p!(
     GF16,
@@ -87,32 +84,17 @@ impl RingWithExceptionalSequence for GF16 {
     }
 }
 
-static LAGRANGE_STORE: LazyLock<RwLock<LagrangeMap<GF16>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+/// Pre-computed Lagrange basis for all (ordered) non-empty subsets of GF16 that exclude 0.
+/// Size is \sum_{k=1}^{15} C(15, k) * k = 15 * 2^14 = 245760, which is small enough to be pre-computed and stored in memory.
+pub(crate) static LAGRANGE_STORE: LazyLock<LagrangeMap<GF16>> = LazyLock::new(|| {
+    build_lagrange_map::<GF16>(NonZero::new(15).expect("15 is non-zero"), 0).expect(
+        "Initialization of the GF16 Lagrange basis can't fail with 15 parties and threshold 0",
+    )
+});
 
 impl Field for GF16 {
-    fn memoize_lagrange(points: &[Self]) -> anyhow::Result<Vec<Poly<Self>>> {
-        if let Ok(lock_lagrange_store) = LAGRANGE_STORE.read() {
-            match lock_lagrange_store.get(points) {
-                Some(v) => Ok(v.clone()),
-                None => {
-                    drop(lock_lagrange_store);
-                    if let Ok(mut lock_lagrange_store) = LAGRANGE_STORE.write() {
-                        let lagrange_pols = lagrange_polynomials(points);
-                        lock_lagrange_store.insert(points.to_vec(), lagrange_pols.clone());
-                        Ok(lagrange_pols)
-                    } else {
-                        Err(anyhow_error_and_log(
-                            "Error writing LAGRANGE_STORE".to_string(),
-                        ))
-                    }
-                }
-            }
-        } else {
-            Err(anyhow_error_and_log(
-                "Error reading LAGRANGE_STORE".to_string(),
-            ))
-        }
+    fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
+        LAGRANGE_STORE.get(points).map(|v| v.as_slice())
     }
 }
 

core/threshold-algebra/src/galois_fields/gf256.rs

@@ -1,14 +1,11 @@
-use crate::{galois_fields::LagrangeMap, poly::lagrange_polynomials};
-use error_utils::anyhow_error_and_log;
-use std::collections::HashMap;
-
 use crate::{
+    galois_fields::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::{LazyLock, RwLock};
+use std::sync::{LazyLock, OnceLock};
 
 g2p!(
     GF256,
@@ -87,32 +84,11 @@ impl RingWithExceptionalSequence for GF256 {
     }
 }
 
-static LAGRANGE_STORE: LazyLock<RwLock<LagrangeMap<GF256>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+pub(crate) static LAGRANGE_STORE: OnceLock<LagrangeMap<GF256>> = OnceLock::new();
 
 impl Field for GF256 {
-    fn memoize_lagrange(points: &[Self]) -> anyhow::Result<Vec<Poly<Self>>> {
-        if let Ok(lock_lagrange_store) = LAGRANGE_STORE.read() {
-            match lock_lagrange_store.get(points) {
-                Some(v) => Ok(v.clone()),
-                None => {
-                    drop(lock_lagrange_store);
-                    if let Ok(mut lock_lagrange_store) = LAGRANGE_STORE.write() {
-                        let lagrange_pols = lagrange_polynomials(points);
-                        lock_lagrange_store.insert(points.to_vec(), lagrange_pols.clone());
-                        Ok(lagrange_pols)
-                    } else {
-                        Err(anyhow_error_and_log(
-                            "Error writing LAGRANGE_STORE".to_string(),
-                        ))
-                    }
-                }
-            }
-        } else {
-            Err(anyhow_error_and_log(
-                "Error reading LAGRANGE_STORE".to_string(),
-            ))
-        }
+    fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
+        LAGRANGE_STORE.get()?.get(points).map(|v| v.as_slice())
     }
 }
 

core/threshold-algebra/src/galois_fields/gf32.rs

@@ -1,14 +1,11 @@
-use std::collections::HashMap;
-
-use crate::{galois_fields::LagrangeMap, poly::lagrange_polynomials};
 use crate::{
+    galois_fields::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
-use error_utils::anyhow_error_and_log;
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::{LazyLock, RwLock};
+use std::sync::{LazyLock, OnceLock};
 
 g2p!(
     GF32,
@@ -87,32 +84,11 @@ impl RingWithExceptionalSequence for GF32 {
     }
 }
 
-static LAGRANGE_STORE: LazyLock<RwLock<LagrangeMap<GF32>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+pub(crate) static LAGRANGE_STORE: OnceLock<LagrangeMap<GF32>> = OnceLock::new();
 
 impl Field for GF32 {
-    fn memoize_lagrange(points: &[Self]) -> anyhow::Result<Vec<Poly<Self>>> {
-        if let Ok(lock_lagrange_store) = LAGRANGE_STORE.read() {
-            match lock_lagrange_store.get(points) {
-                Some(v) => Ok(v.clone()),
-                None => {
-                    drop(lock_lagrange_store);
-                    if let Ok(mut lock_lagrange_store) = LAGRANGE_STORE.write() {
-                        let lagrange_pols = lagrange_polynomials(points);
-                        lock_lagrange_store.insert(points.to_vec(), lagrange_pols.clone());
-                        Ok(lagrange_pols)
-                    } else {
-                        Err(anyhow_error_and_log(
-                            "Error writing LAGRANGE_STORE".to_string(),
-                        ))
-                    }
-                }
-            }
-        } else {
-            Err(anyhow_error_and_log(
-                "Error reading LAGRANGE_STORE".to_string(),
-            ))
-        }
+    fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
+        LAGRANGE_STORE.get()?.get(points).map(|v| v.as_slice())
     }
 }