1use crate::error::{FFTError, FFTResult};
8use rustfft::FftPlanner;
9use scirs2_core::numeric::Complex64;
10use std::collections::HashMap;
11use std::sync::{Arc, Mutex, OnceLock};
12
13pub trait FftBackend: Send + Sync {
15 fn name(&self) -> &str;
17
18 fn description(&self) -> &str;
20
21 fn is_available(&self) -> bool;
23
24 fn fft(&self, input: &[Complex64], output: &mut [Complex64]) -> FFTResult<()>;
26
27 fn ifft(&self, input: &[Complex64], output: &mut [Complex64]) -> FFTResult<()>;
29
30 fn fft_sized(
32 &self,
33 input: &[Complex64],
34 output: &mut [Complex64],
35 size: usize,
36 ) -> FFTResult<()>;
37
38 fn ifft_sized(
40 &self,
41 input: &[Complex64],
42 output: &mut [Complex64],
43 size: usize,
44 ) -> FFTResult<()>;
45
46 fn supports_feature(&self, feature: &str) -> bool;
48}
49
50pub struct RustFftBackend {
52 planner: Arc<Mutex<FftPlanner<f64>>>,
53}
54
55impl RustFftBackend {
56 pub fn new() -> Self {
57 Self {
58 planner: Arc::new(Mutex::new(FftPlanner::new())),
59 }
60 }
61}
62
63impl Default for RustFftBackend {
64 fn default() -> Self {
65 Self::new()
66 }
67}
68
69impl FftBackend for RustFftBackend {
70 fn name(&self) -> &str {
71 "rustfft"
72 }
73
74 fn description(&self) -> &str {
75 "Pure Rust FFT implementation using RustFFT library"
76 }
77
78 fn is_available(&self) -> bool {
79 true
80 }
81
82 fn fft(&self, input: &[Complex64], output: &mut [Complex64]) -> FFTResult<()> {
83 self.fft_sized(input, output, input.len())
84 }
85
86 fn ifft(&self, input: &[Complex64], output: &mut [Complex64]) -> FFTResult<()> {
87 self.ifft_sized(input, output, input.len())
88 }
89
90 fn fft_sized(
91 &self,
92 input: &[Complex64],
93 output: &mut [Complex64],
94 size: usize,
95 ) -> FFTResult<()> {
96 if input.len() != size || output.len() != size {
97 return Err(FFTError::ValueError(
98 "Input and output sizes must match the specified size".to_string(),
99 ));
100 }
101
102 let mut planner = self.planner.lock().unwrap();
104 let fft = planner.plan_fft_forward(size);
105
106 let mut buffer: Vec<rustfft::num_complex::Complex<f64>> = input
108 .iter()
109 .map(|&c| rustfft::num_complex::Complex::new(c.re, c.im))
110 .collect();
111
112 fft.process(&mut buffer);
114
115 for (i, &c) in buffer.iter().enumerate() {
117 output[i] = Complex64::new(c.re, c.im);
118 }
119
120 Ok(())
121 }
122
123 fn ifft_sized(
124 &self,
125 input: &[Complex64],
126 output: &mut [Complex64],
127 size: usize,
128 ) -> FFTResult<()> {
129 if input.len() != size || output.len() != size {
130 return Err(FFTError::ValueError(
131 "Input and output sizes must match the specified size".to_string(),
132 ));
133 }
134
135 let mut planner = self.planner.lock().unwrap();
137 let fft = planner.plan_fft_inverse(size);
138
139 let mut buffer: Vec<rustfft::num_complex::Complex<f64>> = input
141 .iter()
142 .map(|&c| rustfft::num_complex::Complex::new(c.re, c.im))
143 .collect();
144
145 fft.process(&mut buffer);
147
148 let scale = 1.0 / size as f64;
150 for (i, &c) in buffer.iter().enumerate() {
151 output[i] = Complex64::new(c.re * scale, c.im * scale);
152 }
153
154 Ok(())
155 }
156
157 fn supports_feature(&self, feature: &str) -> bool {
158 matches!(feature, "1d_fft" | "2d_fft" | "nd_fft" | "cached_plans")
159 }
160}
161
162pub struct BackendManager {
164 backends: Arc<Mutex<HashMap<String, Arc<dyn FftBackend>>>>,
165 current_backend: Arc<Mutex<String>>,
166}
167
168impl BackendManager {
169 pub fn new() -> Self {
171 let mut backends = HashMap::new();
172
173 let rustfft_backend = Arc::new(RustFftBackend::new()) as Arc<dyn FftBackend>;
175 backends.insert("rustfft".to_string(), rustfft_backend);
176
177 Self {
178 backends: Arc::new(Mutex::new(backends)),
179 current_backend: Arc::new(Mutex::new("rustfft".to_string())),
180 }
181 }
182
183 pub fn register_backend(&self, name: String, backend: Arc<dyn FftBackend>) -> FFTResult<()> {
185 let mut backends = self.backends.lock().unwrap();
186 if backends.contains_key(&name) {
187 return Err(FFTError::ValueError(format!(
188 "Backend '{name}' already exists"
189 )));
190 }
191 backends.insert(name, backend);
192 Ok(())
193 }
194
195 pub fn list_backends(&self) -> Vec<String> {
197 let backends = self.backends.lock().unwrap();
198 backends.keys().cloned().collect()
199 }
200
201 pub fn set_backend(&self, name: &str) -> FFTResult<()> {
203 let backends = self.backends.lock().unwrap();
204 if !backends.contains_key(name) {
205 return Err(FFTError::ValueError(format!("Backend '{name}' not found")));
206 }
207
208 if let Some(backend) = backends.get(name) {
210 if !backend.is_available() {
211 return Err(FFTError::ValueError(format!(
212 "Backend '{name}' is not available"
213 )));
214 }
215 }
216
217 *self.current_backend.lock().unwrap() = name.to_string();
218 Ok(())
219 }
220
221 pub fn get_backend_name(&self) -> String {
223 self.current_backend.lock().unwrap().clone()
224 }
225
226 pub fn get_backend(&self) -> Arc<dyn FftBackend> {
228 let current_name = self.current_backend.lock().unwrap();
229 let backends = self.backends.lock().unwrap();
230 backends
231 .get(&*current_name)
232 .cloned()
233 .expect("Current backend should always exist")
234 }
235
236 pub fn get_backend_info(&self, name: &str) -> Option<BackendInfo> {
238 let backends = self.backends.lock().unwrap();
239 backends.get(name).map(|backend| BackendInfo {
240 name: backend.name().to_string(),
241 description: backend.description().to_string(),
242 available: backend.is_available(),
243 })
244 }
245}
246
247impl Default for BackendManager {
248 fn default() -> Self {
249 Self::new()
250 }
251}
252
253#[derive(Debug, Clone)]
255pub struct BackendInfo {
256 pub name: String,
257 pub description: String,
258 pub available: bool,
259}
260
261impl std::fmt::Display for BackendInfo {
262 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
263 write!(
264 f,
265 "{} - {} ({})",
266 self.name,
267 self.description,
268 if self.available {
269 "available"
270 } else {
271 "not available"
272 }
273 )
274 }
275}
276
277static GLOBAL_BACKEND_MANAGER: OnceLock<BackendManager> = OnceLock::new();
279
280#[allow(dead_code)]
282pub fn get_backend_manager() -> &'static BackendManager {
283 GLOBAL_BACKEND_MANAGER.get_or_init(BackendManager::new)
284}
285
286#[allow(dead_code)]
288pub fn init_backend_manager(manager: BackendManager) -> Result<(), &'static str> {
289 GLOBAL_BACKEND_MANAGER
290 .set(manager)
291 .map_err(|_| "Global backend _manager already initialized")
292}
293
294#[allow(dead_code)]
296pub fn list_backends() -> Vec<String> {
297 get_backend_manager().list_backends()
298}
299
300#[allow(dead_code)]
302pub fn set_backend(name: &str) -> FFTResult<()> {
303 get_backend_manager().set_backend(name)
304}
305
306#[allow(dead_code)]
308pub fn get_backend_name() -> String {
309 get_backend_manager().get_backend_name()
310}
311
312#[allow(dead_code)]
314pub fn get_backend_info(name: &str) -> Option<BackendInfo> {
315 get_backend_manager().get_backend_info(name)
316}
317
318pub struct BackendContext {
320 previous_backend: String,
321 manager: &'static BackendManager,
322}
323
324impl BackendContext {
325 pub fn new(_backendname: &str) -> FFTResult<Self> {
327 let manager = get_backend_manager();
328 let previous_backend = manager.get_backend_name();
329
330 manager.set_backend(_backendname)?;
332
333 Ok(Self {
334 previous_backend,
335 manager,
336 })
337 }
338}
339
340impl Drop for BackendContext {
341 fn drop(&mut self) {
342 let _ = self.manager.set_backend(&self.previous_backend);
344 }
345}
346
347#[cfg(test)]
348mod tests {
349 use super::*;
350
351 #[test]
352 fn test_rustfft_backend() {
353 let backend = RustFftBackend::new();
354 assert_eq!(backend.name(), "rustfft");
355 assert!(backend.is_available());
356 assert!(backend.supports_feature("1d_fft"));
357 }
358
359 #[test]
360 fn test_backend_manager() {
361 let manager = BackendManager::new();
362
363 assert_eq!(manager.get_backend_name(), "rustfft");
365
366 let backends = manager.list_backends();
368 assert!(backends.contains(&"rustfft".to_string()));
369
370 let info = manager.get_backend_info("rustfft").unwrap();
372 assert!(info.available);
373 }
374
375 #[test]
376 fn test_backend_context() {
377 let manager = get_backend_manager();
378 let original = manager.get_backend_name();
379
380 {
381 let _ctx = BackendContext::new("rustfft").unwrap();
382 assert_eq!(manager.get_backend_name(), "rustfft");
383 }
384
385 assert_eq!(manager.get_backend_name(), original);
387 }
388}