From 4d59c897885e4f35e413a10844f1f4b8e4a3b1ce Mon Sep 17 00:00:00 2001
From: Simon Eudeline <simon.eudeline@zama.ai>
Date: Mon, 13 Apr 2026 14:24:16 +0200
Subject: [PATCH 01/11] chore: use once lock for lagrange store

---
 .../src/galois_fields/gf128.rs                | 47 +++++++------------
 .../src/galois_fields/gf16.rs                 | 45 +++++++-----------
 .../src/galois_fields/gf256.rs                | 45 +++++++-----------
 .../src/galois_fields/gf32.rs                 | 45 +++++++-----------
 .../src/galois_fields/gf64.rs                 | 47 +++++++------------
 .../src/galois_fields/gf8.rs                  | 46 +++++++-----------
 .../src/galois_fields/mod.rs                  |  6 ---
 core/threshold-algebra/src/poly.rs            | 14 ++++--
 .../threshold-algebra/src/structure_traits.rs |  2 +-
 .../threshold-experimental/benches/algebra.rs | 35 ++++++--------
 .../src/algebra/levels.rs                     | 40 ++++++----------
 11 files changed, 143 insertions(+), 229 deletions(-)

diff --git a/core/threshold-algebra/src/galois_fields/gf128.rs b/core/threshold-algebra/src/galois_fields/gf128.rs
index 47019e05f..06536320e 100644
--- a/core/threshold-algebra/src/galois_fields/gf128.rs
+++ b/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::poly::lagrange_polynomials;
 use crate::{
     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;
 
 g2p!(
     GF128,
@@ -89,31 +84,23 @@ impl RingWithExceptionalSequence for GF128 {
     }
 }
 
-static LAGRANGE_STORE: LazyLock<RwLock<LagrangeMap<GF128>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+#[cfg(not(test))]
+use std::sync::OnceLock;
+
+#[cfg(not(test))]
+static LAGRANGE_STORE: OnceLock<Vec<Poly<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 memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
+        #[cfg(test)]
+        {
+            lagrange_polynomials(points)
+        }
+        #[cfg(not(test))]
+        {
+            LAGRANGE_STORE
+                .get_or_init(|| lagrange_polynomials(points))
+                .clone()
         }
     }
 }
diff --git a/core/threshold-algebra/src/galois_fields/gf16.rs b/core/threshold-algebra/src/galois_fields/gf16.rs
index 823e14e78..de7e17a79 100644
--- a/core/threshold-algebra/src/galois_fields/gf16.rs
+++ b/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::poly::lagrange_polynomials;
 use crate::{
     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;
 
 g2p!(
     GF16,
@@ -87,31 +84,23 @@ impl RingWithExceptionalSequence for GF16 {
     }
 }
 
-static LAGRANGE_STORE: LazyLock<RwLock<LagrangeMap<GF16>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+#[cfg(not(test))]
+use std::sync::OnceLock;
+
+#[cfg(not(test))]
+static LAGRANGE_STORE: OnceLock<Vec<Poly<GF16>>> = OnceLock::new();
 
 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 memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
+        #[cfg(test)]
+        {
+            lagrange_polynomials(points)
+        }
+        #[cfg(not(test))]
+        {
+            LAGRANGE_STORE
+                .get_or_init(|| lagrange_polynomials(points))
+                .clone()
         }
     }
 }
diff --git a/core/threshold-algebra/src/galois_fields/gf256.rs b/core/threshold-algebra/src/galois_fields/gf256.rs
index d06070ee4..ffca08811 100644
--- a/core/threshold-algebra/src/galois_fields/gf256.rs
+++ b/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::poly::lagrange_polynomials;
 use crate::{
     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;
 
 g2p!(
     GF256,
@@ -87,31 +84,23 @@ impl RingWithExceptionalSequence for GF256 {
     }
 }
 
-static LAGRANGE_STORE: LazyLock<RwLock<LagrangeMap<GF256>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+#[cfg(not(test))]
+use std::sync::OnceLock;
+
+#[cfg(not(test))]
+static LAGRANGE_STORE: OnceLock<Vec<Poly<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 memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
+        #[cfg(test)]
+        {
+            lagrange_polynomials(points)
+        }
+        #[cfg(not(test))]
+        {
+            LAGRANGE_STORE
+                .get_or_init(|| lagrange_polynomials(points))
+                .clone()
         }
     }
 }
diff --git a/core/threshold-algebra/src/galois_fields/gf32.rs b/core/threshold-algebra/src/galois_fields/gf32.rs
index 42ee806af..a957bc331 100644
--- a/core/threshold-algebra/src/galois_fields/gf32.rs
+++ b/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::poly::lagrange_polynomials;
 use crate::{
     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;
 
 g2p!(
     GF32,
@@ -87,31 +84,23 @@ impl RingWithExceptionalSequence for GF32 {
     }
 }
 
-static LAGRANGE_STORE: LazyLock<RwLock<LagrangeMap<GF32>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+#[cfg(not(test))]
+use std::sync::OnceLock;
+
+#[cfg(not(test))]
+static LAGRANGE_STORE: OnceLock<Vec<Poly<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 memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
+        #[cfg(test)]
+        {
+            lagrange_polynomials(points)
+        }
+        #[cfg(not(test))]
+        {
+            LAGRANGE_STORE
+                .get_or_init(|| lagrange_polynomials(points))
+                .clone()
         }
     }
 }
diff --git a/core/threshold-algebra/src/galois_fields/gf64.rs b/core/threshold-algebra/src/galois_fields/gf64.rs
index e1c4b46f3..22d28496b 100644
--- a/core/threshold-algebra/src/galois_fields/gf64.rs
+++ b/core/threshold-algebra/src/galois_fields/gf64.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::poly::lagrange_polynomials;
 use crate::{
     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;
 
 g2p!(
     GF64,
@@ -89,31 +84,23 @@ impl RingWithExceptionalSequence for GF64 {
     }
 }
 
-static LAGRANGE_STORE: LazyLock<RwLock<LagrangeMap<GF64>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+#[cfg(not(test))]
+use std::sync::OnceLock;
+
+#[cfg(not(test))]
+static LAGRANGE_STORE: OnceLock<Vec<Poly<GF64>>> = OnceLock::new();
 
 impl Field for GF64 {
-    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 memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
+        #[cfg(test)]
+        {
+            lagrange_polynomials(points)
+        }
+        #[cfg(not(test))]
+        {
+            LAGRANGE_STORE
+                .get_or_init(|| lagrange_polynomials(points))
+                .clone()
         }
     }
 }
diff --git a/core/threshold-algebra/src/galois_fields/gf8.rs b/core/threshold-algebra/src/galois_fields/gf8.rs
index df0ca2969..12f1630cc 100644
--- a/core/threshold-algebra/src/galois_fields/gf8.rs
+++ b/core/threshold-algebra/src/galois_fields/gf8.rs
@@ -1,16 +1,12 @@
-use std::collections::HashMap;
-
+use crate::poly::lagrange_polynomials;
 use crate::structure_traits::RingWithExceptionalSequence;
-use crate::{galois_fields::LagrangeMap, poly::lagrange_polynomials};
-use error_utils::anyhow_error_and_log;
-
 use crate::{
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, Sample, Zero},
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::{LazyLock, RwLock};
+use std::sync::LazyLock;
 
 g2p!(
     GF8,
@@ -89,31 +85,23 @@ impl RingWithExceptionalSequence for GF8 {
     }
 }
 
-static LAGRANGE_STORE: LazyLock<RwLock<LagrangeMap<GF8>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+#[cfg(not(test))]
+use std::sync::OnceLock;
+
+#[cfg(not(test))]
+static LAGRANGE_STORE: OnceLock<Vec<Poly<GF8>>> = OnceLock::new();
 
 impl Field for GF8 {
-    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 memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
+        #[cfg(test)]
+        {
+            lagrange_polynomials(points)
+        }
+        #[cfg(not(test))]
+        {
+            LAGRANGE_STORE
+                .get_or_init(|| lagrange_polynomials(points))
+                .clone()
         }
     }
 }
diff --git a/core/threshold-algebra/src/galois_fields/mod.rs b/core/threshold-algebra/src/galois_fields/mod.rs
index 2a80f019d..bfc0fd8f8 100644
--- a/core/threshold-algebra/src/galois_fields/mod.rs
+++ b/core/threshold-algebra/src/galois_fields/mod.rs
@@ -1,9 +1,3 @@
-use crate::poly::Poly;
-use std::collections::HashMap;
-
-/// Map from evaluation points to their precomputed Lagrange polynomials.
-pub(crate) type LagrangeMap<F> = HashMap<Vec<F>, Vec<Poly<F>>>;
-
 pub mod common;
 #[cfg(feature = "extension_degree_7")]
 pub mod gf128;
