Add support for Option<T>

Author: robertwayne

rust/bufferfish-core/src/decodable.rs

@@ -243,3 +243,22 @@ impl<T: Decodable> Decodable for Vec<T> {
         T::max_bytes_allowed().map(|max_t_size| 2 + (u16::MAX as usize * max_t_size))
     }
 }
+
+impl<T: Decodable> Decodable for Option<T> {
+    fn decode_value(bf: &mut Bufferfish) -> Result<Option<T>, BufferfishError> {
+        let flag = bf.read_u8()?;
+        match flag {
+            0 => Ok(None),
+            1 => Ok(Some(T::decode_value(bf)?)),
+            _ => Err(BufferfishError::InvalidEnumVariant),
+        }
+    }
+
+    fn min_bytes_required() -> Option<usize> {
+        Some(1)
+    }
+
+    fn max_bytes_allowed() -> Option<usize> {
+        T::max_bytes_allowed().map(|max_t_size| 1 + max_t_size)
+    }
+}

rust/bufferfish-core/src/encodable.rs

@@ -105,3 +105,15 @@ impl<T: Encodable> Encodable for Vec<T> {
         bf.write_array(self)
     }
 }
+
+impl<T: Encodable> Encodable for Option<T> {
+    fn encode_value(&self, bf: &mut Bufferfish) -> Result<(), BufferfishError> {
+        match self {
+            Some(value) => {
+                bf.write_u8(1)?;
+                value.encode_value(bf)
+            }
+            None => bf.write_u8(0),
+        }
+    }
+}

rust/bufferfish-core/src/lib.rs

@@ -372,6 +372,27 @@ impl Bufferfish {
         Ok(())
     }
 
+    /// Writes an `Option<T>` to the buffer, where `T` implements the
+    /// Encodable trait. If the option is `Some`, it will write a u8 with value
+    /// 1, followed by the encoded value. If the option is `None`, it will write
+    /// a u8 with value 0.
+    pub fn write_option<T: Encodable>(
+        &mut self,
+        option: &Option<T>,
+    ) -> Result<(), BufferfishError> {
+        match option {
+            Some(value) => {
+                self.write_u8(1)?;
+                value.encode_value(self)?;
+            }
+            None => {
+                self.write_u8(0)?;
+            }
+        }
+
+        Ok(())
+    }
+
     /// Writes an array of raw bytes to the buffer. Useful for encoding
     /// distinct structs into byte arrays and appending them to a buffer later.
     pub fn write_raw_bytes(&mut self, bytes: &[u8]) -> Result<(), BufferfishError> {
@@ -548,6 +569,19 @@ impl Bufferfish {
     }
 }
 
+/// Reads an `Option<T>` from the buffer, where `T` implements the
+/// Decodable trait. If the option is `Some`, it will read a u8 with value
+/// 1, followed by the decoded value. If the option is `None`, it will read
+/// a u8 with value 0.
+pub fn read_option<T: Decodable>(bf: &mut Bufferfish) -> Result<Option<T>, BufferfishError> {
+    let flag = bf.read_u8()?;
+    match flag {
+        0 => Ok(None),
+        1 => Ok(Some(T::decode(bf)?)),
+        _ => Err(BufferfishError::InvalidEnumVariant),
+    }
+}
+
 impl std::fmt::Display for Bufferfish {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         let inner = self.inner.get_ref();

rust/bufferfish/src/compiler.rs

@@ -381,6 +381,24 @@ fn get_bufferfish_write_fn(ty: Type, value_accessor: &str) -> String {
                 }
             }
 
