From ff15061f86b94c52f5e4c1a810b6cf6a46625e3a Mon Sep 17 00:00:00 2001
From: Matt Walker <matt.g.d.walker@gmail.com>
Date: Thu, 26 Jun 2025 15:44:14 -0400
Subject: [PATCH] fix: Update proof-systems with memory-safe Lagrange
 commitment functions

Merge branch 'master' into dw/bumping-up-wasm-bindgen and add memory leak fixes
for Lagrange commitment operations that prevent memory accumulation during proof
generation processes.
---
 plonk-wasm/src/rayon.rs |  2 ++
 plonk-wasm/src/srs.rs   | 49 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/plonk-wasm/src/rayon.rs b/plonk-wasm/src/rayon.rs
index 5127f1e726..a6e9f591ed 100644
--- a/plonk-wasm/src/rayon.rs
+++ b/plonk-wasm/src/rayon.rs
@@ -18,6 +18,8 @@ use wasm_bindgen::prelude::*;
 #[cfg(feature = "nodejs")]
 use js_sys::JsString;
 
+// IMPORTANT: This thread pool must be cleaned up by calling exit_thread_pool() from JavaScript
+// when the WASM module is being disposed. Otherwise, the threads will leak memory.
 static mut THREAD_POOL: Option<rayon::ThreadPool> = None;
 
 pub fn run_in_pool<OP, R>(op: OP) -> R
diff --git a/plonk-wasm/src/srs.rs b/plonk-wasm/src/srs.rs
index 4a38717f76..cc51cc45c3 100644
--- a/plonk-wasm/src/srs.rs
+++ b/plonk-wasm/src/srs.rs
@@ -142,7 +142,51 @@ macro_rules! impl_srs {
                 Box::into_raw(boxed_comm)
             }
 
+            /// Direct version that doesn't use raw pointers, for single-threaded use.
+            #[wasm_bindgen]
+            pub fn [<$name:snake _lagrange_commitments_whole_domain>](
+                srs: &[<Wasm $field_name:camel Srs>],
+                domain_size: i32,
+            ) -> WasmVector<$WasmPolyComm> {
+                srs.get_lagrange_basis_from_domain_size(domain_size as usize)
+                    .clone()
+                    .into_iter()
+                    .map(|x| x.into())
+                    .collect()
+            }
+
+            /// Takes ownership of the lagrange commitments from a raw pointer.
+            /// This properly deallocates the memory and should be used instead of read_from_ptr.
+            ///
+            /// # Safety
+            ///
+            /// This function is unsafe because it dereferences a raw pointer.
+            /// The pointer must be valid and must not be used after this call.
+            #[wasm_bindgen]
+            pub unsafe fn [<$name:snake _lagrange_commitments_whole_domain_take>](
+                ptr: *mut WasmVector<$WasmPolyComm>,
+            ) -> WasmVector<$WasmPolyComm> {
+                // Take ownership of the Box and move out its contents
+                let b = unsafe { Box::from_raw(ptr) };
+                *b  // Move the value out, Box is properly dropped
+            }
+
+            /// Frees the memory allocated for lagrange commitments without returning the data.
+            ///
+            /// # Safety
+            ///
+            /// This function is unsafe because it dereferences a raw pointer.
+            /// The pointer must be valid and must not be used after this call.
+            #[wasm_bindgen]
+            pub unsafe fn [<$name:snake _lagrange_commitments_whole_domain_free>](
+                ptr: *mut WasmVector<$WasmPolyComm>,
+            ) {
+                // Reconstruct and immediately drop the Box to free memory
+                let _ = unsafe { Box::from_raw(ptr) };
+            }
+
             /// Reads the lagrange commitments from a raw pointer.
+            /// DEPRECATED: This function has undefined behavior. Use lagrange_commitments_whole_domain_take instead.
             ///
             /// # Safety
             ///
@@ -154,8 +198,11 @@ macro_rules! impl_srs {
             ) -> WasmVector<$WasmPolyComm> {
                 // read the commitment at the pointers address, hack for the web
                 // worker implementation (see o1js web worker impl for details)
+                // WARNING: This implementation has been fixed to avoid use-after-free
                 let b = unsafe { Box::from_raw(ptr) };
-                b.as_ref().clone()
+                let data = b.as_ref().clone();
+                Box::into_raw(b); // Re-leak to avoid double-free for backward compatibility
+                data
             }
 
             #[wasm_bindgen]