From a12e42681a7966741c05ac1de14a5c342ff0c07c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Bourdeauducq?= <sb@m-labs.hk>
Date: Fri, 4 Oct 2024 18:59:46 +0800
Subject: [PATCH] add FFT function

---
 Cargo.lock  | 60 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 Cargo.toml  |  2 ++
 src/main.rs | 41 ++++++++++++++++++++++++++++++++++++
 3 files changed, 103 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index d691216..c7b630b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -698,7 +698,9 @@ dependencies = [
  "egui_plot",
  "libloading 0.8.5",
  "nac3core",
+ "num-complex",
  "parking_lot",
+ "rustfft",
  "tempfile",
 ]
 
@@ -2213,12 +2215,30 @@ version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2bf50223579dc7cdcfb3bfcacf7069ff68243f8c363f62ffa99cf000a6b9c451"
 
+[[package]]
+name = "num-complex"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "num-conv"
 version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9"
 
+[[package]]
+name = "num-integer"
+version = "0.1.46"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "num-traits"
 version = "0.2.19"
@@ -2673,6 +2693,15 @@ version = "0.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e8cf8e6a8aa66ce33f63993ffc4ea4271eb5b0530a9002db8455ea6050c77bfa"
 
+[[package]]
+name = "primal-check"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc0d895b311e3af9902528fbb8f928688abbd95872819320517cc24ca6b2bd08"
+dependencies = [
+ "num-integer",
+]
+
 [[package]]
 name = "proc-macro-crate"
 version = "1.3.1"
@@ -2864,6 +2893,21 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
 
+[[package]]
+name = "rustfft"
+version = "6.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "43806561bc506d0c5d160643ad742e3161049ac01027b5e6d7524091fd401d86"
+dependencies = [
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "primal-check",
+ "strength_reduce",
+ "transpose",
+ "version_check",
+]
+
 [[package]]
 name = "rustix"
 version = "0.37.27"
@@ -3153,6 +3197,12 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
 
+[[package]]
+name = "strength_reduce"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fe895eb47f22e2ddd4dabc02bce419d2e643c8e3b585c78158b349195bc24d82"
+
 [[package]]
 name = "strict-num"
 version = "0.1.1"
@@ -3428,6 +3478,16 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "transpose"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1ad61aed86bc3faea4300c7aee358b4c6d0c8d6ccc36524c96e4c92ccf26e77e"
+dependencies = [
+ "num-integer",
+ "strength_reduce",
+]
+
 [[package]]
 name = "ttf-parser"
 version = "0.24.1"
diff --git a/Cargo.toml b/Cargo.toml
index 67b84df..967a4b7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -10,6 +10,8 @@ egui_plot = "0.28"
 parking_lot = "0.12"
 tempfile = "3.12"
 libloading = "0.8"
+num-complex = "0.4"
+rustfft = "6.2"
 
 [dependencies.nac3core]
 git = "https://git.m-labs.hk/M-Labs/nac3"
diff --git a/src/main.rs b/src/main.rs
index ec89798..d8ad52f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -347,6 +347,30 @@ fn compile(code: &String, run_symbol: &String, output_filename: &Path) -> Result
             vars: typedef::VarMap::new(),
         },
     );
+    register_cells_function(
+        &mut composer,
+        resolver.clone(),
+        internal_resolver.as_ref(),
+        "fft",
+        typedef::FunSignature {
+            args: vec![
+                typedef::FuncArg {
+                    name: "reals".into(),
+                    ty: list_float,
+                    default_value: None,
+                    is_vararg: false,
+                },
+                typedef::FuncArg {
+                    name: "imags".into(),
+                    ty: list_float,
+                    default_value: None,
+                    is_vararg: false,
+                },
+            ],
+            ret: primitive.none,
+            vars: typedef::VarMap::new(),
+        },
+    );
 
     let parser_result =
         match nac3parser::parser::parse_program(code.as_str(), String::from("cell1").into()) {
@@ -658,6 +682,23 @@ pub extern "C" fn __nac3_cells_plot(data: *const List) {
         .show(ui, |plot_ui| plot_ui.line(line));
 }
 
+#[no_mangle]
+pub extern "C" fn __nac3_cells_fft(reals: *mut List, imags: *mut List) {
+    let reals_slice = unsafe { std::slice::from_raw_parts_mut((*reals).data, (*reals).length) };
+    let imags_slice = unsafe { std::slice::from_raw_parts_mut((*imags).data, (*imags).length) };
+    let mut planner = rustfft::FftPlanner::<f64>::new();
+    let fft = planner.plan_fft_forward(reals_slice.len());
+    let mut buffer = vec![num_complex::Complex { re: 0.0, im: 0.0 }; reals_slice.len()];
+    for (buf, (re, im)) in buffer.iter_mut().zip(reals_slice.iter().zip(imags_slice.iter())) {
+        *buf = num_complex::Complex { re: *re, im: *im };
+    }
+    fft.process(&mut buffer);
+    for (buf, (re, im)) in buffer.iter().zip(reals_slice.iter_mut().zip(imags_slice.iter_mut())) {
+        *re = buf.re;
+        *im = buf.im;
+    }
+}
+
 impl Cells {
     fn new() -> Self {
         Self {