+            if path.segments.len() == 1
+                && path.segments[0].ident == "Option"
+                && let PathArguments::AngleBracketed(ref args) = path.segments[0].arguments
+                && let Some(GenericArgument::Type(inner_ty)) = args.args.first()
+            {
+                let inner_write_call = if is_primitive_type(inner_ty) {
+                    let inner_write_fn = get_element_write_fn(inner_ty.clone());
+                    format!("bf.{inner_write_fn}({value_accessor})")
+                } else {
+                    let inner_write_fn = get_element_write_fn(inner_ty.clone());
+                    format!("{inner_write_fn}(bf, {value_accessor})")
+                };
+
+                return format!(
+                    "if ({value_accessor} !== null) {{\n        bf.writeUint8(1)\n        {inner_write_call}\n    }} else {{\n        bf.writeUint8(0)\n    }}",
+                );
+            }
+
             match path.get_ident().map(|ident| ident.to_string()).as_deref() {
                 Some("u8") => format!("bf.writeUint8({value_accessor})"),
                 Some("u16") => format!("bf.writeUint16({value_accessor})"),
@@ -491,6 +509,15 @@ fn get_typescript_type(ty: Type) -> String {
                 return format!("Array<{inner_ts_type}>");
             }
 
+            if path.segments.len() == 1
+                && path.segments[0].ident == "Option"
+                && let PathArguments::AngleBracketed(ref args) = path.segments[0].arguments
+                && let Some(GenericArgument::Type(inner_ty)) = args.args.first()
+            {
+                let inner_ts_type = get_typescript_type(inner_ty.clone());
+                return format!("{inner_ts_type} | null");
+            }
+
             match path.get_ident().map(|ident| ident.to_string()).as_deref() {
                 #[rustfmt::skip]
                 Some("u8") | Some("u16") | Some("u32") | Some("i8") | Some("i16") | Some("i32") 
@@ -626,6 +653,18 @@ fn get_bufferfish_fn(ty: Type) -> String {
                 return format!("bf.readArray(() => {inner_fn}) as Array<{inner_ts_type}>");
             }
 
+            if path.segments.len() == 1
+                && path.segments[0].ident == "Option"
+                && let PathArguments::AngleBracketed(ref args) = path.segments[0].arguments
+                && let Some(GenericArgument::Type(inner_ty)) = args.args.first()
+            {
+                let inner_fn = get_bufferfish_fn(inner_ty.clone());
+                let inner_ts_type = get_typescript_type(inner_ty.clone());
+                return format!(
+                    "(bf.readUint8() === 1 ? {inner_fn} : null) as {inner_ts_type} | null"
+                );
+            }
+
             match path.get_ident().map(|ident| ident.to_string()).as_deref() {
                 Some("u8") => "bf.readUint8() as number".to_string(),
                 Some("u16") => "bf.readUint16() as number".to_string(),
@@ -699,6 +738,7 @@ mod tests {
 pub enum MessageId {
     Join = 0,
     Leave,
+    Ping,
     Unknown = 255,
 }
 
@@ -716,6 +756,12 @@ pub struct LeaveMessage;
 #[derive(Encode)]
 #[bufferfish(MessageId::Unknown)]
 pub struct UnknownMessage(pub u8, pub u16);
+
+#[derive(Encode, Decode)]
+#[bufferfish(MessageId::Ping)]
+pub struct PingMessage {
+    pub timestamp: Option<u64>,
+}
     "#;
 
         let expected_output = r#"/* AUTOGENERATED BUFFERFISH FILE, DO NOT EDIT */
@@ -724,6 +770,7 @@ import { Bufferfish } from 'bufferfish'
 export enum MessageId {
     Join = 0,
     Leave = 1,
+    Ping = 2,
     Unknown = 255,
 }
 
@@ -770,6 +817,26 @@ export function encodeUnknownMessage(bf: Bufferfish, value: UnknownMessage): voi
     encodeMessageId(bf, MessageId.Unknown)
     bf.writeUint8(value[0])
     bf.writeUint16(value[1])
+}
+
+export interface PingMessage {
+    timestamp: bigint | null
+}
+
+export function decodePingMessage(bf: Bufferfish): PingMessage {
+    return {
+        timestamp: (bf.readUint8() === 1 ? bf.readUint64() as bigint : null) as bigint | null,
+    }
+}
+
+export function encodePingMessage(bf: Bufferfish, value: PingMessage): void {
+    encodeMessageId(bf, MessageId.Ping)
+    if (value.timestamp !== null) {
+        bf.writeUint8(1)
+        bf.writeUint64(value.timestamp)
+    } else {
+        bf.writeUint8(0)
+    }
 }"#;
 
         let items = syn::parse_file(test_file)

rust/bufferfish/src/lib.rs

@@ -1121,4 +1121,99 @@ mod tests {
         let decoded = CustomType::decode(&mut bf).unwrap();
         assert_eq!(decoded.value, 42);
     }
+
+    #[test]
+    fn test_optional_encode() {
+        use bufferfish_core as bufferfish;
+        use bufferfish_core::Encodable;
+
+        #[derive(Debug, Encode)]
+        struct OptionContainer {
+            value: Option<u32>,
+        }
+
+        let mut bf = Bufferfish::new();
+        let container_with_value = OptionContainer { value: Some(42) };
+        let container_without_value = OptionContainer { value: None };
+
+        container_with_value.encode(&mut bf).unwrap();
+        container_without_value.encode(&mut bf).unwrap();
+
+        assert_eq!(bf.as_ref(), &[1, 0, 0, 0, 42, 0]);
+    }
+
+    #[test]
+    fn test_optional_decode_some() {
+        use bufferfish_core as bufferfish;
+        use bufferfish_core::Decodable;
+
+        #[derive(Debug, Decode)]
+        struct OptionContainer {
+            value: Option<u32>,
+        }
+
+        let mut bf = Bufferfish::new();
+        bf.write_option(&Some(42)).unwrap();
+
+        let container_with_value = OptionContainer::decode(&mut bf).unwrap();
+
+        assert_eq!(container_with_value.value, Some(42));
+    }
+
+    #[test]
+    fn test_optional_decode_none() {
+        use bufferfish_core as bufferfish;
+        use bufferfish_core::Decodable;
+
+        #[derive(Debug, Decode)]
+        struct OptionContainer {
+            value: Option<u32>,
+        }
+
+        let mut bf = Bufferfish::new();
+        bf.write_option(&None::<u32>).unwrap();
+
+        let container_without_value = OptionContainer::decode(&mut bf).unwrap();
+
+        assert_eq!(container_without_value.value, None);
+    }
+
+    #[test]
+    fn test_encode_decode_with_option() {
+        use bufferfish_core as bufferfish;
+        use bufferfish_core::{Decodable, Encodable};
+
+        #[derive(Debug, Encode, Decode)]
+        struct User {
+            id: u32,
+            name: String,
+            age: Option<u8>,
+        }
+
+        let mut bf = Bufferfish::new();
+        let user_with_age = User {
+            id: 1,
+            name: "Bufferfish".to_string(),
+            age: Some(10),
+        };
+        let user_without_age = User {
+            id: 2,
+            name: "Bufferfish2".to_string(),
+            age: None,
+        };
+
+        user_with_age.encode(&mut bf).unwrap();
+        user_without_age.encode(&mut bf).unwrap();
+
+        let decoded_user_with_age = User::decode(&mut bf).unwrap();
+        let decoded_user_without_age = User::decode(&mut bf).unwrap();
+
+        assert_eq!(decoded_user_with_age.id, 1);
+        assert_eq!(decoded_user_with_age.name, "Bufferfish");
+        assert_eq!(decoded_user_with_age.age, Some(10));
+
+        assert_eq!(decoded_user_without_age.id, 2);
+        assert_eq!(decoded_user_without_age.name, "Bufferfish2");
+        assert_eq!(decoded_user_without_age.age, None);
+    }
 }