Add a test for (rust.zerocopy)=true in proto3 (#151)

dropbox · Oct 18, 2023 · aabaaf8 · aabaaf8
1 parent 4d8c43f
commit aabaaf8
Show file tree

Hide file tree

Showing 3 changed files with 167 additions and 0 deletions.
diff --git a/pb-test/gen/pb-jelly/proto_pbtest/src/pbtest3.rs.expected b/pb-test/gen/pb-jelly/proto_pbtest/src/pbtest3.rs.expected
@@ -8887,6 +8887,152 @@ impl ::pb_jelly::Reflection for TestProto3Optional {
   }
 }
 
+#[derive(Clone, Debug, PartialEq)]
+pub struct TestProto3Zerocopy {
+  pub data1: ::pb_jelly::Lazy<::bytes::Bytes>,
+  pub data2: ::pb_jelly::Lazy<::bytes::Bytes>,
+}
+impl ::std::default::Default for TestProto3Zerocopy {
+  fn default() -> Self {
+    TestProto3Zerocopy {
+      data1: ::std::default::Default::default(),
+      data2: ::std::default::Default::default(),
+    }
+  }
+}
+lazy_static! {
+  pub static ref TestProto3Zerocopy_default: TestProto3Zerocopy = TestProto3Zerocopy::default();
+}
+impl ::pb_jelly::Message for TestProto3Zerocopy {
+  fn descriptor(&self) -> ::std::option::Option<::pb_jelly::MessageDescriptor> {
+    Some(::pb_jelly::MessageDescriptor {
+      name: "TestProto3Zerocopy",
+      full_name: "pbtest.TestProto3Zerocopy",
+      fields: &[
+        ::pb_jelly::FieldDescriptor {
+          name: "data1",
+          full_name: "pbtest.TestProto3Zerocopy.data1",
+          index: 0,
+          number: 1,
+          typ: ::pb_jelly::wire_format::Type::LengthDelimited,
+          label: ::pb_jelly::Label::Optional,
+          oneof_index: None,
+        },
+        ::pb_jelly::FieldDescriptor {
+          name: "data2",
+          full_name: "pbtest.TestProto3Zerocopy.data2",
+          index: 1,
+          number: 2,
+          typ: ::pb_jelly::wire_format::Type::LengthDelimited,
+          label: ::pb_jelly::Label::Optional,
+          oneof_index: None,
+        },
+      ],
+      oneofs: &[
+      ],
+    })
+  }
+  fn compute_size(&self) -> usize {
+    let mut size = 0;
+    let mut data1_size = 0;
+    if self.data1 != <::pb_jelly::Lazy<::bytes::Bytes> as ::std::default::Default>::default() {
+      let val = &self.data1;
+      let l = ::pb_jelly::Message::compute_size(val);
+      data1_size += ::pb_jelly::wire_format::serialized_length(1);
+      data1_size += ::pb_jelly::varint::serialized_length(l as u64);
+      data1_size += l;
+    }
+    size += data1_size;
+    let mut data2_size = 0;
+    if self.data2 != <::pb_jelly::Lazy<::bytes::Bytes> as ::std::default::Default>::default() {
+      let val = &self.data2;
+      let l = ::pb_jelly::Message::compute_size(val);
+      data2_size += ::pb_jelly::wire_format::serialized_length(2);
+      data2_size += ::pb_jelly::varint::serialized_length(l as u64);
+      data2_size += l;
+    }
+    size += data2_size;
+    size
+  }
+  fn compute_grpc_slices_size(&self) -> usize {
+    let mut size = 0;
+    if self.data1 != <::pb_jelly::Lazy<::bytes::Bytes> as ::std::default::Default>::default() {
+      let val = &self.data1;
+      size += ::pb_jelly::Message::compute_grpc_slices_size(val);
+    }
+    if self.data2 != <::pb_jelly::Lazy<::bytes::Bytes> as ::std::default::Default>::default() {
+      let val = &self.data2;
+      size += ::pb_jelly::Message::compute_grpc_slices_size(val);
+    }
+    size
+  }
+  fn serialize<W: ::pb_jelly::PbBufferWriter>(&self, w: &mut W) -> ::std::io::Result<()> {
+    if self.data1 != <::pb_jelly::Lazy<::bytes::Bytes> as ::std::default::Default>::default() {
+      let val = &self.data1;
+      ::pb_jelly::wire_format::write(1, ::pb_jelly::wire_format::Type::LengthDelimited, w)?;
+      let l = ::pb_jelly::Message::compute_size(val);
+      ::pb_jelly::varint::write(l as u64, w)?;
+      ::pb_jelly::Message::serialize(val, w)?;
+    }
+    if self.data2 != <::pb_jelly::Lazy<::bytes::Bytes> as ::std::default::Default>::default() {
+      let val = &self.data2;
+      ::pb_jelly::wire_format::write(2, ::pb_jelly::wire_format::Type::LengthDelimited, w)?;
+      let l = ::pb_jelly::Message::compute_size(val);
+      ::pb_jelly::varint::write(l as u64, w)?;
+      ::pb_jelly::Message::serialize(val, w)?;
+    }
+    Ok(())
+  }
+  fn deserialize<B: ::pb_jelly::PbBufferReader>(&mut self, mut buf: &mut B) -> ::std::io::Result<()> {
+    while let Some((field_number, typ)) = ::pb_jelly::wire_format::read(&mut buf)? {
+      match field_number {
+        1 => {
+          ::pb_jelly::ensure_wire_format(typ, ::pb_jelly::wire_format::Type::LengthDelimited, "TestProto3Zerocopy", 1)?;
+          let len = ::pb_jelly::varint::ensure_read(&mut buf)?;
+          let mut next = ::pb_jelly::ensure_split(buf, len as usize)?;
+          let mut val: ::pb_jelly::Lazy<::bytes::Bytes> = ::std::default::Default::default();
+          ::pb_jelly::Message::deserialize(&mut val, &mut next)?;
+          self.data1 = val;
+        }
+        2 => {
+          ::pb_jelly::ensure_wire_format(typ, ::pb_jelly::wire_format::Type::LengthDelimited, "TestProto3Zerocopy", 2)?;
+          let len = ::pb_jelly::varint::ensure_read(&mut buf)?;
+          let mut next = ::pb_jelly::ensure_split(buf, len as usize)?;
+          let mut val: ::pb_jelly::Lazy<::bytes::Bytes> = ::std::default::Default::default();
+          ::pb_jelly::Message::deserialize(&mut val, &mut next)?;
+          self.data2 = val;
+        }
+        _ => {
+          ::pb_jelly::skip(typ, &mut buf)?;
+        }
+      }
+    }
+    Ok(())
+  }
+}
+impl ::pb_jelly::Reflection for TestProto3Zerocopy {
+  fn which_one_of(&self, oneof_name: &str) -> ::std::option::Option<&'static str> {
+    match oneof_name {
+      _ => {
+        panic!("unknown oneof name given");
+      }
+    }
+  }
+  fn get_field_mut(&mut self, field_name: &str) -> ::pb_jelly::reflection::FieldMut<'_> {
+    match field_name {
+      "data1" => {
+        ::pb_jelly::reflection::FieldMut::Value(&mut self.data1)
+      }
+      "data2" => {
+        ::pb_jelly::reflection::FieldMut::Value(&mut self.data2)
+      }
+      _ => {
+        panic!("unknown field name given")
+      }
+    }
+  }
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct RecursiveOneof {
   /// This field should be boxed automatically.

diff --git a/pb-test/proto/packages/pbtest/pbtest3.proto b/pb-test/proto/packages/pbtest/pbtest3.proto
@@ -352,6 +352,11 @@ message TestProto3Optional {
     }
 }
 
