Enable wav_dump for beamforming

BUG=b:340217962
FIXED=b:340217962
TEST=bazel test //...

Change-Id: I83bc31f2e169a2041e8b6ad5a93dad7c7cd78f05
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/adhd/+/5535080
Commit-Queue: Li-Yu Yu <aaronyu@google.com>
Tested-by: Li-Yu Yu <aaronyu@google.com>
Reviewed-by: Chih-Yang Hsia <paulhsia@chromium.org>
Tested-by: chromeos-cop-builder@chromeos-cop.iam.gserviceaccount.com <chromeos-cop-builder@chromeos-cop.iam.gserviceaccount.com>

2 files changed, 112 insertions, 4 deletions

diff --git a/audio_processor/src/processors/plugin/loader.rs b/audio_processor/src/processors/plugin/loader.rs
index 0ceb9717..42151cca 100644
--- a/audio_processor/src/processors/plugin/loader.rs
+++ b/audio_processor/src/processors/plugin/loader.rs
@@ -2,6 +2,8 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+use std::path::PathBuf;
+
 use anyhow::anyhow;
 use anyhow::Context;
 
@@ -9,9 +11,18 @@ use crate::processors::ChunkWrapper;
 use crate::processors::DynamicPluginProcessor;
 use crate::processors::SpeexResampler;
 use crate::AudioProcessor;
+use crate::Pipeline;
 use crate::ProcessorVec;
 use crate::Shape;
 
+/// Configuration for plugin dump.
+pub struct PluginDumpConfig {
+    /// Path to store the pre-processing WAVE dump.
+    pub pre_processing_wav_dump: PathBuf,
+    /// Path to store the post-processing WAVE dump.
+    pub post_processing_wav_dump: PathBuf,
+}
+
 // TODO(aaronyu): This is the wrong abstraction. We should have an API to
 // chain processors with incompatible shapes and rates easily instead,
 // when we have something that looks more like a graph.
@@ -24,6 +35,7 @@ pub struct PluginLoader<'a> {
     pub outer_block_size: usize,
     pub inner_block_size: usize,
     pub allow_chunk_wrapper: bool,
+    pub dump_config: Option<&'a PluginDumpConfig>,
 }
 
 impl<'a> PluginLoader<'a> {
@@ -37,6 +49,26 @@ impl<'a> PluginLoader<'a> {
         )
         .with_context(|| "DynamicPluginProcessor::new failed")?;
 
+        let mut maybe_dump_processor: ProcessorVec = vec![];
+        match self.dump_config {
+            Some(dump_config) => {
+                maybe_dump_processor.add_wav_dump(
+                    &dump_config.pre_processing_wav_dump,
+                    self.channels,
+                    self.inner_rate,
+                )?;
+                maybe_dump_processor.add(processor);
+                maybe_dump_processor.add_wav_dump(
+                    &dump_config.post_processing_wav_dump,
+                    self.channels,
+                    self.inner_rate,
+                )?;
+            }
+            None => {
+                maybe_dump_processor.add(processor);
+            }
+        }
+
         let maybe_wrapped_processor: Box<dyn AudioProcessor<I = f32, O = f32> + Send> = if self
             .outer_block_size
             * self.inner_rate
@@ -47,13 +79,13 @@ impl<'a> PluginLoader<'a> {
                 return Err(anyhow!("ChunkWrapper is not allowed but required: outer rate={}, block_size={}; inner rate={}, block_size={}", self.outer_rate, self.outer_block_size, self.inner_rate, self.inner_block_size));
             }
             Box::new(ChunkWrapper::new(
-                processor,
+                maybe_dump_processor,
                 self.inner_block_size,
                 self.channels,
                 self.channels,
             ))
         } else {
-            Box::new(processor)
+            Box::new(maybe_dump_processor)
         };
 
         let processors = if self.outer_rate == self.inner_rate {
@@ -89,3 +121,63 @@ impl<'a> PluginLoader<'a> {
         Ok(processors)
     }
 }
