From a2b802eb80bc8151a1b0b4d148f409059dd5af1c Mon Sep 17 00:00:00 2001
From: "Mayeul@Zama" <69792125+mayeul-zama@users.noreply.github.com>
Date: Thu, 11 Dec 2025 16:03:16 +0100
Subject: [PATCH] refactor(core): simplify PlanMap

---
 tfhe/src/core_crypto/commons/plan.rs | 36 +++++++++++++---------------
 1 file changed, 17 insertions(+), 19 deletions(-)

diff --git a/tfhe/src/core_crypto/commons/plan.rs b/tfhe/src/core_crypto/commons/plan.rs
index ded3f0ccc..609fbc4ed 100644
--- a/tfhe/src/core_crypto/commons/plan.rs
+++ b/tfhe/src/core_crypto/commons/plan.rs
@@ -1,33 +1,31 @@
-use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::hash::Hash;
-use std::sync::{Arc, OnceLock, RwLock};
+use std::sync::{Arc, RwLock};
 
-pub type PlanMap<Key, Value> = RwLock<HashMap<Key, Arc<OnceLock<Arc<Value>>>>>;
+pub type PlanMap<Key, Value> = RwLock<HashMap<Key, Arc<Value>>>;
 
 pub fn new_from_plan_map<Key: Eq + Hash + Copy, Value>(
     values_map: &PlanMap<Key, Value>,
     key: Key,
     new_value: impl Fn(Key) -> Value,
 ) -> Arc<Value> {
-    let get_plan = || {
-        let plans = values_map.read().unwrap();
-        let plan = plans.get(&key).cloned();
-        drop(plans);
+    let values = values_map.read().unwrap();
 
-        plan.map(|p| p.get_or_init(|| Arc::new(new_value(key))).clone())
-    };
+    let value = values.get(&key).cloned();
+    drop(values);
 
-    get_plan().unwrap_or_else(|| {
-        // If we don't find a plan for the given size, we insert a new OnceLock,
-        // drop the write lock on the map and then let get_plan() initialize the OnceLock
-        // (without holding the write lock on the map).
-        let mut plans = values_map.write().unwrap();
-        if let Entry::Vacant(v) = plans.entry(key) {
-            v.insert(Arc::new(OnceLock::new()));
-        }
-        drop(plans);
+    value.unwrap_or_else(|| {
+        // If we don't find a plan for the given polynomial size and modulus, we insert a
+        // new one (if we still don't find it after getting a write lock)
 
-        get_plan().unwrap()
+        let new_value = Arc::new(new_value(key));
+
+        let mut values = values_map.write().unwrap();
+
+        let value = Arc::clone(values.entry(key).or_insert(new_value));
+
+        drop(values);
+
+        value
     })
 }