From 82924aa96ed4f0a1d9040daaf167bfb9d484949b Mon Sep 17 00:00:00 2001
From: Victor Koenders <gpg@trangar.com>
Date: Sat, 16 Oct 2021 11:03:25 +0200
Subject: [PATCH] Added support for Option<T> and Result<T, U>

---
 derive/src/derive_enum.rs        | 12 ++++----
 derive/src/generate/generator.rs |  4 +++
 src/de/impls.rs                  | 48 ++++++++++++++++++++++++++++++++
 src/enc/impls.rs                 | 33 ++++++++++++++++++++++
 src/error.rs                     |  3 ++
 tests/basic_types.rs             |  7 +++++
 6 files changed, 102 insertions(+), 5 deletions(-)

diff --git a/derive/src/derive_enum.rs b/derive/src/derive_enum.rs
index 6e06319..394ae2f 100644
--- a/derive/src/derive_enum.rs
+++ b/derive/src/derive_enum.rs
@@ -80,6 +80,7 @@ impl DeriveEnum {
 
     pub fn generate_decodable(self, generator: &mut Generator) -> Result<()> {
         let DeriveEnum { variants } = self;
+        let enum_name = generator.target_name().to_string();
 
         if generator.has_lifetimes() {
             // enum has a lifetime, implement BorrowDecodable
@@ -125,14 +126,14 @@ impl DeriveEnum {
 
                     // invalid idx
                     variant_case.push_parsed(format!(
-                        "variant => return Err(bincode::error::DecodeError::UnexpectedVariant {{ min: 0, max: {}, found: variant }})",
-                        variants.len() - 1
+                        "variant => return Err(bincode::error::DecodeError::UnexpectedVariant {{ min: 0, max: {}, found: variant, type_name: {:?} }})",
+                        variants.len() - 1,
+                        enum_name.to_string()
                     ));
                 });
             });
         } else {
             // enum has no lifetimes, implement Decodable
-
             generator.impl_for("bincode::de::Decodable")
                 .generate_fn("decode")
                 .with_generic("D", ["bincode::de::Decode"])
@@ -175,8 +176,9 @@ impl DeriveEnum {
 
                 // invalid idx
                 variant_case.push_parsed(format!(
-                    "variant => return Err(bincode::error::DecodeError::UnexpectedVariant {{ min: 0, max: {}, found: variant }})",
-                    variants.len() - 1
+                    "variant => return Err(bincode::error::DecodeError::UnexpectedVariant {{ min: 0, max: {}, found: variant, type_name: {:?} }})",
+                    variants.len() - 1,
+                    enum_name.to_string()
                 ));
             });
         });
diff --git a/derive/src/generate/generator.rs b/derive/src/generate/generator.rs
index 94edf7d..cf56fa6 100644
--- a/derive/src/generate/generator.rs
+++ b/derive/src/generate/generator.rs
@@ -24,6 +24,10 @@ impl Generator {
         }
     }
 
+    pub fn target_name(&self) -> &Ident {
+        &self.name
+    }
+
     pub fn impl_for<'a>(&'a mut self, trait_name: &str) -> ImplFor<'a> {
         ImplFor::new(self, trait_name)
     }
diff --git a/src/de/impls.rs b/src/de/impls.rs
index 7f683f0..b5b0933 100644
--- a/src/de/impls.rs
+++ b/src/de/impls.rs
@@ -127,6 +127,54 @@ impl<'de, T> Decodable for core::marker::PhantomData<T> {
     }
 }
 
+impl<'de, T> Decodable for Option<T>
+where
+    T: Decodable,
+{
+    fn decode<D: Decode>(mut decoder: D) -> Result<Self, DecodeError> {
+        let is_some = u8::decode(&mut decoder)?;
+        match is_some {
+            0 => Ok(None),
+            1 => {
+                let val = T::decode(decoder)?;
+                Ok(Some(val))
+            }
+            x => Err(DecodeError::UnexpectedVariant {
+                found: x as u32,
+                max: 1,
+                min: 0,
+                type_name: core::any::type_name::<Option<T>>(),
+            }),
+        }
+    }
+}
+
+impl<'de, T, U> Decodable for Result<T, U>
+where
+    T: Decodable,
+    U: Decodable,
+{
+    fn decode<D: Decode>(mut decoder: D) -> Result<Self, DecodeError> {
+        let is_ok = u8::decode(&mut decoder)?;
+        match is_ok {
+            0 => {
+                let t = T::decode(decoder)?;
+                Ok(Ok(t))
+            }
+            1 => {
+                let u = U::decode(decoder)?;
+                Ok(Err(u))
+            }
+            x => Err(DecodeError::UnexpectedVariant {
+                found: x as u32,
+                max: 1,
+                min: 0,
+                type_name: core::any::type_name::<Result<T, U>>(),
+            }),
+        }
+    }
+}
+
 impl<'a, 'de, T> Decode for &'a mut T
 where
     T: Decode,
diff --git a/src/enc/impls.rs b/src/enc/impls.rs
index 5b57ea6..6f7c431 100644
--- a/src/enc/impls.rs
+++ b/src/enc/impls.rs
@@ -115,6 +115,39 @@ impl<const N: usize> Encodeable for [u8; N] {
     }
 }
 
+impl<T> Encodeable for Option<T>
+where
+    T: Encodeable,
+{
+    fn encode<E: Encode>(&self, mut encoder: E) -> Result<(), EncodeError> {
+        if let Some(val) = self {
+            1u8.encode(&mut encoder)?;
+            val.encode(encoder)
+        } else {
+            0u8.encode(encoder)
+        }
+    }
+}
+
+impl<T, U> Encodeable for Result<T, U>
+where
+    T: Encodeable,
+    U: Encodeable,
+{
+    fn encode<E: Encode>(&self, mut encoder: E) -> Result<(), EncodeError> {
+        match self {
+            Ok(val) => {
+                0u8.encode(&mut encoder)?;
+                val.encode(encoder)
+            }
+            Err(err) => {
+                1u8.encode(&mut encoder)?;
+                err.encode(encoder)
+            }
+        }
+    }
+}
+
 impl<'a, T> Encode for &'a mut T
 where
     T: Encode,
diff --git a/src/error.rs b/src/error.rs
index eb0c66c..e687cac 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -34,6 +34,9 @@ pub enum DecodeError {
 
     /// Invalid enum variant was found. The decoder tried to decode variant index `found`, but the variant index should be between `min` and `max`.
     UnexpectedVariant {
+        /// The type name that was being decoded.
+        type_name: &'static str,
+
         /// The min index of the enum. Usually this is `0`.
         min: u32,
 
diff --git a/tests/basic_types.rs b/tests/basic_types.rs
index d5cc68d..afe338d 100644
--- a/tests/basic_types.rs
+++ b/tests/basic_types.rs
@@ -151,6 +151,13 @@ fn test_numbers() {
         225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240,
         241, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255
     ]);
+
+    // Option and Result
+    the_same(Option::<u32>::None);
+    the_same(Option::<u32>::Some(1234));
+
+    the_same(Result::<u32, u8>::Ok(1555));
+    the_same(Result::<u32, u8>::Err(15));
 }
 
 #[test]