From f63848aa0ee139d4b587303e4ea24540d804d7c4 Mon Sep 17 00:00:00 2001
From: Rutherther <rutherther@ditigal.xyz>
Date: Fri, 31 Oct 2025 14:46:57 +0100
Subject: [PATCH] refactor(lib): add methods to get wrapped perturbations
 instead of traits.

Remove WrapperPerturbation and ListWrapperPerturbation in favor of
wrapped and wrapped_mut functions on PerturbationOperator. This
allows for easier handling as it's not easy to downcast Any into
WrapperPerturbation or ListWrapperPerturbation. So instead
always return dyn PerturbationOperator instead of dyn Any.
Use dyn Any only for downstasting to concrete types.

This finally allows for apply_to_mutations to be implemented
under MutationPerturbation, allowing to modify the arguments.
---
 codes/eoa_lib/src/perturbation/mod.rs | 109 ++++++++++++++------------
 1 file changed, 57 insertions(+), 52 deletions(-)

diff --git a/codes/eoa_lib/src/perturbation/mod.rs b/codes/eoa_lib/src/perturbation/mod.rs
index 730e6564675a968ee8edcd3de951f29f4c3b8b8c..501044e98082fbde6bf472e6cee9c5e23aa260f7 100644
--- a/codes/eoa_lib/src/perturbation/mod.rs
+++ b/codes/eoa_lib/src/perturbation/mod.rs
@@ -11,10 +11,30 @@ pub trait AnyPerturbationOperator: Any {
     fn as_any_mut(&mut self) -> &mut dyn Any;
 }
 
+pub enum Wrapped<'a, T> {
+    Single,
+    Wrapped(&'a dyn PerturbationOperator<Chromosome = T>),
+    ListWrapped(Vec<&'a dyn PerturbationOperator<Chromosome = T>>),
+}
+
+pub enum WrappedMut<'a, T> {
+    Single,
+    Wrapped(&'a mut dyn PerturbationOperator<Chromosome = T>),
+    ListWrapped(Vec<&'a mut dyn PerturbationOperator<Chromosome = T>>),
+}
+
 pub trait PerturbationOperator: AnyPerturbationOperator {
     type Chromosome;
 
     fn perturb(&self, chromosome: &mut Self::Chromosome, rng: &mut dyn RngCore);
+
+    fn wrapped(&self) -> Wrapped<'_, Self::Chromosome> {
+        Wrapped::Single
+    }
+
+    fn wrapped_mut(&mut self) -> WrappedMut<'_, Self::Chromosome> {
+        WrappedMut::Single
+    }
 }
 
 impl<T: PerturbationOperator> AnyPerturbationOperator for T
@@ -24,44 +44,11 @@ where
     fn as_any(&self) -> &dyn Any {
         self
     }
-
     fn as_any_mut(&mut self) -> &mut dyn Any {
         self
     }
 }
 
-pub trait WrapperPerturbation: PerturbationOperator {
-    fn borrow_any(&self) -> &dyn Any;
-    fn borrow_any_mut(&mut self) -> &mut dyn Any;
-
-    fn borrow<TPerturbation: 'static>(&self) -> Option<&TPerturbation> {
-        self.borrow_any().downcast_ref()
-    }
-
-    fn borrow_mut<TPerturbation: 'static>(&mut self) -> Option<&mut TPerturbation> {
-        self.borrow_any_mut().downcast_mut()
-    }
-}
-
-pub trait ListWrapperPerturbation: PerturbationOperator {
-    fn borrow_list_any<'a>(&'a self) -> impl Iterator<Item = &'a dyn Any>;
-    fn borrow_list_any_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut dyn Any>;
-
-    fn borrow<'a, T: 'static>(&'a self, index: usize) -> Option<&'a T> {
-        self.borrow_list_any()
-            .nth(index)
-            .map(|perturbation| perturbation.downcast_ref())
-            .flatten()
-    }
-
-    fn borrow_mut<'a, T: 'static>(&'a mut self, index: usize) -> Option<&'a mut T> {
-        self.borrow_list_any_mut()
-            .nth(index)
-            .map(|perturbation| perturbation.downcast_mut())
-            .flatten()
-    }
-}
-
 pub struct IdentityPerturbation<TChromosome> {
     _phantom: PhantomData<TChromosome>
 }
@@ -316,15 +303,13 @@ where
             BoundedPerturbationStrategy::Retry(retries) => self.retry_perturb(chromosome, Some(retries), rng)
         }
     }