diff --git a/core/threshold-algebra/src/poly.rs b/core/threshold-algebra/src/poly.rs
index e041cec9d..81bcf61f1 100644
--- a/core/threshold-algebra/src/poly.rs
+++ b/core/threshold-algebra/src/poly.rs
@@ -600,15 +600,23 @@ pub fn lagrange_polynomials<F: Field>(points: &[F]) -> Vec<Poly<F>> {
 
 /// interpolate a polynomial through coordinates where points holds the x-coordinates and values holds the y-coordinates
 pub fn lagrange_interpolation<F: Field>(points: &[F], values: &[F]) -> anyhow::Result<Poly<F>> {
-    let ls = F::memoize_lagrange(points)?;
-    if ls.len() != values.len() {
+    let lagrange_polys = F::memoize_lagrange(points);
+    lagrange_interpolation_with_polys(&lagrange_polys, values)
+}
+
+/// interpolate a polynomial using pre-computed Lagrange basis polynomials and y-coordinates
+pub fn lagrange_interpolation_with_polys<F: Field>(
+    lagrange_polys: &[Poly<F>],
+    values: &[F],
+) -> anyhow::Result<Poly<F>> {
+    if lagrange_polys.len() != values.len() {
         return Err(anyhow_error_and_log(
             "Lagrange interpolation failure: mismatch between number of points and values"
                 .to_string(),
         ));
     }
     let mut res = Poly::zero();
-    for (li, vi) in ls.into_iter().zip_eq(values.iter()) {
+    for (li, vi) in lagrange_polys.iter().zip_eq(values.iter()) {
         let term = li * vi;
         res = res + term;
     }
diff --git a/core/threshold-algebra/src/structure_traits.rs b/core/threshold-algebra/src/structure_traits.rs
index bb51b47d7..c9d36595b 100644
--- a/core/threshold-algebra/src/structure_traits.rs
+++ b/core/threshold-algebra/src/structure_traits.rs
@@ -92,7 +92,7 @@ pub trait Field
 where
     Self: RingWithExceptionalSequence + Div<Self, Output = Self> + DivAssign<Self>,
 {
-    fn memoize_lagrange(points: &[Self]) -> anyhow::Result<Vec<Poly<Self>>>;
+    fn memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>>;
 
     /// computes the multiplicative inverse of the field element
     fn invert(&self) -> Self {
diff --git a/core/threshold-experimental/benches/algebra.rs b/core/threshold-experimental/benches/algebra.rs
index 4a1fffc91..00124192e 100644
--- a/core/threshold-experimental/benches/algebra.rs
+++ b/core/threshold-experimental/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();
diff --git a/core/threshold-experimental/src/algebra/levels.rs b/core/threshold-experimental/src/algebra/levels.rs
index 25a7335eb..0116b431a 100644
--- a/core/threshold-experimental/src/algebra/levels.rs
+++ b/core/threshold-experimental/src/algebra/levels.rs
@@ -21,11 +21,10 @@ use itertools::Itertools;
 use rand::CryptoRng;
 use rand::Rng;
 use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
 use std::iter::Sum;
 use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
-use std::sync::LazyLock;
-use std::sync::RwLock;
+#[cfg(not(test))]
+use std::sync::OnceLock;
 
 use crate::algebra::crt::LevelKswCrtRepresentation;
 use crate::algebra::crt::from_crt;
@@ -349,31 +348,20 @@ macro_rules! impl_field_level {
                 }
             }
 
-            static [<LAGRANGE_STORE_BGV_ $name:upper>]: LazyLock<RwLock<HashMap<Vec<$name>, Vec<Poly<$name>>>>> =
-                LazyLock::new(|| RwLock::new(HashMap::new()));
+            #[cfg(not(test))]
+            static [<LAGRANGE_STORE_BGV_ $name:upper>]: OnceLock<Vec<Poly<$name>>> = OnceLock::new();
 
             impl Field for $name {
-                fn memoize_lagrange(points: &[Self]) -> anyhow::Result<Vec<Poly<Self>>> {
-                    if let Ok(lock_lagrange_store) = [<LAGRANGE_STORE_BGV_ $name:upper>].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_BGV_ $name:upper>].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 memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
+                    #[cfg(test)]
+                    {
+                        lagrange_polynomials(points)
+                    }
+                    #[cfg(not(test))]
+                    {
+                        [<LAGRANGE_STORE_BGV_ $name:upper>]
+                            .get_or_init(|| lagrange_polynomials(points))
+                            .clone()
                     }
                 }
 

From 370997f2f59c30377ce15a270d0581404d16c92e Mon Sep 17 00:00:00 2001
From: Simon Eudeline <simon.eudeline@zama.ai>
Date: Mon, 13 Apr 2026 14:32:51 +0200
Subject: [PATCH 02/11] chore: code cleaning

---
 .../src/galois_fields/gf128.rs                | 19 +++----------------
 .../src/galois_fields/gf16.rs                 | 19 +++----------------
 .../src/galois_fields/gf256.rs                | 19 +++----------------
 .../src/galois_fields/gf32.rs                 | 19 +++----------------
 .../src/galois_fields/gf64.rs                 | 19 +++----------------
 .../src/galois_fields/gf8.rs                  | 19 +++----------------
 core/threshold-algebra/src/poly.rs            | 11 +++++++++--
 .../threshold-algebra/src/structure_traits.rs |  2 +-
 .../src/algebra/levels.rs                     | 16 +++-------------
 9 files changed, 31 insertions(+), 112 deletions(-)

diff --git a/core/threshold-algebra/src/galois_fields/gf128.rs b/core/threshold-algebra/src/galois_fields/gf128.rs
index 06536320e..09a130e43 100644
--- a/core/threshold-algebra/src/galois_fields/gf128.rs
+++ b/core/threshold-algebra/src/galois_fields/gf128.rs
@@ -5,7 +5,7 @@ use crate::{
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::LazyLock;
+use std::sync::{LazyLock, OnceLock};
 
 g2p!(
     GF128,
@@ -84,24 +84,11 @@ impl RingWithExceptionalSequence for GF128 {
     }
 }
 
-#[cfg(not(test))]
-use std::sync::OnceLock;
-
-#[cfg(not(test))]
 static LAGRANGE_STORE: OnceLock<Vec<Poly<GF128>>> = OnceLock::new();
 
 impl Field for GF128 {
-    fn memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
-        #[cfg(test)]
-        {
-            lagrange_polynomials(points)
-        }
-        #[cfg(not(test))]
-        {
-            LAGRANGE_STORE
-                .get_or_init(|| lagrange_polynomials(points))
-                .clone()
-        }
+    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
+        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/gf16.rs b/core/threshold-algebra/src/galois_fields/gf16.rs
index de7e17a79..dd1864e2a 100644
--- a/core/threshold-algebra/src/galois_fields/gf16.rs
+++ b/core/threshold-algebra/src/galois_fields/gf16.rs
@@ -5,7 +5,7 @@ use crate::{
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::LazyLock;
+use std::sync::{LazyLock, OnceLock};
 
 g2p!(
     GF16,
@@ -84,24 +84,11 @@ impl RingWithExceptionalSequence for GF16 {
     }
 }
 
-#[cfg(not(test))]
-use std::sync::OnceLock;
-
-#[cfg(not(test))]
 static LAGRANGE_STORE: OnceLock<Vec<Poly<GF16>>> = OnceLock::new();
 
 impl Field for GF16 {
-    fn memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
-        #[cfg(test)]
-        {
-            lagrange_polynomials(points)
-        }
-        #[cfg(not(test))]
-        {
-            LAGRANGE_STORE
-                .get_or_init(|| lagrange_polynomials(points))
-                .clone()
-        }
+    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
+        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/gf256.rs b/core/threshold-algebra/src/galois_fields/gf256.rs
index ffca08811..585f78ca2 100644
--- a/core/threshold-algebra/src/galois_fields/gf256.rs
+++ b/core/threshold-algebra/src/galois_fields/gf256.rs
@@ -5,7 +5,7 @@ use crate::{
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::LazyLock;
+use std::sync::{LazyLock, OnceLock};
 
 g2p!(
     GF256,
@@ -84,24 +84,11 @@ impl RingWithExceptionalSequence for GF256 {
     }
 }
 
-#[cfg(not(test))]
-use std::sync::OnceLock;
-
-#[cfg(not(test))]
 static LAGRANGE_STORE: OnceLock<Vec<Poly<GF256>>> = OnceLock::new();
 
 impl Field for GF256 {
-    fn memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
-        #[cfg(test)]
-        {
-            lagrange_polynomials(points)
-        }
-        #[cfg(not(test))]
-        {
-            LAGRANGE_STORE
-                .get_or_init(|| lagrange_polynomials(points))
-                .clone()
-        }
+    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
+        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/gf32.rs b/core/threshold-algebra/src/galois_fields/gf32.rs
index a957bc331..430d5adfe 100644
--- a/core/threshold-algebra/src/galois_fields/gf32.rs
+++ b/core/threshold-algebra/src/galois_fields/gf32.rs
@@ -5,7 +5,7 @@ use crate::{
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::LazyLock;
+use std::sync::{LazyLock, OnceLock};
 
 g2p!(
     GF32,
@@ -84,24 +84,11 @@ impl RingWithExceptionalSequence for GF32 {
     }
 }
 
-#[cfg(not(test))]
-use std::sync::OnceLock;
-
-#[cfg(not(test))]
 static LAGRANGE_STORE: OnceLock<Vec<Poly<GF32>>> = OnceLock::new();
 
 impl Field for GF32 {
-    fn memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
-        #[cfg(test)]
-        {
-            lagrange_polynomials(points)
-        }
-        #[cfg(not(test))]
-        {
-            LAGRANGE_STORE
-                .get_or_init(|| lagrange_polynomials(points))
-                .clone()
-        }
+    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
+        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/gf64.rs b/core/threshold-algebra/src/galois_fields/gf64.rs
index 22d28496b..a0817e471 100644
--- a/core/threshold-algebra/src/galois_fields/gf64.rs
+++ b/core/threshold-algebra/src/galois_fields/gf64.rs
@@ -5,7 +5,7 @@ use crate::{
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::LazyLock;
+use std::sync::{LazyLock, OnceLock};
 
 g2p!(
     GF64,
@@ -84,24 +84,11 @@ impl RingWithExceptionalSequence for GF64 {
     }
 }
 
-#[cfg(not(test))]
-use std::sync::OnceLock;
-
-#[cfg(not(test))]
 static LAGRANGE_STORE: OnceLock<Vec<Poly<GF64>>> = OnceLock::new();
 
 impl Field for GF64 {
-    fn memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
-        #[cfg(test)]
-        {
-            lagrange_polynomials(points)
-        }
-        #[cfg(not(test))]
-        {
-            LAGRANGE_STORE
-                .get_or_init(|| lagrange_polynomials(points))
-                .clone()
-        }
+    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
+        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/gf8.rs b/core/threshold-algebra/src/galois_fields/gf8.rs
index 12f1630cc..28c1c9254 100644
--- a/core/threshold-algebra/src/galois_fields/gf8.rs
+++ b/core/threshold-algebra/src/galois_fields/gf8.rs
@@ -6,7 +6,7 @@ use crate::{
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::LazyLock;
+use std::sync::{LazyLock, OnceLock};
 
 g2p!(
     GF8,
@@ -85,24 +85,11 @@ impl RingWithExceptionalSequence for GF8 {
     }
 }
 
-#[cfg(not(test))]
-use std::sync::OnceLock;
-
-#[cfg(not(test))]
 static LAGRANGE_STORE: OnceLock<Vec<Poly<GF8>>> = OnceLock::new();
 
 impl Field for GF8 {
-    fn memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
-        #[cfg(test)]
-        {
-            lagrange_polynomials(points)
-        }
-        #[cfg(not(test))]
-        {
-            LAGRANGE_STORE
-                .get_or_init(|| lagrange_polynomials(points))
-                .clone()
-        }
+    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
+        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
     }
 }
 
diff --git a/core/threshold-algebra/src/poly.rs b/core/threshold-algebra/src/poly.rs
index 81bcf61f1..98a1e458b 100644
--- a/core/threshold-algebra/src/poly.rs
+++ b/core/threshold-algebra/src/poly.rs
@@ -600,15 +600,22 @@ pub fn lagrange_polynomials<F: Field>(points: &[F]) -> Vec<Poly<F>> {
 
 /// interpolate a polynomial through coordinates where points holds the x-coordinates and values holds the y-coordinates
 pub fn lagrange_interpolation<F: Field>(points: &[F], values: &[F]) -> anyhow::Result<Poly<F>> {
+    // In tests, bypass memoization to allow different point sets per test.
+    // memoize_lagrange uses OnceLock which caches a single point set per field type.
+    #[cfg(test)]
+    let lagrange_polys = lagrange_polynomials(points);
+    #[cfg(not(test))]
     let lagrange_polys = F::memoize_lagrange(points);
-    lagrange_interpolation_with_polys(&lagrange_polys, values)
+
+    lagrange_interpolation_with_polys(lagrange_polys, values)
 }
 
 /// interpolate a polynomial using pre-computed Lagrange basis polynomials and y-coordinates
 pub fn lagrange_interpolation_with_polys<F: Field>(
-    lagrange_polys: &[Poly<F>],
+    lagrange_polys: impl AsRef<[Poly<F>]>,
     values: &[F],
 ) -> anyhow::Result<Poly<F>> {
+    let lagrange_polys = lagrange_polys.as_ref();
     if lagrange_polys.len() != values.len() {
         return Err(anyhow_error_and_log(
             "Lagrange interpolation failure: mismatch between number of points and values"
diff --git a/core/threshold-algebra/src/structure_traits.rs b/core/threshold-algebra/src/structure_traits.rs
index c9d36595b..76227bd68 100644
--- a/core/threshold-algebra/src/structure_traits.rs
+++ b/core/threshold-algebra/src/structure_traits.rs
@@ -92,7 +92,7 @@ pub trait Field
 where
     Self: RingWithExceptionalSequence + Div<Self, Output = Self> + DivAssign<Self>,
 {
-    fn memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>>;
+    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>];
 
     /// computes the multiplicative inverse of the field element
     fn invert(&self) -> Self {
diff --git a/core/threshold-experimental/src/algebra/levels.rs b/core/threshold-experimental/src/algebra/levels.rs
index 0116b431a..7faf4eee7 100644
--- a/core/threshold-experimental/src/algebra/levels.rs
+++ b/core/threshold-experimental/src/algebra/levels.rs
@@ -23,7 +23,6 @@ use rand::Rng;
 use serde::{Deserialize, Serialize};
 use std::iter::Sum;
 use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
-#[cfg(not(test))]
 use std::sync::OnceLock;
 
 use crate::algebra::crt::LevelKswCrtRepresentation;
@@ -348,21 +347,12 @@ macro_rules! impl_field_level {
                 }
             }
 
-            #[cfg(not(test))]
             static [<LAGRANGE_STORE_BGV_ $name:upper>]: OnceLock<Vec<Poly<$name>>> = OnceLock::new();
 
             impl Field for $name {
-                fn memoize_lagrange(points: &[Self]) -> Vec<Poly<Self>> {
-                    #[cfg(test)]
-                    {
-                        lagrange_polynomials(points)
-                    }
-                    #[cfg(not(test))]
-                    {
-                        [<LAGRANGE_STORE_BGV_ $name:upper>]
-                            .get_or_init(|| lagrange_polynomials(points))
-                            .clone()
-                    }
+                fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
+                    [<LAGRANGE_STORE_BGV_ $name:upper>]
+                        .get_or_init(|| lagrange_polynomials(points))
                 }
 
                 fn invert(&self) -> Self {

From beae4debd5aaf41f3a713ec4e2d2a345bda3d44e Mon Sep 17 00:00:00 2001
From: Simon Eudeline <simon.eudeline@zama.ai>
Date: Mon, 13 Apr 2026 15:47:25 +0200
Subject: [PATCH 03/11] ci: update check-changes

---
 .github/workflows/main.yml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 7877d1bed..785ddbf7a 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -80,12 +80,15 @@ jobs:
             core-service:
               - 'core/grpc/**'
               - 'core/service/**'
+              - 'core/thread-handles/**'
               - 'core/threshold/**'
+              - 'core/threshold-*/**'
               - 'observability/**'
               - 'Cargo.toml'
               - 'Cargo.lock'
             core-threshold:
               - 'core/threshold/**'
+              - 'core/threshold-*/**'
               - 'observability/**'
               - 'Cargo.toml'
               - 'Cargo.lock'

From d0942dcb3cba268b54e66b5ad358ad9ac12c3314 Mon Sep 17 00:00:00 2001
From: Simon Eudeline <simon.eudeline@zama.ai>
Date: Tue, 14 Apr 2026 15:17:39 +0200
Subject: [PATCH 04/11] chore: pre-compute lagrange store

---
 .gitignore                                    |  1 +
 core/service/src/bin/kms-server.rs            |  7 +++
 .../src/galois_fields/common.rs               | 57 ++++++++++++++++++-
 .../src/galois_fields/gf128.rs                |  8 +--
 .../src/galois_fields/gf16.rs                 |  8 +--
 .../src/galois_fields/gf256.rs                |  8 +--
 .../src/galois_fields/gf32.rs                 |  8 +--
 .../src/galois_fields/gf64.rs                 |  8 +--
 .../src/galois_fields/gf8.rs                  | 11 ++--
 .../src/galois_fields/mod.rs                  |  6 ++
 core/threshold-algebra/src/poly.rs            | 13 ++---
 .../threshold-algebra/src/structure_traits.rs |  4 +-
 .../src/algebra/levels.rs                     | 12 ++--
 core/threshold/src/grpc/server.rs             |  7 +++
 14 files changed, 115 insertions(+), 43 deletions(-)

diff --git a/.gitignore b/.gitignore
index ae934ab8c..15c6b23a5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,6 @@
 CLAUDE.md
 .claude/
+.codex
 /target*
 **/temp/*
 /keys
diff --git a/core/service/src/bin/kms-server.rs b/core/service/src/bin/kms-server.rs
index 9c2577204..ca0bd1f70 100644
--- a/core/service/src/bin/kms-server.rs
+++ b/core/service/src/bin/kms-server.rs
@@ -1,3 +1,4 @@
+use algebra::galois_fields::common::init_all_lagrange_stores;
 use anyhow::ensure;
 use clap::Parser;
 use futures_util::future::OptionFuture;
@@ -367,6 +368,12 @@ async fn main_exec() -> anyhow::Result<()> {
 
     tracing::info!("Starting KMS Server with core config: {:?}", &core_config);
 
+    if let Some(threshold_config) = core_config.threshold.as_ref()
+        && let Some(peers) = threshold_config.peers.as_ref()
+    {
+        init_all_lagrange_stores(peers.len(), 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(),
diff --git a/core/threshold-algebra/src/galois_fields/common.rs b/core/threshold-algebra/src/galois_fields/common.rs
index f1f094559..da473eb1b 100644
--- a/core/threshold-algebra/src/galois_fields/common.rs
+++ b/core/threshold-algebra/src/galois_fields/common.rs
@@ -1,4 +1,59 @@
-use crate::{sharing::shamir::ShamirFieldPoly, structure_traits::Field, syndrome::decode_syndrome};
+use super::LagrangeMap;
+use crate::{
+    poly::lagrange_polynomials, sharing::shamir::ShamirFieldPoly, structure_traits::Field,
+    syndrome::decode_syndrome,
+};
+use itertools::Itertools;
+
+/// Builds a LagrangeMap containing pre-computed Lagrange polynomials for all sorted subsets
+/// of `{embed(1), ..., embed(num_parties)}` with size ≥ `threshold + 1`.
+fn build_lagrange_map<F: Field>(
+    num_parties: usize,
+    threshold: usize,
+) -> anyhow::Result<LagrangeMap<F>> {
+    let all_points: Vec<F> = (1..=num_parties)
+        .map(F::get_from_exceptional_sequence)
+        .collect::<anyhow::Result<Vec<_>>>()?;
+    let mut map = LagrangeMap::new();
+    for size in (threshold + 1)..=num_parties {
+        for subset in all_points.iter().combinations(size) {
+            let points: Vec<F> = subset.into_iter().copied().collect();
+            let polys = lagrange_polynomials(&points);
+            map.insert(points, polys);
+        }
+    }
+    Ok(map)
+}
+
+/// Pre-computes Lagrange polynomial stores for all enabled Galois field types.
+/// Must be called at startup with the known number of parties and threshold.
+pub fn init_all_lagrange_stores(num_parties: usize, threshold: usize) -> anyhow::Result<()> {
+    #[cfg(feature = "extension_degree_3")]
+    super::gf8::LAGRANGE_STORE
+        .set(build_lagrange_map(num_parties, threshold)?)
+        .ok();
+    #[cfg(feature = "extension_degree_4")]
+    super::gf16::LAGRANGE_STORE
+        .set(build_lagrange_map(num_parties, threshold)?)
+        .ok();
+    #[cfg(feature = "extension_degree_5")]
+    super::gf32::LAGRANGE_STORE
+        .set(build_lagrange_map(num_parties, threshold)?)
+        .ok();
+    #[cfg(feature = "extension_degree_6")]
+    super::gf64::LAGRANGE_STORE
+        .set(build_lagrange_map(num_parties, threshold)?)
+        .ok();
+    #[cfg(feature = "extension_degree_7")]
+    super::gf128::LAGRANGE_STORE
+        .set(build_lagrange_map(num_parties, threshold)?)
+        .ok();
+    #[cfg(feature = "extension_degree_8")]
+    super::gf256::LAGRANGE_STORE
+        .set(build_lagrange_map(num_parties, threshold)?)
+        .ok();
+    Ok(())
+}
 
 pub fn syndrome_decoding_z2<F: Field + From<u8>>(
     parties: &[usize],
diff --git a/core/threshold-algebra/src/galois_fields/gf128.rs b/core/threshold-algebra/src/galois_fields/gf128.rs
index 09a130e43..d7bc03888 100644
--- a/core/threshold-algebra/src/galois_fields/gf128.rs
+++ b/core/threshold-algebra/src/galois_fields/gf128.rs
@@ -1,5 +1,5 @@
-use crate::poly::lagrange_polynomials;
 use crate::{
+    galois_fields::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
@@ -84,11 +84,11 @@ impl RingWithExceptionalSequence for GF128 {
     }
 }
 
-static LAGRANGE_STORE: OnceLock<Vec<Poly<GF128>>> = OnceLock::new();
+pub(crate) static LAGRANGE_STORE: OnceLock<LagrangeMap<GF128>> = OnceLock::new();
 
 impl Field for GF128 {
-    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
-        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
+    fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
+        LAGRANGE_STORE.get()?.get(points).map(|v| v.as_slice())
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/gf16.rs b/core/threshold-algebra/src/galois_fields/gf16.rs
index dd1864e2a..113c6fce7 100644
--- a/core/threshold-algebra/src/galois_fields/gf16.rs
+++ b/core/threshold-algebra/src/galois_fields/gf16.rs
@@ -1,5 +1,5 @@
-use crate::poly::lagrange_polynomials;
 use crate::{
+    galois_fields::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
@@ -84,11 +84,11 @@ impl RingWithExceptionalSequence for GF16 {
     }
 }
 
-static LAGRANGE_STORE: OnceLock<Vec<Poly<GF16>>> = OnceLock::new();
+pub(crate) static LAGRANGE_STORE: OnceLock<LagrangeMap<GF16>> = OnceLock::new();
 
 impl Field for GF16 {
-    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
-        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
+    fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
+        LAGRANGE_STORE.get()?.get(points).map(|v| v.as_slice())
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/gf256.rs b/core/threshold-algebra/src/galois_fields/gf256.rs
index 585f78ca2..7f61618a9 100644
--- a/core/threshold-algebra/src/galois_fields/gf256.rs
+++ b/core/threshold-algebra/src/galois_fields/gf256.rs
@@ -1,5 +1,5 @@
-use crate::poly::lagrange_polynomials;
 use crate::{
+    galois_fields::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
@@ -84,11 +84,11 @@ impl RingWithExceptionalSequence for GF256 {
     }
 }
 
-static LAGRANGE_STORE: OnceLock<Vec<Poly<GF256>>> = OnceLock::new();
+pub(crate) static LAGRANGE_STORE: OnceLock<LagrangeMap<GF256>> = OnceLock::new();
 
 impl Field for GF256 {
-    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
-        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
+    fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
+        LAGRANGE_STORE.get()?.get(points).map(|v| v.as_slice())
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/gf32.rs b/core/threshold-algebra/src/galois_fields/gf32.rs
index 430d5adfe..9c68b9af6 100644
--- a/core/threshold-algebra/src/galois_fields/gf32.rs
+++ b/core/threshold-algebra/src/galois_fields/gf32.rs
@@ -1,5 +1,5 @@
-use crate::poly::lagrange_polynomials;
 use crate::{
+    galois_fields::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
@@ -84,11 +84,11 @@ impl RingWithExceptionalSequence for GF32 {
     }
 }
 
-static LAGRANGE_STORE: OnceLock<Vec<Poly<GF32>>> = OnceLock::new();
+pub(crate) static LAGRANGE_STORE: OnceLock<LagrangeMap<GF32>> = OnceLock::new();
 
 impl Field for GF32 {
-    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
-        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
+    fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
+        LAGRANGE_STORE.get()?.get(points).map(|v| v.as_slice())
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/gf64.rs b/core/threshold-algebra/src/galois_fields/gf64.rs
index a0817e471..5239ff913 100644
--- a/core/threshold-algebra/src/galois_fields/gf64.rs
+++ b/core/threshold-algebra/src/galois_fields/gf64.rs
@@ -1,5 +1,5 @@
-use crate::poly::lagrange_polynomials;
 use crate::{
+    galois_fields::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
@@ -84,11 +84,11 @@ impl RingWithExceptionalSequence for GF64 {
     }
 }
 
-static LAGRANGE_STORE: OnceLock<Vec<Poly<GF64>>> = OnceLock::new();
+pub(crate) static LAGRANGE_STORE: OnceLock<LagrangeMap<GF64>> = OnceLock::new();
 
 impl Field for GF64 {
-    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
-        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
+    fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
+        LAGRANGE_STORE.get()?.get(points).map(|v| v.as_slice())
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/gf8.rs b/core/threshold-algebra/src/galois_fields/gf8.rs
index 28c1c9254..1fd688fd4 100644
--- a/core/threshold-algebra/src/galois_fields/gf8.rs
+++ b/core/threshold-algebra/src/galois_fields/gf8.rs
@@ -1,8 +1,7 @@
-use crate::poly::lagrange_polynomials;
-use crate::structure_traits::RingWithExceptionalSequence;
 use crate::{
+    galois_fields::LagrangeMap,
     poly::Poly,
-    structure_traits::{Field, FromU128, One, Ring, Sample, Zero},
+    structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
@@ -85,11 +84,11 @@ impl RingWithExceptionalSequence for GF8 {
     }
 }
 
-static LAGRANGE_STORE: OnceLock<Vec<Poly<GF8>>> = OnceLock::new();
+pub(crate) static LAGRANGE_STORE: OnceLock<LagrangeMap<GF8>> = OnceLock::new();
 
 impl Field for GF8 {
-    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
-        LAGRANGE_STORE.get_or_init(|| lagrange_polynomials(points))
+    fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
+        LAGRANGE_STORE.get()?.get(points).map(|v| v.as_slice())
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/mod.rs b/core/threshold-algebra/src/galois_fields/mod.rs
index bfc0fd8f8..2a80f019d 100644
--- a/core/threshold-algebra/src/galois_fields/mod.rs
+++ b/core/threshold-algebra/src/galois_fields/mod.rs
@@ -1,3 +1,9 @@
+use crate::poly::Poly;
+use std::collections::HashMap;
+
+/// Map from evaluation points to their precomputed Lagrange polynomials.
+pub(crate) type LagrangeMap<F> = HashMap<Vec<F>, Vec<Poly<F>>>;
+
 pub mod common;
 #[cfg(feature = "extension_degree_7")]
 pub mod gf128;
diff --git a/core/threshold-algebra/src/poly.rs b/core/threshold-algebra/src/poly.rs
index 98a1e458b..e03e222ba 100644
--- a/core/threshold-algebra/src/poly.rs
+++ b/core/threshold-algebra/src/poly.rs
@@ -600,14 +600,11 @@ pub fn lagrange_polynomials<F: Field>(points: &[F]) -> Vec<Poly<F>> {
 
 /// interpolate a polynomial through coordinates where points holds the x-coordinates and values holds the y-coordinates
 pub fn lagrange_interpolation<F: Field>(points: &[F], values: &[F]) -> anyhow::Result<Poly<F>> {
-    // In tests, bypass memoization to allow different point sets per test.
-    // memoize_lagrange uses OnceLock which caches a single point set per field type.
-    #[cfg(test)]
-    let lagrange_polys = lagrange_polynomials(points);
-    #[cfg(not(test))]
-    let lagrange_polys = F::memoize_lagrange(points);
-
-    lagrange_interpolation_with_polys(lagrange_polys, values)
+    if let Some(cached) = F::cached_lagrange_polys(points) {
+        lagrange_interpolation_with_polys(cached, values)
+    } else {
+        lagrange_interpolation_with_polys(lagrange_polynomials(points), values)
+    }
 }
 
 /// interpolate a polynomial using pre-computed Lagrange basis polynomials and y-coordinates
diff --git a/core/threshold-algebra/src/structure_traits.rs b/core/threshold-algebra/src/structure_traits.rs
index 76227bd68..dd64df0e7 100644
--- a/core/threshold-algebra/src/structure_traits.rs
+++ b/core/threshold-algebra/src/structure_traits.rs
@@ -92,7 +92,9 @@ pub trait Field
 where
     Self: RingWithExceptionalSequence + Div<Self, Output = Self> + DivAssign<Self>,
 {
-    fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>];
+    /// Looks up pre-computed Lagrange polynomials for the given point set.
+    /// Returns `None` if the store has not been initialized.
+    fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]>;
 
     /// computes the multiplicative inverse of the field element
     fn invert(&self) -> Self {
diff --git a/core/threshold-experimental/src/algebra/levels.rs b/core/threshold-experimental/src/algebra/levels.rs
index 7faf4eee7..504513151 100644
--- a/core/threshold-experimental/src/algebra/levels.rs
+++ b/core/threshold-experimental/src/algebra/levels.rs
@@ -1,7 +1,7 @@
 use algebra::{
     PRSSConversions,
     error_correction::error_correction,
-    poly::{Poly, lagrange_polynomials},
+    poly::Poly,
     sharing::{shamir::ShamirSharings, share::Share},
     structure_traits::{
         ErrorCorrect, Field, FromU128, Invert, One, Ring, RingWithExceptionalSequence, Sample,
@@ -23,7 +23,6 @@ use rand::Rng;
 use serde::{Deserialize, Serialize};
 use std::iter::Sum;
 use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
-use std::sync::OnceLock;
 
 use crate::algebra::crt::LevelKswCrtRepresentation;
 use crate::algebra::crt::from_crt;
@@ -347,12 +346,11 @@ macro_rules! impl_field_level {
                 }
             }
 
-            static [<LAGRANGE_STORE_BGV_ $name:upper>]: OnceLock<Vec<Poly<$name>>> = OnceLock::new();
-
             impl Field for $name {
-                fn memoize_lagrange(points: &[Self]) -> &'static [Poly<Self>] {
-                    [<LAGRANGE_STORE_BGV_ $name:upper>]
-                        .get_or_init(|| lagrange_polynomials(points))
+                fn cached_lagrange_polys(_points: &[Self]) -> Option<&'static [Poly<Self>]> {
+                    // BGV level types are experimental; no pre-computed store.
+                    // Falls back to direct computation in lagrange_interpolation.
+                    None
                 }
 
                 fn invert(&self) -> Self {
diff --git a/core/threshold/src/grpc/server.rs b/core/threshold/src/grpc/server.rs
index 73d3c4ac8..0f23bb519 100644
--- a/core/threshold/src/grpc/server.rs
+++ b/core/threshold/src/grpc/server.rs
@@ -1,6 +1,7 @@
 use crate::{choreography::grpc::GrpcChoreography, conf::party::PartyConf};
 use algebra::{
     base_ring::{Z64, Z128},
+    galois_fields::common::init_all_lagrange_stores,
     galois_rings::common::ResiduePoly,
     structure_traits::{Derive, ErrorCorrect, Invert, Solve, Syndrome},
 };
@@ -36,6 +37,12 @@ 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_all_lagrange_stores(num_parties, threshold)?;
+    }
 
     let tls_conf = settings
         .certpaths

From 6ec797f4f96587d24b25bd431699ec1afad5692f Mon Sep 17 00:00:00 2001
From: Titouan Tanguy <titouan.tanguy@zama.ai>
Date: Fri, 24 Apr 2026 14:30:59 +0200
Subject: [PATCH 05/11] chore fully init the smaller fields' cache of lagrange
 basis

---
 core/service/src/bin/kms-server.rs                 | 13 ++++++++-----
 core/threshold-algebra/src/galois_fields/common.rs | 14 +++++---------
 core/threshold-algebra/src/galois_fields/gf16.rs   | 12 ++++++++----
 core/threshold-algebra/src/galois_fields/gf8.rs    | 12 ++++++++----
 4 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/core/service/src/bin/kms-server.rs b/core/service/src/bin/kms-server.rs
index ca0bd1f70..23f069f79 100644
--- a/core/service/src/bin/kms-server.rs
+++ b/core/service/src/bin/kms-server.rs
@@ -368,11 +368,14 @@ async fn main_exec() -> anyhow::Result<()> {
 
     tracing::info!("Starting KMS Server with core config: {:?}", &core_config);
 
-    if let Some(threshold_config) = core_config.threshold.as_ref()
-        && let Some(peers) = threshold_config.peers.as_ref()
-    {
-        init_all_lagrange_stores(peers.len(), threshold_config.threshold as usize)?;
-    }
+    // NOTE: Cache for GF16 (which we use here -- for now) is fully
+    // hence why the block below is commented out.
+
+    //if let Some(threshold_config) = core_config.threshold.as_ref()
+    //    && let Some(peers) = threshold_config.peers.as_ref()
+    //{
+    //    init_all_lagrange_stores(peers.len(), threshold_config.threshold as usize)?;
+    //}
 
     tracing::info!(
         "Multi-threading values: tokio::num_workers: {}, rayon_num_threads: {}, total_num_cpus: {}",
diff --git a/core/threshold-algebra/src/galois_fields/common.rs b/core/threshold-algebra/src/galois_fields/common.rs
index da473eb1b..e2adbe378 100644
--- a/core/threshold-algebra/src/galois_fields/common.rs
+++ b/core/threshold-algebra/src/galois_fields/common.rs
@@ -7,7 +7,7 @@ use itertools::Itertools;
 
 /// Builds a LagrangeMap containing pre-computed Lagrange polynomials for all sorted subsets
 /// of `{embed(1), ..., embed(num_parties)}` with size ≥ `threshold + 1`.
-fn build_lagrange_map<F: Field>(
+pub(crate) fn build_lagrange_map<F: Field>(
     num_parties: usize,
     threshold: usize,
 ) -> anyhow::Result<LagrangeMap<F>> {
@@ -25,17 +25,13 @@ fn build_lagrange_map<F: Field>(
     Ok(map)
 }
 
+#[allow(unused_variables)]
 /// Pre-computes Lagrange polynomial stores for all enabled Galois field types.
 /// Must be called at startup with the known number of parties and threshold.
+///
+/// __NOTE__: GF8 and GF16 are small enough that we can pre-compute all possible lagrange basis.
+/// so they are skipped here.
 pub fn init_all_lagrange_stores(num_parties: usize, threshold: usize) -> anyhow::Result<()> {
-    #[cfg(feature = "extension_degree_3")]
-    super::gf8::LAGRANGE_STORE
-        .set(build_lagrange_map(num_parties, threshold)?)
-        .ok();
-    #[cfg(feature = "extension_degree_4")]
-    super::gf16::LAGRANGE_STORE
-        .set(build_lagrange_map(num_parties, threshold)?)
-        .ok();
     #[cfg(feature = "extension_degree_5")]
     super::gf32::LAGRANGE_STORE
         .set(build_lagrange_map(num_parties, threshold)?)
diff --git a/core/threshold-algebra/src/galois_fields/gf16.rs b/core/threshold-algebra/src/galois_fields/gf16.rs
index 113c6fce7..3e56285e1 100644
--- a/core/threshold-algebra/src/galois_fields/gf16.rs
+++ b/core/threshold-algebra/src/galois_fields/gf16.rs
@@ -1,11 +1,11 @@
 use crate::{
-    galois_fields::LagrangeMap,
+    galois_fields::{LagrangeMap, common::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, OnceLock};
+use std::sync::LazyLock;
 
 g2p!(
     GF16,
@@ -84,11 +84,15 @@ impl RingWithExceptionalSequence for GF16 {
     }
 }
 
-pub(crate) static LAGRANGE_STORE: OnceLock<LagrangeMap<GF16>> = OnceLock::new();
+pub(crate) static LAGRANGE_STORE: LazyLock<LagrangeMap<GF16>> = LazyLock::new(|| {
+    build_lagrange_map::<GF16>(0, 15).expect(
+        " Initializaiton of the GF16 Lagrange basis can't fail with 15 parties and threshold 0",
+    )
+});
 
 impl Field for GF16 {
     fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
-        LAGRANGE_STORE.get()?.get(points).map(|v| v.as_slice())
+        LAGRANGE_STORE.get(points).map(|v| v.as_slice())
     }
 }
 
diff --git a/core/threshold-algebra/src/galois_fields/gf8.rs b/core/threshold-algebra/src/galois_fields/gf8.rs
index 1fd688fd4..a36cf96a3 100644
--- a/core/threshold-algebra/src/galois_fields/gf8.rs
+++ b/core/threshold-algebra/src/galois_fields/gf8.rs
@@ -1,11 +1,11 @@
 use crate::{
-    galois_fields::LagrangeMap,
+    galois_fields::{LagrangeMap, common::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, OnceLock};
+use std::sync::LazyLock;
 
 g2p!(
     GF8,
@@ -84,11 +84,15 @@ impl RingWithExceptionalSequence for GF8 {
     }
 }
 
-pub(crate) static LAGRANGE_STORE: OnceLock<LagrangeMap<GF8>> = OnceLock::new();
+pub(crate) static LAGRANGE_STORE: LazyLock<LagrangeMap<GF8>> = LazyLock::new(|| {
+    build_lagrange_map::<GF8>(0, 7).expect(
+        " Initializaiton of the GF8 Lagrange basis can't fail with 7 parties and threshold 0",
+    )
+});
 
 impl Field for GF8 {
     fn cached_lagrange_polys(points: &[Self]) -> Option<&'static [Poly<Self>]> {
-        LAGRANGE_STORE.get()?.get(points).map(|v| v.as_slice())
+        LAGRANGE_STORE.get(points).map(|v| v.as_slice())
     }
 }
 

From 536cd2e50709d19fcd1facde237e5ddae64dc5b6 Mon Sep 17 00:00:00 2001
From: Titouan Tanguy <titouan.tanguy@zama.ai>
Date: Fri, 24 Apr 2026 14:43:20 +0200
Subject: [PATCH 06/11] chore: lint

---
 core/service/src/bin/kms-server.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/core/service/src/bin/kms-server.rs b/core/service/src/bin/kms-server.rs
index 23f069f79..97d99a7d8 100644
--- a/core/service/src/bin/kms-server.rs
+++ b/core/service/src/bin/kms-server.rs
@@ -1,4 +1,3 @@
-use algebra::galois_fields::common::init_all_lagrange_stores;
 use anyhow::ensure;
 use clap::Parser;
 use futures_util::future::OptionFuture;

From af896e0e246a1d54d003e6d49a70af57792693f6 Mon Sep 17 00:00:00 2001
From: Titouan Tanguy <titouan.tanguy@zama.ai>
Date: Fri, 24 Apr 2026 16:34:09 +0200
Subject: [PATCH 07/11] chore: fix my mistakes

---
 core/service/src/bin/kms-server.rs            | 22 ++++++++++-------
 .../src/galois_fields/common.rs               | 24 ++++++++++++-------
 .../src/galois_fields/gf16.rs                 |  8 ++++---
 .../src/galois_fields/gf8.rs                  |  8 ++++---
 core/threshold/src/grpc/server.rs             |  7 ++++--
 5 files changed, 44 insertions(+), 25 deletions(-)

diff --git a/core/service/src/bin/kms-server.rs b/core/service/src/bin/kms-server.rs
index 97d99a7d8..c5052508b 100644
--- a/core/service/src/bin/kms-server.rs
+++ b/core/service/src/bin/kms-server.rs
@@ -1,3 +1,4 @@
+use algebra::galois_fields::common::init_all_lagrange_stores;
 use anyhow::ensure;
 use clap::Parser;
 use futures_util::future::OptionFuture;
@@ -367,14 +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
-    // hence why the block below is commented out.
-
-    //if let Some(threshold_config) = core_config.threshold.as_ref()
-    //    && let Some(peers) = threshold_config.peers.as_ref()
-    //{
-    //    init_all_lagrange_stores(peers.len(), threshold_config.threshold as usize)?;
-    //}
+    // 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()
+    {
+        if peers.len() > 0 {
+            init_all_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: {}",
diff --git a/core/threshold-algebra/src/galois_fields/common.rs b/core/threshold-algebra/src/galois_fields/common.rs
index e2adbe378..556d3cb05 100644
--- a/core/threshold-algebra/src/galois_fields/common.rs
+++ b/core/threshold-algebra/src/galois_fields/common.rs
@@ -1,3 +1,5 @@
+use std::num::NonZero;
+
 use super::LagrangeMap;
 use crate::{
     poly::lagrange_polynomials, sharing::shamir::ShamirFieldPoly, structure_traits::Field,
@@ -6,16 +8,17 @@ use crate::{
 use itertools::Itertools;
 
 /// Builds a LagrangeMap containing pre-computed Lagrange polynomials for all sorted subsets
-/// of `{embed(1), ..., embed(num_parties)}` with size ≥ `threshold + 1`.
+/// of `{embed(1), ..., embed(num_parties)}` with size ≥ `min_threshold + 1`.
 pub(crate) fn build_lagrange_map<F: Field>(
-    num_parties: usize,
-    threshold: usize,
+    num_parties: NonZero<usize>,
+    min_threshold: usize,
 ) -> anyhow::Result<LagrangeMap<F>> {
+    let num_parties = num_parties.get();
     let all_points: Vec<F> = (1..=num_parties)
         .map(F::get_from_exceptional_sequence)
         .collect::<anyhow::Result<Vec<_>>>()?;
     let mut map = LagrangeMap::new();
-    for size in (threshold + 1)..=num_parties {
+    for size in (min_threshold + 1)..=num_parties {
         for subset in all_points.iter().combinations(size) {
             let points: Vec<F> = subset.into_iter().copied().collect();
             let polys = lagrange_polynomials(&points);
@@ -31,22 +34,25 @@ pub(crate) fn build_lagrange_map<F: Field>(
 ///
 /// __NOTE__: GF8 and GF16 are small enough that we can pre-compute all possible lagrange basis.
 /// so they are skipped here.
-pub fn init_all_lagrange_stores(num_parties: usize, threshold: usize) -> anyhow::Result<()> {
+pub fn init_all_lagrange_stores(
+    num_parties: NonZero<usize>,
+    min_threshold: usize,
+) -> anyhow::Result<()> {
     #[cfg(feature = "extension_degree_5")]
     super::gf32::LAGRANGE_STORE
-        .set(build_lagrange_map(num_parties, threshold)?)
+        .set(build_lagrange_map(num_parties, min_threshold)?)
         .ok();
     #[cfg(feature = "extension_degree_6")]
     super::gf64::LAGRANGE_STORE
-        .set(build_lagrange_map(num_parties, threshold)?)
+        .set(build_lagrange_map(num_parties, min_threshold)?)
         .ok();
     #[cfg(feature = "extension_degree_7")]
     super::gf128::LAGRANGE_STORE
-        .set(build_lagrange_map(num_parties, threshold)?)
+        .set(build_lagrange_map(num_parties, min_threshold)?)
         .ok();
     #[cfg(feature = "extension_degree_8")]
     super::gf256::LAGRANGE_STORE
-        .set(build_lagrange_map(num_parties, threshold)?)
+        .set(build_lagrange_map(num_parties, min_threshold)?)
         .ok();
     Ok(())
 }
diff --git a/core/threshold-algebra/src/galois_fields/gf16.rs b/core/threshold-algebra/src/galois_fields/gf16.rs
index 3e56285e1..ab6aa6f43 100644
--- a/core/threshold-algebra/src/galois_fields/gf16.rs
+++ b/core/threshold-algebra/src/galois_fields/gf16.rs
@@ -5,7 +5,7 @@ use crate::{
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::LazyLock;
+use std::{num::NonZero, sync::LazyLock};
 
 g2p!(
     GF16,
@@ -84,9 +84,11 @@ impl RingWithExceptionalSequence for GF16 {
     }
 }
 
+/// 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>(0, 15).expect(
-        " Initializaiton of the GF16 Lagrange basis can't fail with 15 parties and threshold 0",
+    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",
     )
 });
 
diff --git a/core/threshold-algebra/src/galois_fields/gf8.rs b/core/threshold-algebra/src/galois_fields/gf8.rs
index a36cf96a3..e3133d396 100644
--- a/core/threshold-algebra/src/galois_fields/gf8.rs
+++ b/core/threshold-algebra/src/galois_fields/gf8.rs
@@ -5,7 +5,7 @@ use crate::{
 };
 use g2p::{GaloisField, g2p};
 use serde::{Deserialize, Serialize};
-use std::sync::LazyLock;
+use std::{num::NonZero, sync::LazyLock};
 
 g2p!(
     GF8,
@@ -84,9 +84,11 @@ impl RingWithExceptionalSequence for GF8 {
     }
 }
 
+/// Pre-computed Lagrange basis for all (ordered) non-empty subsets of GF8 that exclude 0.
+/// Size is \sum_{k=1}^{7} C(7, k) * k = 7 * 2^6 = 448, which is small enough to be pre-computed and stored in memory.
 pub(crate) static LAGRANGE_STORE: LazyLock<LagrangeMap<GF8>> = LazyLock::new(|| {
-    build_lagrange_map::<GF8>(0, 7).expect(
-        " Initializaiton of the GF8 Lagrange basis can't fail with 7 parties and threshold 0",
+    build_lagrange_map::<GF8>(NonZero::new(7).expect("7 is non-zero"), 0).expect(
+        "Initialization of the GF8 Lagrange basis can't fail with 7 parties and threshold 0",
     )
 });
 
diff --git a/core/threshold/src/grpc/server.rs b/core/threshold/src/grpc/server.rs
index 0f23bb519..b81f73d87 100644
--- a/core/threshold/src/grpc/server.rs
+++ b/core/threshold/src/grpc/server.rs
@@ -6,7 +6,7 @@ use algebra::{
     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::{
     large_execution::offline::SecureLargePreprocessing,
     online::preprocessing::{PreprocessorFactory, create_memory_factory, create_redis_factory},
@@ -41,7 +41,10 @@ where
         let num_parties = peers.len() + 1;
         // The threshold-fhe config currently assumes n = 3t + 1.
         let threshold = (num_parties - 1) / 3;
-        init_all_lagrange_stores(num_parties, threshold)?;
+        init_all_lagrange_stores(
+            NonZero::new(num_parties).expect("num_parties must be non-zero"),
+            threshold,
+        )?;
     }
 
     let tls_conf = settings

From 03fe74919e026cea0872ad9e80d72ecb503345de Mon Sep 17 00:00:00 2001
From: Titouan Tanguy <titouan.tanguy@zama.ai>
Date: Fri, 24 Apr 2026 17:06:26 +0200
Subject: [PATCH 08/11] chore: missing import

---
 core/service/src/bin/kms-server.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/service/src/bin/kms-server.rs b/core/service/src/bin/kms-server.rs
index c5052508b..76dbd648d 100644
--- a/core/service/src/bin/kms-server.rs
+++ b/core/service/src/bin/kms-server.rs
@@ -33,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;

From 468101bd6c2006426116f3a67b6fd26b9c12f103 Mon Sep 17 00:00:00 2001
From: Titouan Tanguy <titouan.tanguy@zama.ai>
Date: Fri, 24 Apr 2026 17:28:22 +0200
Subject: [PATCH 09/11] chore: clippy

---
 core/service/src/bin/kms-server.rs | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/core/service/src/bin/kms-server.rs b/core/service/src/bin/kms-server.rs
index 76dbd648d..461a4bf3d 100644
--- a/core/service/src/bin/kms-server.rs
+++ b/core/service/src/bin/kms-server.rs
@@ -372,14 +372,13 @@ async fn main_exec() -> anyhow::Result<()> {
     // 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()
     {
-        if peers.len() > 0 {
-            init_all_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,
-            )?;
-        }
+        init_all_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!(

From eee013cb20ce0e93106c9d1905d9c39f5b221fdd Mon Sep 17 00:00:00 2001
From: Titouan Tanguy <titouan.tanguy@zama.ai>
Date: Mon, 27 Apr 2026 16:02:05 +0200
Subject: [PATCH 10/11] chore: move alias

---
 core/threshold-algebra/src/galois_fields/common.rs | 2 +-
 core/threshold-algebra/src/galois_fields/gf128.rs  | 2 +-
 core/threshold-algebra/src/galois_fields/gf16.rs   | 2 +-
 core/threshold-algebra/src/galois_fields/gf256.rs  | 2 +-
 core/threshold-algebra/src/galois_fields/gf32.rs   | 2 +-
 core/threshold-algebra/src/galois_fields/gf64.rs   | 2 +-
 core/threshold-algebra/src/galois_fields/gf8.rs    | 2 +-
 core/threshold-algebra/src/galois_fields/mod.rs    | 2 --
 8 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/core/threshold-algebra/src/galois_fields/common.rs b/core/threshold-algebra/src/galois_fields/common.rs
index 556d3cb05..7b01498ba 100644
--- a/core/threshold-algebra/src/galois_fields/common.rs
+++ b/core/threshold-algebra/src/galois_fields/common.rs
@@ -1,12 +1,12 @@
 use std::num::NonZero;
 
-use super::LagrangeMap;
 use crate::{
     poly::lagrange_polynomials, sharing::shamir::ShamirFieldPoly, structure_traits::Field,
     syndrome::decode_syndrome,
 };
 use itertools::Itertools;
 
+pub(crate) type LagrangeMap<F> = std::collections::HashMap<Vec<F>, Vec<crate::poly::Poly<F>>>;
 /// Builds a LagrangeMap containing pre-computed Lagrange polynomials for all sorted subsets
 /// of `{embed(1), ..., embed(num_parties)}` with size ≥ `min_threshold + 1`.
 pub(crate) fn build_lagrange_map<F: Field>(
diff --git a/core/threshold-algebra/src/galois_fields/gf128.rs b/core/threshold-algebra/src/galois_fields/gf128.rs
index d7bc03888..61f1990f3 100644
--- a/core/threshold-algebra/src/galois_fields/gf128.rs
+++ b/core/threshold-algebra/src/galois_fields/gf128.rs
@@ -1,5 +1,5 @@
 use crate::{
-    galois_fields::LagrangeMap,
+    galois_fields::common::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
diff --git a/core/threshold-algebra/src/galois_fields/gf16.rs b/core/threshold-algebra/src/galois_fields/gf16.rs
index ab6aa6f43..2a5b1126a 100644
--- a/core/threshold-algebra/src/galois_fields/gf16.rs
+++ b/core/threshold-algebra/src/galois_fields/gf16.rs
@@ -1,5 +1,5 @@
 use crate::{
-    galois_fields::{LagrangeMap, common::build_lagrange_map},
+    galois_fields::common::{LagrangeMap, build_lagrange_map},
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
diff --git a/core/threshold-algebra/src/galois_fields/gf256.rs b/core/threshold-algebra/src/galois_fields/gf256.rs
index 7f61618a9..c32cc7634 100644
--- a/core/threshold-algebra/src/galois_fields/gf256.rs
+++ b/core/threshold-algebra/src/galois_fields/gf256.rs
@@ -1,5 +1,5 @@
 use crate::{
-    galois_fields::LagrangeMap,
+    galois_fields::common::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
diff --git a/core/threshold-algebra/src/galois_fields/gf32.rs b/core/threshold-algebra/src/galois_fields/gf32.rs
index 9c68b9af6..da7419ff9 100644
--- a/core/threshold-algebra/src/galois_fields/gf32.rs
+++ b/core/threshold-algebra/src/galois_fields/gf32.rs
@@ -1,5 +1,5 @@
 use crate::{
-    galois_fields::LagrangeMap,
+    galois_fields::common::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
diff --git a/core/threshold-algebra/src/galois_fields/gf64.rs b/core/threshold-algebra/src/galois_fields/gf64.rs
index 5239ff913..63934e174 100644
--- a/core/threshold-algebra/src/galois_fields/gf64.rs
+++ b/core/threshold-algebra/src/galois_fields/gf64.rs
@@ -1,5 +1,5 @@
 use crate::{
-    galois_fields::LagrangeMap,
+    galois_fields::common::LagrangeMap,
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
diff --git a/core/threshold-algebra/src/galois_fields/gf8.rs b/core/threshold-algebra/src/galois_fields/gf8.rs
index e3133d396..c482a21ed 100644
--- a/core/threshold-algebra/src/galois_fields/gf8.rs
+++ b/core/threshold-algebra/src/galois_fields/gf8.rs
@@ -1,5 +1,5 @@
 use crate::{
-    galois_fields::{LagrangeMap, common::build_lagrange_map},
+    galois_fields::common::{LagrangeMap, build_lagrange_map},
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
diff --git a/core/threshold-algebra/src/galois_fields/mod.rs b/core/threshold-algebra/src/galois_fields/mod.rs
index ed08365e7..32702ef1a 100644
--- a/core/threshold-algebra/src/galois_fields/mod.rs
+++ b/core/threshold-algebra/src/galois_fields/mod.rs
@@ -7,8 +7,6 @@
     feature = "extension_degree_7",
     feature = "extension_degree_8",
 ))]
-pub(crate) type LagrangeMap<F> = std::collections::HashMap<Vec<F>, Vec<crate::poly::Poly<F>>>;
-
 pub mod common;
 #[cfg(feature = "extension_degree_7")]
 pub mod gf128;

From 3e03b400abcf8f104ca07810cb3ce0d2b1a3d40a Mon Sep 17 00:00:00 2001
From: David Palm <dvdplm@gmail.com>
Date: Mon, 27 Apr 2026 16:19:58 +0200
Subject: [PATCH 11/11] chore: sort out features. move syndrome related code to
 syndrome, rename common to lagrange, feature gate lagrange, rename init

---
 core/experiments/src/choreography/server.rs     |  4 ++--
 core/service/src/bin/kms-server.rs              |  4 ++--
 .../threshold-algebra/src/galois_fields/gf16.rs |  2 +-
 core/threshold-algebra/src/galois_fields/gf8.rs |  2 +-
 .../galois_fields/{common.rs => lagrange.rs}    | 17 ++---------------
 core/threshold-algebra/src/galois_fields/mod.rs | 10 +++++++++-
 .../src/galois_rings/common.rs                  |  3 +--
 core/threshold-algebra/src/syndrome.rs          | 12 ++++++++++++
 8 files changed, 30 insertions(+), 24 deletions(-)
 rename core/threshold-algebra/src/galois_fields/{common.rs => lagrange.rs} (80%)

diff --git a/core/experiments/src/choreography/server.rs b/core/experiments/src/choreography/server.rs
index d1cb6c3e6..0f246808d 100644
--- a/core/experiments/src/choreography/server.rs
+++ b/core/experiments/src/choreography/server.rs
@@ -1,7 +1,7 @@
 use crate::conf::party::PartyConf;
 use algebra::{
     base_ring::{Z64, Z128},
-    galois_fields::common::init_all_lagrange_stores,
+    galois_fields::lagrange::init_lagrange_stores,
     galois_rings::common::ResiduePoly,
     structure_traits::{Derive, ErrorCorrect, Invert, Solve, Syndrome},
 };
@@ -39,7 +39,7 @@ where
         let num_parties = peers.len() + 1;
         // The threshold-fhe config currently assumes n = 3t + 1.
         let threshold = (num_parties - 1) / 3;
-        init_all_lagrange_stores(
+        init_lagrange_stores(
             NonZero::new(num_parties).expect("num_parties must be non-zero"),
             threshold,
         )?;
diff --git a/core/service/src/bin/kms-server.rs b/core/service/src/bin/kms-server.rs
index 461a4bf3d..a1d5e336f 100644
--- a/core/service/src/bin/kms-server.rs
+++ b/core/service/src/bin/kms-server.rs
@@ -1,4 +1,4 @@
-use algebra::galois_fields::common::init_all_lagrange_stores;
+use algebra::galois_fields::lagrange::init_lagrange_stores;
 use anyhow::ensure;
 use clap::Parser;
 use futures_util::future::OptionFuture;
@@ -374,7 +374,7 @@ async fn main_exec() -> anyhow::Result<()> {
         && let Some(peers) = threshold_config.peers.as_ref()
         && !peers.is_empty()
     {
-        init_all_lagrange_stores(
+        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,
diff --git a/core/threshold-algebra/src/galois_fields/gf16.rs b/core/threshold-algebra/src/galois_fields/gf16.rs
index ab6aa6f43..2a686dff3 100644
--- a/core/threshold-algebra/src/galois_fields/gf16.rs
+++ b/core/threshold-algebra/src/galois_fields/gf16.rs
@@ -1,5 +1,5 @@
 use crate::{
-    galois_fields::{LagrangeMap, common::build_lagrange_map},
+    galois_fields::{LagrangeMap, lagrange::build_lagrange_map},
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
diff --git a/core/threshold-algebra/src/galois_fields/gf8.rs b/core/threshold-algebra/src/galois_fields/gf8.rs
index e3133d396..868f4b91b 100644
--- a/core/threshold-algebra/src/galois_fields/gf8.rs
+++ b/core/threshold-algebra/src/galois_fields/gf8.rs
@@ -1,5 +1,5 @@
 use crate::{
-    galois_fields::{LagrangeMap, common::build_lagrange_map},
+    galois_fields::{LagrangeMap, lagrange::build_lagrange_map},
     poly::Poly,
     structure_traits::{Field, FromU128, One, Ring, RingWithExceptionalSequence, Sample, Zero},
 };
diff --git a/core/threshold-algebra/src/galois_fields/common.rs b/core/threshold-algebra/src/galois_fields/lagrange.rs
similarity index 80%
rename from core/threshold-algebra/src/galois_fields/common.rs
rename to core/threshold-algebra/src/galois_fields/lagrange.rs
index 556d3cb05..aa8eebfea 100644
--- a/core/threshold-algebra/src/galois_fields/common.rs
+++ b/core/threshold-algebra/src/galois_fields/lagrange.rs
@@ -1,10 +1,7 @@
 use std::num::NonZero;
 
 use super::LagrangeMap;
-use crate::{
-    poly::lagrange_polynomials, sharing::shamir::ShamirFieldPoly, structure_traits::Field,
-    syndrome::decode_syndrome,
-};
+use crate::{poly::lagrange_polynomials, structure_traits::Field};
 use itertools::Itertools;
 
 /// Builds a LagrangeMap containing pre-computed Lagrange polynomials for all sorted subsets
@@ -34,7 +31,7 @@ pub(crate) fn build_lagrange_map<F: Field>(
 ///
 /// __NOTE__: GF8 and GF16 are small enough that we can pre-compute all possible lagrange basis.
 /// so they are skipped here.
-pub fn init_all_lagrange_stores(
+pub fn init_lagrange_stores(
     num_parties: NonZero<usize>,
     min_threshold: usize,
 ) -> anyhow::Result<()> {
@@ -56,13 +53,3 @@ pub fn init_all_lagrange_stores(
         .ok();
     Ok(())
 }
-
-pub fn syndrome_decoding_z2<F: Field + From<u8>>(
-    parties: &[usize],
-    syndrome: &ShamirFieldPoly<F>,
-    threshold: usize,
-) -> Vec<F> {
-    let xs: Vec<F> = parties.iter().map(|s| F::from(*s as u8)).collect();
-    let r = parties.len() - (threshold + 1);
-    decode_syndrome(syndrome, &xs, r)
-}
diff --git a/core/threshold-algebra/src/galois_fields/mod.rs b/core/threshold-algebra/src/galois_fields/mod.rs
index ed08365e7..c6b704973 100644
--- a/core/threshold-algebra/src/galois_fields/mod.rs
+++ b/core/threshold-algebra/src/galois_fields/mod.rs
@@ -9,7 +9,6 @@
 ))]
 pub(crate) type LagrangeMap<F> = std::collections::HashMap<Vec<F>, Vec<crate::poly::Poly<F>>>;
 
-pub mod common;
 #[cfg(feature = "extension_degree_7")]
 pub mod gf128;
 #[cfg(feature = "extension_degree_4")]
@@ -22,3 +21,12 @@ pub mod gf32;
 pub mod gf64;
 #[cfg(feature = "extension_degree_3")]
 pub mod gf8;
+#[cfg(any(
+    feature = "extension_degree_3",
+    feature = "extension_degree_4",
+    feature = "extension_degree_5",
+    feature = "extension_degree_6",
+    feature = "extension_degree_7",
+    feature = "extension_degree_8",
+))]
+pub mod lagrange;
diff --git a/core/threshold-algebra/src/galois_rings/common.rs b/core/threshold-algebra/src/galois_rings/common.rs
index 3938d7145..4f893428d 100644
--- a/core/threshold-algebra/src/galois_rings/common.rs
+++ b/core/threshold-algebra/src/galois_rings/common.rs
@@ -7,7 +7,7 @@ use crate::{
         BaseRing, Derive, FromU128, Invert, One, QuotientMaximalIdeal, Ring,
         RingWithExceptionalSequence, Sample, Solve, Solve1, Syndrome, ZConsts, Zero,
     },
-    syndrome::lagrange_numerators,
+    syndrome::{lagrange_numerators, syndrome_decoding_z2},
 };
 use error_utils::anyhow_error_and_log;
 
@@ -17,7 +17,6 @@ use threshold_types::role::Role;
 
 use crate::{
     base_ring::{Z64, Z128},
-    galois_fields::common::syndrome_decoding_z2,
     sharing::shamir::{ShamirFieldPoly, ShamirSharings},
     sharing::share::Share,
 };
diff --git a/core/threshold-algebra/src/syndrome.rs b/core/threshold-algebra/src/syndrome.rs
index 3de9badba..784ba2e4f 100644
--- a/core/threshold-algebra/src/syndrome.rs
+++ b/core/threshold-algebra/src/syndrome.rs
@@ -1,6 +1,7 @@
 use super::{
     bivariate::compute_powers_list,
     poly::Poly,
+    sharing::shamir::ShamirFieldPoly,
     structure_traits::{Field, Ring},
 };
 use std::ops::Neg;
@@ -167,6 +168,17 @@ pub fn decode_syndrome<F: Field>(syndrome: &Poly<F>, x_alpha: &[F], r: usize) ->
     e
 }
 
+/// Decodes a syndrome over a binary extension field using party indices as evaluation points.
+pub fn syndrome_decoding_z2<F: Field + From<u8>>(
+    parties: &[usize],
+    syndrome: &ShamirFieldPoly<F>,
+    threshold: usize,
+) -> Vec<F> {
+    let xs: Vec<F> = parties.iter().map(|s| F::from(*s as u8)).collect();
+    let r = parties.len() - (threshold + 1);
+    decode_syndrome(syndrome, &xs, r)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;