Fix segfault in upb pathway in terra rust

We had previously commented out the upb portion of simple_nested_test.
This is because nonmutable getters have submessages being NULL by default.
This means that trying to fetch anything, like a simple scalar from that nested message would segfault.

This CL makes the externC return an Option<RawMessage> since we've discovered that upb can return NULL. This way, we can check for `None` and handle the NULL case appropriately.

We know that the NULL pathway can only come from terra upb, since
cpp automagically constructs submsgs if they don't exist.

We've augmented upb.rs to contain a scratch space that allocates a zeroed-out contiguous chunk of memory @64KB. Since a block of zeroed-out memory is a legit message from upb's point of view, we can provide $pbr$::ScratchSpace::zeroed_block() to upb in order to get the default submessage behavior we want from upb.

This block is lazily allocated upon first request. This means that a consumer of the cpp kernel will not incur an additional cost.

PiperOrigin-RevId: 573840755

rust/test/shared/BUILD

@@ -220,10 +220,8 @@ rust_test(
     deps = ["//rust/test:nested_cc_rust_proto"],
 )
 
-# TODO: enable once we emit default messages for upb at the rust level
-# The default submsg that upb emits is NULL (whereas cpp returns a default constructed msg)
-#rust_test(
-#    name = "simple_nested_upb_test",
-#    srcs = ["simple_nested_test.rs"],
-#    deps = ["//rust/test:nested_upb_rust_proto"],
-#)
+rust_test(
+    name = "simple_nested_upb_test",
+    srcs = ["simple_nested_test.rs"],
+    deps = ["//rust/test:nested_upb_rust_proto"],
+)

rust/test/shared/simple_nested_test.rs

@@ -10,7 +10,6 @@ use nested_proto::nest::Outer;
 #[test]
 fn test_simple_nested_proto() {
     let outer_msg = Outer::new();
-    // TODO: passing cpp, segfaulting upb
     assert_eq!(outer_msg.inner().num(), 0);
     assert!(!outer_msg.inner().boolean());
 }

rust/upb.rs

@@ -17,6 +17,7 @@ use std::mem::MaybeUninit;
 use std::ops::Deref;
 use std::ptr::{self, NonNull};
 use std::slice;
+use std::sync::Once;
 
 /// See `upb/port/def.inc`.
 const UPB_MALLOC_ALIGN: usize = 8;
@@ -134,6 +135,38 @@ impl Drop for Arena {
     }
 }
 
+static mut INTERNAL_PTR: Option<RawMessage> = None;
+static INIT: Once = Once::new();
+
+// TODO:(b/304577017)
+const ALIGN: usize = 32;
+const UPB_SCRATCH_SPACE_BYTES: usize = 64_000;
+
+/// Holds a zero-initialized block of memory for use by upb.
+/// By default, if a message is not set in cpp, a default message is created.
+/// upb departs from this and returns a null ptr. However, since contiguous
+/// chunks of memory filled with zeroes are legit messages from upb's point of
+/// view, we can allocate a large block and refer to that when dealing
+/// with readonly access.
+pub struct ScratchSpace;
+impl ScratchSpace {
+    pub fn zeroed_block() -> RawMessage {
+        unsafe {
+            INIT.call_once(|| {
+                let layout =
+                    std::alloc::Layout::from_size_align(UPB_SCRATCH_SPACE_BYTES, ALIGN).unwrap();
+                let Some(ptr) =
+                    crate::__internal::RawMessage::new(std::alloc::alloc_zeroed(layout).cast())
+                else {
+                    std::alloc::handle_alloc_error(layout)
+                };
+                INTERNAL_PTR = Some(ptr)
+            });
+            INTERNAL_PTR.unwrap()
+        }
+    }
+}
+
 /// Serialized Protobuf wire format data.
 ///
 /// It's typically produced by `<Message>::serialize()`.

src/google/protobuf/compiler/rust/accessors/singular_message.cc

@@ -7,7 +7,6 @@
 
 #include "absl/strings/match.h"
 #include "absl/strings/string_view.h"
-#include "absl/strings/strip.h"
 #include "google/protobuf/compiler/cpp/helpers.h"
 #include "google/protobuf/compiler/rust/accessors/accessor_generator.h"
 #include "google/protobuf/compiler/rust/context.h"
@@ -28,26 +27,56 @@ void SingularMessage::InMsgImpl(Context<FieldDescriptor> field) const {
     // TODO: Handle imports correctly, default to $Msg$View for now
     prefix = field.desc().containing_type()->name();
   }
-  field.Emit(
-      {
-          {"prefix", prefix},
-          {"field", field.desc().name()},
-          {"getter_thunk", Thunk(field, "get")},
-      },
-      R"rs(
+  if (field.is_cpp()) {
+    field.Emit({{"prefix", prefix},
+                {"field", field.desc().name()},
+                {"getter_thunk", Thunk(field, "get")}},
+               R"rs(
           pub fn r#$field$(&self) -> $prefix$View {
-            $prefix$View::new($pbi$::Private, unsafe { $getter_thunk$(self.inner.msg) } )
+            // For C++ kernel, getters automatically return the
+            // default_instance if the field is unset.
+            let submsg = unsafe { $getter_thunk$(self.inner.msg) };
+            $prefix$View::new($pbi$::Private, submsg)
           }
         )rs");
+  } else {
+    field.Emit({{"prefix", prefix},
+                {"field", field.desc().name()},
+                {"getter_thunk", Thunk(field, "get")}},
+               R"rs(
+          pub fn r#$field$(&self) -> $prefix$View {
+            let submsg = unsafe { $getter_thunk$(self.inner.msg) };
+            // For upb, getters return null if the field is unset, so we need to
+            // check for null and return the default instance manually. Note that
+            // a null ptr received from upb manifests as Option::None
+            match submsg {
+                // TODO:(b/304357029)
+                None => $prefix$View::new($pbi$::Private, $pbr$::ScratchSpace::zeroed_block()),
+                Some(field) => $prefix$View::new($pbi$::Private, field),
+              }
+          }
+        )rs");
+  }
 }
 
 void SingularMessage::InExternC(Context<FieldDescriptor> field) const {
   field.Emit(
       {
           {"getter_thunk", Thunk(field, "get")},
+          {"ReturnType",
+           [&] {
+             if (field.is_cpp()) {
+               // guaranteed to have a nonnull submsg for the cpp kernel
+               field.Emit({}, "$pbi$::RawMessage;");
+             } else {
+               // upb kernel may return NULL for a submsg, we can detect this
+               // in terra rust if the option returned is None
+               field.Emit({}, "Option<$pbi$::RawMessage>;");
+             }
+           }},
       },
       R"rs(
-                  fn $getter_thunk$(raw_msg: $pbi$::RawMessage) -> $pbi$::RawMessage;
+                  fn $getter_thunk$(raw_msg: $pbi$::RawMessage) -> $ReturnType$;
                )rs");
 }