From e8a9106f22c304c153420738b66da032ba9ccb39 Mon Sep 17 00:00:00 2001
From: Easton Potokar <contagon6@gmail.com>
Date: Wed, 10 Jul 2024 11:52:46 -0400
Subject: [PATCH] Updated NoiseModelSafe serialization. Blanket impl is back,
 conditioned on serial

---
 src/lib.rs            | 37 ++++++++++++++++++++++++++++
 src/noise/gaussian.rs |  3 ++-
 src/noise/mod.rs      | 56 +++++++++++++++++++++++++------------------
 src/noise/unit.rs     |  7 ++++--
 4 files changed, 77 insertions(+), 26 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index e21fb76..8867266 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -19,3 +19,40 @@ pub mod variables;
 
 #[cfg(feature = "rerun")]
 pub mod rerun;
+
+#[cfg(feature = "serde")]
+pub mod serde {
+    pub trait Tagged: serde::Serialize {
+        const TAG: &'static str;
+    }
+
+    #[macro_export]
+    macro_rules! register_typetag {
+        ($trait:path, $ty:ty) => {
+            // TODO: It'd be great if this was a blanket implementation, but
+            // I had problems getting it to run over const generics
+            impl $crate::serde::Tagged for $ty {
+                const TAG: &'static str = stringify!($ty);
+            }
+
+            typetag::__private::inventory::submit! {
+                <dyn $trait>::typetag_register(
+                    <$ty as $crate::noise::Tagged>::TAG,
+                    (|deserializer| typetag::__private::Result::Ok(
+                        typetag::__private::Box::new(
+                            typetag::__private::erased_serde::deserialize::<$ty>(deserializer)?
+                        ),
+                    )) as typetag::__private::DeserializeFn<<dyn $trait as typetag::__private::Strictest>::Object>
+                )
+            }
+        };
+    }
+}
+
+#[cfg(not(feature = "serde"))]
+pub mod serde {
+    #[macro_export]
+    macro_rules! register_typetag {
+        ($trait:path, $ty:ty) => {};
+    }
+}
diff --git a/src/noise/gaussian.rs b/src/noise/gaussian.rs
index 6264338..c9e0476 100644
--- a/src/noise/gaussian.rs
+++ b/src/noise/gaussian.rs
@@ -4,9 +4,10 @@ use super::{NoiseModel, UnitNoise};
 use crate::{
     dtype,
     linalg::{Const, Matrix, MatrixView, MatrixX, Vector, VectorView, VectorX},
+    register_noise,
 };
 
-impl_safe_noise!(
+register_noise!(
     GaussianNoise<1>,
     GaussianNoise<2>,
     GaussianNoise<3>,
diff --git a/src/noise/mod.rs b/src/noise/mod.rs
index 8184743..3b57e79 100644
--- a/src/noise/mod.rs
+++ b/src/noise/mod.rs
@@ -1,6 +1,7 @@
 use std::fmt::{Debug, Display};
 
 use crate::linalg::{DimName, MatrixX, VectorX};
+
 pub trait NoiseModel: Debug + Display {
     type Dim: DimName;
 
@@ -22,30 +23,39 @@ pub trait NoiseModelSafe: Debug + Display {
     fn whiten_mat(&self, m: MatrixX) -> MatrixX;
 }
 
+impl<
+        #[cfg(not(feature = "serde"))] T: NoiseModel,
+        #[cfg(feature = "serde")] T: NoiseModel + crate::serde::Tagged,
+    > NoiseModelSafe for T
+{
+    fn dim(&self) -> usize {
+        NoiseModel::dim(self)
+    }
+
+    fn whiten_vec(&self, v: VectorX) -> VectorX {
+        NoiseModel::whiten_vec(self, v)
+    }
+
+    fn whiten_mat(&self, m: MatrixX) -> MatrixX {
+        NoiseModel::whiten_mat(self, m)
+    }
+
+    #[doc(hidden)]
+    #[cfg(feature = "serde")]
+    fn typetag_name(&self) -> &'static str {
+        Self::TAG
+    }
+
+    #[doc(hidden)]
+    #[cfg(feature = "serde")]
+    fn typetag_deserialize(&self) {}
+}
+
 #[macro_export]
-macro_rules! impl_safe_noise {
-    ($($var:ident < $num:literal > ),* $(,)?) => {
-        use paste::paste;
-        paste!{
-            $(
-                type [<$var $num>] = $var< $num >;
-                #[cfg_attr(feature = "serde", typetag::serde)]
-                impl $crate::noise::NoiseModelSafe for [<$var $num>]{
-                    fn dim(&self) -> usize {
-                        $crate::noise::NoiseModel::dim(self)
-                    }
-
-                    fn whiten_vec(&self, v: VectorX) -> VectorX {
-                        $crate::noise::NoiseModel::whiten_vec(self, v)
-                    }
-
-                    fn whiten_mat(&self, m: MatrixX) -> MatrixX {
-                        $crate::noise::NoiseModel::whiten_mat(self, m)
-                    }
-                }
-            )*
-        }
-    };
+macro_rules! register_noise {
+    ($($ty:ty),* $(,)?) => {$(
+        $crate::register_typetag!($crate::noise::NoiseModelSafe, $ty);
+    )*};
 }
 
 mod gaussian;
diff --git a/src/noise/unit.rs b/src/noise/unit.rs
index d2ff7bc..fba8505 100644
--- a/src/noise/unit.rs
+++ b/src/noise/unit.rs
@@ -1,7 +1,10 @@
 use super::NoiseModel;
-use crate::linalg::{Const, MatrixX, VectorX};
+use crate::{
+    linalg::{Const, MatrixX, VectorX},
+    register_noise,
+};
 
-impl_safe_noise!(
+register_noise!(
     UnitNoise<1>,
     UnitNoise<2>,
     UnitNoise<3>,