-}
 
-impl<const LEN: usize, T: PerturbationOperator<Chromosome = SVector<f64, LEN>> + 'static> WrapperPerturbation for BoundedPerturbation<LEN, T> {
-    fn borrow_any(&self) -> &dyn Any {
-        self.perturbation.as_any()
+    fn wrapped(&self) -> Wrapped<'_, Self::Chromosome> {
+        Wrapped::Wrapped(&self.perturbation)
     }
 
-    fn borrow_any_mut(&mut self) -> &mut dyn Any {
-        self.perturbation.as_any_mut()
+    fn wrapped_mut(&mut self) -> WrappedMut<'_, Self::Chromosome> {
+        WrappedMut::Wrapped(&mut self.perturbation)
     }
 }
 
@@ -334,13 +319,31 @@ pub struct MutationPerturbation<T> {
     probability: f64
 }
 
-impl<T> MutationPerturbation<T> {
+impl<T: 'static> MutationPerturbation<T> {
     pub fn new(perturbation: Box<dyn PerturbationOperator<Chromosome = T>>, probability: f64) -> Self {
         Self {
             perturbation,
             probability
         }
     }
+
+    fn apply_to_mutations(base_perturbation: &mut dyn PerturbationOperator<Chromosome = T>, apply: &dyn Fn(&mut Self)) {
+        if let Some(mutation) = base_perturbation.as_any_mut().downcast_mut::<MutationPerturbation<T>>() {
+            apply(mutation);
+        }
+
+        match base_perturbation.wrapped_mut() {
+            WrappedMut::Single => (),
+            WrappedMut::Wrapped(wrapped) => {
+                Self::apply_to_mutations(wrapped, apply);
+            },
+            WrappedMut::ListWrapped(wrapped) => {
+                for wrapped in wrapped {
+                    Self::apply_to_mutations(wrapped, apply);
+                }
+            }
+        };
+    }
 }
 
 impl<T: 'static> PerturbationOperator for MutationPerturbation<T> {
@@ -351,15 +354,13 @@ impl<T: 'static> PerturbationOperator for MutationPerturbation<T> {
             self.perturbation.perturb(chromosome, rng);
         }
     }
-}
 
-impl<T: 'static> WrapperPerturbation for MutationPerturbation<T> {
-    fn borrow_any(&self) -> &dyn Any {
-        self.perturbation.as_any()
+    fn wrapped(&self) -> Wrapped<'_, Self::Chromosome> {
+        Wrapped::Wrapped(self.perturbation.as_ref())
     }
 
-    fn borrow_any_mut(&mut self) -> &mut dyn Any {
-        self.perturbation.as_any_mut()
+    fn wrapped_mut(&mut self) -> WrappedMut<'_, Self::Chromosome> {
+        WrappedMut::Wrapped(self.perturbation.as_mut())
     }
 }
 
@@ -383,15 +384,19 @@ impl<T: 'static> PerturbationOperator for CombinedPerturbation<T> {
             perturbation.perturb(chromosome, rng);
         }
     }
-}
 
-impl<T: 'static> ListWrapperPerturbation for CombinedPerturbation<T> {
-    fn borrow_list_any<'a>(&'a self) -> impl Iterator<Item = &'a dyn Any> {
-        self.perturbations.iter().map(|p| p.as_any())
+    fn wrapped(&self) -> Wrapped<'_, Self::Chromosome> {
+        Wrapped::ListWrapped(
+            self.perturbations
+                .iter()
+                .map(|p| p.as_ref()).collect())
     }
 
-    fn borrow_list_any_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut dyn Any> {
-        self.perturbations.iter_mut().map(|p| p.as_any_mut())
+    fn wrapped_mut(&mut self) -> WrappedMut<'_, Self::Chromosome> {
+        WrappedMut::ListWrapped(
+            self.perturbations
+                .iter_mut()
+                .map(|p| p.as_mut()).collect())
     }
 }