+message TestProto3Zerocopy {
+    bytes data1 = 1 [(rust.zero_copy) = true];
+    bytes data2 = 2 [(rust.zero_copy) = true];
+}
+
 message RecursiveOneof {
     oneof oneof_field {
         // This field should be boxed automatically.

diff --git a/pb-test/src/pbtest.rs b/pb-test/src/pbtest.rs
@@ -3,6 +3,7 @@ use std::fs::File;
 use std::io::Cursor;
 use std::io::Read;
 
+use bytes::Bytes;
 use pb_jelly::reflection::FieldMut;
 use pb_jelly::wire_format::Type;
 use pb_jelly::{
@@ -872,6 +873,21 @@ fn test_proto3_optional() {
     assert_eq!(proto.a_int32, Some(456));
 }
 
+#[test]
+fn test_proto3_zerocopy_read() {
+    let data = Bytes::from_static(b"\x0a\x08zerocopy\x12\x09zerocopy2");
+
+    let mut proto = TestProto3Zerocopy::default();
+    proto.deserialize(&mut Cursor::new(data.clone())).unwrap();
+    let data1 = proto.data1.into_buffer();
+    let data2 = proto.data2.into_buffer();
+    assert_eq!(data1, &b"zerocopy"[..]);
+    assert_eq!(data2, &b"zerocopy2"[..]);
+    // The deserialized buffers should point into `data`.
+    data.slice_ref(data1.as_ref());
+    data.slice_ref(data2.as_ref());
+}
+
 // Test that boxing works properly for oneof fields.
 #[test]
 fn test_recursive_oneof() {