From a1b6dd9247b5fd7fca9db035209a6b19310c437a Mon Sep 17 00:00:00 2001
From: Jakob Buchgraber <buchgr@google.com>
Date: Thu, 27 Apr 2023 00:30:07 -0700
Subject: [PATCH] Implement Msg.deserialize() for the UPB kernel

PiperOrigin-RevId: 527496306
---
 rust/test/shared/BUILD                        | 48 ++++++++-----------
 ...st_proto_test.rs => serialization_test.rs} | 30 ++++++++++--
 src/google/protobuf/compiler/rust/message.cc  | 25 ++++++++--
 3 files changed, 68 insertions(+), 35 deletions(-)
 rename rust/test/shared/{unittest_proto_test.rs => serialization_test.rs} (67%)

diff --git a/rust/test/shared/BUILD b/rust/test/shared/BUILD
index c163263a3c..ef2c7a93a9 100644
--- a/rust/test/shared/BUILD
+++ b/rust/test/shared/BUILD
@@ -16,34 +16,6 @@
 
 load("@rules_rust//rust:defs.bzl", "rust_test")
 
-rust_test(
-    name = "unittest_proto_upb_test",
-    srcs = ["unittest_proto_test.rs"],
-    tags = [
-        # TODO(b/270274576): Enable testing on arm once we have a Rust Arm toolchain.
-        "not_build:arm",
-        # TODO(b/243126140): Enable tsan once we support sanitizers with Rust.
-        "notsan",
-        # TODO(b/243126140): Enable msan once we support sanitizers with Rust.
-        "nomsan",
-    ],
-    deps = ["//rust/test:unittest_upb_rust_proto"],
-)
-
-rust_test(
-    name = "unittest_proto_cpp_test",
-    srcs = ["unittest_proto_test.rs"],
-    tags = [
-        # TODO(b/270274576): Enable testing on arm once we have a Rust Arm toolchain.
-        "not_build:arm",
-        # TODO(b/243126140): Enable tsan once we support sanitizers with Rust.
-        "notsan",
-        # TODO(b/243126140): Enable msan once we support sanitizers with Rust.
-        "nomsan",
-    ],
-    deps = ["//rust/test:unittest_cc_rust_proto"],
-)
-
 rust_test(
     name = "child_parent_upb_test",
     srcs = ["child_parent_test.rs"],
@@ -141,3 +113,23 @@ rust_test(
     ],
     deps = ["//rust/test:unittest_upb_rust_proto"],
 )
+
+rust_test(
+    name = "serialization_upb_test",
+    srcs = ["serialization_test.rs"],
+    tags = [
+        # TODO(b/270274576): Enable testing on arm once we have a Rust Arm toolchain.
+        "not_build:arm",
+    ],
+    deps = ["//rust/test:unittest_upb_rust_proto"],
+)
+
+rust_test(
+    name = "serialization_cpp_test",
+    srcs = ["serialization_test.rs"],
+    tags = [
+        # TODO(b/270274576): Enable testing on arm once we have a Rust Arm toolchain.
+        "not_build:arm",
+    ],
+    deps = ["//rust/test:unittest_cc_rust_proto"],
+)
diff --git a/rust/test/shared/unittest_proto_test.rs b/rust/test/shared/serialization_test.rs
similarity index 67%
rename from rust/test/shared/unittest_proto_test.rs
rename to rust/test/shared/serialization_test.rs
index 76d05f8fe7..67d285fd16 100644
--- a/rust/test/shared/unittest_proto_test.rs
+++ b/rust/test/shared/serialization_test.rs
@@ -31,7 +31,31 @@
 use unittest_proto::proto2_unittest::TestAllTypes;
 
 #[test]
-fn test_serialization() {
-    let test_all_types = TestAllTypes::new();
-    assert_eq!(*test_all_types.serialize(), []);
+fn serialize_deserialize_message() {
+    let mut msg = TestAllTypes::new();
+    msg.optional_int64_set(Some(42));
+    msg.optional_bool_set(Some(true));
+    msg.optional_bytes_set(Some(b"serialize deserialize test"));
+
+    let serialized = msg.serialize();
+
+    let mut msg2 = TestAllTypes::new();
+    assert!(msg2.deserialize(&serialized).is_ok());
+
+    assert_eq!(msg.optional_int64(), msg2.optional_int64());
+    assert_eq!(msg.optional_bool(), msg2.optional_bool());
+    assert_eq!(msg.optional_bytes(), msg2.optional_bytes());
+}
+
+#[test]
+fn deserialize_empty() {
+    let mut msg = TestAllTypes::new();
+    assert!(msg.deserialize(&[]).is_ok());
+}
+
+#[test]
+fn deserialize_error() {
+    let mut msg = TestAllTypes::new();
+    let data = b"not a serialized proto";
+    assert!(msg.deserialize(&*data).is_err());
 }
diff --git a/src/google/protobuf/compiler/rust/message.cc b/src/google/protobuf/compiler/rust/message.cc
index 8a4627a816..49b3a95bdb 100644
--- a/src/google/protobuf/compiler/rust/message.cc
+++ b/src/google/protobuf/compiler/rust/message.cc
@@ -77,7 +77,7 @@ void MessageNew(Context<Descriptor> msg) {
 
     case Kernel::kUpb:
       msg.Emit({{"new_thunk", Thunk(msg, "new")}}, R"rs(
-        let arena = unsafe { $pbi$::Arena::new() };
+        let arena = $pbi$::Arena::new();
         Self {
           msg: unsafe { $new_thunk$(arena.raw()) },
           arena,
@@ -133,9 +133,24 @@ void MessageDeserialize(Context<Descriptor> msg) {
       return;
 
     case Kernel::kUpb:
-      msg.Emit(R"rs(
-        let _ = data;
-        $std$::unimplemented!()
+      msg.Emit({{"deserialize_thunk", Thunk(msg, "parse")}}, R"rs(
+        let arena = $pbi$::Arena::new();
+        let msg = unsafe {
+          $NonNull$::<u8>::new(
+            $deserialize_thunk$(data.as_ptr(), data.len(), arena.raw())
+          )
+        };
+
+        match msg {
+          None => Err($pb$::ParseError),
+          Some(msg) => {
+            // This assignment causes self.arena to be dropped and to deallocate
+            // any previous message pointed/owned to by self.msg.
+            self.arena = arena;
+            self.msg = msg;
+            Ok(())
+          }
+        }
       )rs");
       return;
   }
@@ -166,10 +181,12 @@ void MessageExterns(Context<Descriptor> msg) {
           {
               {"new_thunk", Thunk(msg, "new")},
               {"serialize_thunk", Thunk(msg, "serialize")},
+              {"deserialize_thunk", Thunk(msg, "parse")},
           },
           R"rs(
           fn $new_thunk$(arena: $pbi$::RawArena) -> $NonNull$<u8>;
           fn $serialize_thunk$(msg: $NonNull$<u8>, arena: $pbi$::RawArena, len: &mut usize) -> $NonNull$<u8>;
+          fn $deserialize_thunk$(data: *const u8, size: usize, arena: $pbi$::RawArena) -> *mut u8;
       )rs");
       return;
   }