+
+#[cfg(feature = "bazel")]
+#[cfg(test)]
+mod tests {
+    use std::env;
+
+    use tempfile::TempDir;
+
+    use super::PluginDumpConfig;
+    use super::PluginLoader;
+    use crate::util::read_wav;
+    use crate::AudioProcessor;
+    use crate::MultiBuffer;
+
+    #[test]
+    fn test_wav_dump() {
+        let plugin_path = env::var("LIBTEST_PLUGINS_SO").unwrap();
+
+        let tmpd = TempDir::new().unwrap();
+
+        let dumps = PluginDumpConfig {
+            pre_processing_wav_dump: tmpd.path().join("pre.wav"),
+            post_processing_wav_dump: tmpd.path().join("post.wav"),
+        };
+        let mut pipeline = PluginLoader {
+            path: &plugin_path,
+            constructor: "negate_processor_create",
+            channels: 2,
+            outer_rate: 16000,
+            inner_rate: 16000,
+            outer_block_size: 3,
+            inner_block_size: 4,
+            allow_chunk_wrapper: true,
+            dump_config: Some(&dumps),
+        }
+        .load_and_wrap()
+        .unwrap();
+
+        let mut input = MultiBuffer::from(vec![vec![1f32, 2., 3.], vec![4., 5., 6.]]);
+        let output = pipeline.process(input.as_multi_slice()).unwrap();
+        assert_eq!(output.into_raw(), [[0.; 3]; 2]);
+
+        let mut input = MultiBuffer::from(vec![vec![7., 8., 9.], vec![10., 11., 12.]]);
+        let output = pipeline.process(input.as_multi_slice()).unwrap();
+        assert_eq!(output.into_raw(), [[0., -1., -2.], [0., -4., -5.]]);
+
+        // Flush wav buffers.
+        drop(pipeline);
+
+        let (pre_spec, pre_wav) = read_wav::<f32>(&dumps.pre_processing_wav_dump).unwrap();
+        let (post_spec, post_wav) = read_wav::<f32>(&dumps.post_processing_wav_dump).unwrap();
+        assert_eq!(pre_spec.sample_rate, 16000);
+        assert_eq!(post_spec.sample_rate, 16000);
+        assert_eq!(pre_wav.to_vecs(), [[1., 2., 3., 7.], [4., 5., 6., 10.]]);
+        assert_eq!(
+            post_wav.to_vecs(),
+            [[-1.0, -2.0, -3.0, -7.0], [-4.0, -5.0, -6.0, -10.0]]
+        );
+    }
+}
diff --git a/cras/src/server/rust/src/cras_processor.rs b/cras/src/server/rust/src/cras_processor.rs
index 563612ba..12985ea4 100644
--- a/cras/src/server/rust/src/cras_processor.rs
+++ b/cras/src/server/rust/src/cras_processor.rs
@@ -14,6 +14,7 @@ use audio_processor::processors::binding::plugin_processor;
 use audio_processor::processors::export_plugin;
 use audio_processor::processors::CheckShape;
 use audio_processor::processors::NegateAudioProcessor;
+use audio_processor::processors::PluginDumpConfig;
 use audio_processor::processors::PluginLoader;
 use audio_processor::processors::PluginProcessor;
 use audio_processor::processors::ShuffleChannels;
@@ -112,6 +113,7 @@ fn create_noise_cancellation_pipeline(
         outer_block_size: config.block_size,
         inner_block_size: 480,
         allow_chunk_wrapper: false,
+        dump_config: None,
     }
     .load_and_wrap()
 }
@@ -139,11 +141,15 @@ fn create_style_transfer_pipeline(config: &CrasProcessorConfig) -> anyhow::Resul
         outer_block_size: config.block_size,
         inner_block_size: 480,
         allow_chunk_wrapper: true,
+        dump_config: None,
     }
     .load_and_wrap()
 }
 
-fn create_beamforming_pipeline(config: &CrasProcessorConfig) -> anyhow::Result<ProcessorVec> {
+fn create_beamforming_pipeline(
+    config: &CrasProcessorConfig,
+    dump_config: Option<&PluginDumpConfig>,
+) -> anyhow::Result<ProcessorVec> {
     // Check shape is supported.
     if config.channels != 3 {
         bail!("unsupported channel count {}", config.channels);
@@ -159,6 +165,7 @@ fn create_beamforming_pipeline(config: &CrasProcessorConfig) -> anyhow::Result<P
         outer_block_size: config.block_size,
         inner_block_size: 256,
         allow_chunk_wrapper: true,
+        dump_config,
     }
     .load_and_wrap()
 }
@@ -256,8 +263,16 @@ impl CrasProcessor {
                 }
             }
             CrasProcessorEffect::Beamforming => {
+                let dump_config = if config.wav_dump {
+                    Some(PluginDumpConfig {
+                        pre_processing_wav_dump: dump_base.join("pre_beamforming.wav"),
+                        post_processing_wav_dump: dump_base.join("post_beamforming.wav"),
+                    })
+                } else {
+                    None
+                };
                 pipeline.extend(
-                    create_beamforming_pipeline(&config)
+                    create_beamforming_pipeline(&config, dump_config.as_ref())
                         .context("failed when creating beamforming pipeline")?,
                 );
             }
@@ -278,6 +293,7 @@ impl CrasProcessor {
                             block_size => block_size as usize,
                         },
                         allow_chunk_wrapper: true,
+                        dump_config: None,
                     }
                     .load_and_wrap()?,
                 );