From 6f4c897dc4c22ad446053153171f491445affa2e Mon Sep 17 00:00:00 2001
From: Chris Macklin <chris.macklin@10xgenomics.com>
Date: Mon, 1 Apr 2024 23:58:10 -0700
Subject: [PATCH] Allow orbit proc to pass data to itself.

---
 enclone_print/src/process_clonotypes.rs | 24 ++++++++++++++++--------
 enclone_ranger/src/main_enclone.rs      |  6 +++---
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/enclone_print/src/process_clonotypes.rs b/enclone_print/src/process_clonotypes.rs
index d4b444161..5ae6bf5c4 100644
--- a/enclone_print/src/process_clonotypes.rs
+++ b/enclone_print/src/process_clonotypes.rs
@@ -22,12 +22,12 @@ use vector_utils::{erase_if, next_diff12_3};
 /// Process clonotypes.
 /// Filter out exact subclonotypes in orbits that appear to be junk.
 /// Write out barcode fates and loupe clonotype files.
-pub fn print_clonotypes<T: Send>(
+pub fn process_clonotypes<T: Send, D: Default>(
     setup: &EncloneSetup,
     enclone_exacts: &EncloneExacts,
     gex_readers: &[Option<GexReaders<'_>>],
     fate: &[BarcodeFates],
-    mut proc: impl OrbitProcessor<T> + Send + Sync,
+    mut proc: impl OrbitProcessor<T, D> + Send + Sync,
 ) -> Result<(), String> {
     let EncloneSetup {
         ctl,
@@ -115,7 +115,7 @@ pub fn print_clonotypes<T: Send>(
         // Let n be the total number of cells in this pass.
         let n: usize = mults.iter().sum();
 
-        if n >= ctl.clono_filt_opt.ncells_low
+        let proc_filter_data = if n >= ctl.clono_filt_opt.ncells_low
             || ctl.clono_group_opt.asymmetric_center == "from_filters"
         {
             // Mark some weak exact subclonotypes for deletion.
@@ -132,8 +132,10 @@ pub fn print_clonotypes<T: Send>(
                 &rsi,
                 &mut bads,
                 true,
-            )?;
-        }
+            )?
+        } else {
+            Default::default()
+        };
 
         // Delete weak exact subclonotypes.
 
@@ -214,6 +216,7 @@ pub fn print_clonotypes<T: Send>(
             &rsi,
             &mut bads,
             in_center,
+            proc_filter_data,
         )?;
 
         Ok((num_cells, loupe_clonotype, res))
@@ -310,7 +313,11 @@ fn sort_exact_clonotypes(
 }
 
 /// Inject a behavior to provide additional filtering and post-processing of each orbit.
-pub trait OrbitProcessor<T> {
+pub trait OrbitProcessor<T, D: Default> {
+    /// Filter performs additional filtering of clonotypes by mutating bads.
+    ///
+    /// The processor may return a data structure of type D that is provided
+    /// to the finalize method.
     #[allow(unused)]
     fn filter(
         &self,
@@ -324,7 +331,7 @@ pub trait OrbitProcessor<T> {
         rsi: &ColInfo,
         bads: &mut [bool],
         in_center: bool,
-    ) -> Result<(), String> {
+    ) -> Result<D, String> {
         let ctl = &setup.ctl;
         // This assertion ensures that we never would have entered code that was
         // moved out of this repo and into enclone proper.
@@ -333,7 +340,7 @@ pub trait OrbitProcessor<T> {
                 && !ctl.gen_opt.complete
                 && ctl.gen_opt.var_def.is_empty()
         );
-        Ok(())
+        Ok(Default::default())
     }
 
     #[allow(unused)]
@@ -349,6 +356,7 @@ pub trait OrbitProcessor<T> {
         rsi: &ColInfo,
         bads: &mut [bool],
         in_center: bool,
+        filter_data: D,
     ) -> Result<Option<T>, String> {
         Ok(None)
     }
diff --git a/enclone_ranger/src/main_enclone.rs b/enclone_ranger/src/main_enclone.rs
index 610071e25..691ab4231 100644
--- a/enclone_ranger/src/main_enclone.rs
+++ b/enclone_ranger/src/main_enclone.rs
@@ -10,7 +10,7 @@ use enclone_args::load_gex::get_gex_info;
 use enclone_args::proc_args::proc_args;
 use enclone_core::defs::EncloneControl;
 use enclone_core::enclone_structs::EncloneSetup;
-use enclone_print::process_clonotypes::{print_clonotypes, OrbitProcessor};
+use enclone_print::process_clonotypes::{process_clonotypes, OrbitProcessor};
 use enclone_stuff::start::main_enclone_start;
 use std::sync::atomic::Ordering::SeqCst;
 use std::{
@@ -77,7 +77,7 @@ pub fn main_enclone_ranger(args: &[String]) -> Result<(), String> {
     let setup = main_enclone_setup_ranger(args)?;
     let (exacts, fate) = main_enclone_start(&setup)?;
     let gex_readers = setup.create_gex_readers();
-    print_clonotypes::<()>(&setup, &exacts, &gex_readers, &fate, NoOpProc)
+    process_clonotypes::<(), ()>(&setup, &exacts, &gex_readers, &fate, NoOpProc)
 }
 
 pub fn main_enclone_setup_ranger(args: &[String]) -> Result<EncloneSetup, String> {
@@ -151,4 +151,4 @@ pub fn main_enclone_setup_ranger(args: &[String]) -> Result<EncloneSetup, String
 
 struct NoOpProc;
 
-impl OrbitProcessor<()> for NoOpProc {}
+impl OrbitProcessor<(), ()> for NoOpProc {}