From 476ee0ef479940bbe03c76d5bc531ed4796a1a37 Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Wed, 28 May 2025 14:55:31 +0200
Subject: [PATCH] Added tmp_mat_znx_dft to Scratch

---
 backend/src/lib.rs         | 18 ++++++++++++++++++
 backend/src/mat_znx_dft.rs | 14 ++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/backend/src/lib.rs b/backend/src/lib.rs
index 89a52ef..9e71c8d 100644
--- a/backend/src/lib.rs
+++ b/backend/src/lib.rs
@@ -248,4 +248,22 @@ impl Scratch {
             Self::new(rem_slice),
         )
     }
+
+    pub fn tmp_mat_znx_dft<B: Backend>(
+        &mut self,
+        module: &Module<B>,
+        rows: usize,
+        cols_in: usize,
+        cols_out: usize,
+        size: usize,
+    ) -> (MatZnxDft<&mut [u8], B>, &mut Self) {
+        let (take_slice, rem_slice) = Self::take_slice_aligned(
+            &mut self.data,
+            module.bytes_of_mat_znx_dft(rows, cols_in, cols_out, size),
+        );
+        (
+            MatZnxDft::from_data(take_slice, module.n(), rows, cols_in, cols_out, size),
+            Self::new(rem_slice),
+        )
+    }
 }
diff --git a/backend/src/mat_znx_dft.rs b/backend/src/mat_znx_dft.rs
index a717f4b..e9d6737 100644
--- a/backend/src/mat_znx_dft.rs
+++ b/backend/src/mat_znx_dft.rs
@@ -198,3 +198,17 @@ where
         }
     }
 }
+
+impl<D, B: Backend> MatZnxDft<D, B> {
+    pub(crate) fn from_data(data: D, n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> Self {
+        Self {
+            data,
+            n,
+            rows,
+            cols_in,
+            cols_out,
+            size,
+            _phantom: PhantomData,
+        }
+    }
+}