Make types for 'inner' (runtime-specific) owned Repeated and Maps

PiperOrigin-RevId: 614995436

rust/cpp.rs

@@ -274,6 +274,18 @@ pub fn copy_bytes_in_arena_if_needed_by_runtime<'msg>(
     val
 }
 
+/// The raw type-erased version of an owned `Repeated`.
+#[derive(Debug)]
+pub struct InnerRepeated {
+    raw: RawRepeatedField,
+}
+
+impl InnerRepeated {
+    pub fn as_mut(&mut self) -> InnerRepeatedMut<'_> {
+        InnerRepeatedMut::new(Private, self.raw)
+    }
+}
+
 /// The raw type-erased pointer version of `RepeatedMut`.
 ///
 /// Contains a `proto2::RepeatedField*` or `proto2::RepeatedPtrField*`.
@@ -363,9 +375,9 @@ macro_rules! impl_repeated_primitives {
             unsafe impl ProxiedInRepeated for $t {
                 #[allow(dead_code)]
                 fn repeated_new(_: Private) -> Repeated<$t> {
-                    unsafe {
-                        Repeated::from_inner(InnerRepeatedMut::new(Private, $new_thunk()))
-                    }
+                    Repeated::from_inner(InnerRepeated {
+                        raw: unsafe { $new_thunk() }
+                    })
                 }
                 #[allow(dead_code)]
                 unsafe fn repeated_free(_: Private, f: &mut Repeated<$t>) {
@@ -441,6 +453,21 @@ pub fn cast_enum_repeated_mut<E: Enum + ProxiedInRepeated>(
     }
 }
 
+#[derive(Debug)]
+pub struct InnerMap {
+    pub(crate) raw: RawMap,
+}
+
+impl InnerMap {
+    pub fn new(_private: Private, raw: RawMap) -> Self {
+        Self { raw }
+    }
+
+    pub fn as_mut(&mut self) -> InnerMapMut<'_> {
+        InnerMapMut { raw: self.raw, _phantom: PhantomData }
+    }
+}
+
 #[derive(Clone, Copy, Debug)]
 pub struct InnerMapMut<'msg> {
     pub(crate) raw: RawMap,
@@ -564,9 +591,8 @@ macro_rules! impl_ProxiedInMapValue_for_non_generated_value_types {
                     unsafe {
                         Map::from_inner(
                             Private,
-                            InnerMapMut {
+                            InnerMap {
                                 raw: [< __rust_proto_thunk__Map_ $key_t _ $t _new >](),
-                                _phantom: PhantomData
                             }
                         )
                     }

rust/map.rs

@@ -8,7 +8,7 @@
 use crate::{
     Mut, MutProxy, Proxied, SettableValue, View, ViewProxy,
     __internal::{Private, RawMap},
-    __runtime::{InnerMapMut, RawMapIter},
+    __runtime::{InnerMap, InnerMapMut, RawMapIter},
 };
 use std::marker::PhantomData;
 
@@ -60,7 +60,7 @@ impl<'msg, K: ?Sized, V: ?Sized> std::fmt::Debug for MapMut<'msg, K, V> {
 }
 
 pub struct Map<K: ?Sized + Proxied, V: ?Sized + ProxiedInMapValue<K>> {
-    inner: InnerMapMut<'static>,
+    inner: InnerMap,
     _phantom: PhantomData<(PhantomData<K>, PhantomData<V>)>,
 }
 
@@ -169,28 +169,20 @@ where
     }
 
     pub fn as_mut(&mut self) -> MapMut<'_, K, V> {
-        MapMut { inner: self.inner, _phantom: PhantomData }
+        MapMut { inner: self.inner.as_mut(), _phantom: PhantomData }
     }
 
     pub fn as_view(&self) -> MapView<'_, K, V> {
         MapView { raw: self.inner.raw, _phantom: PhantomData }
     }
 
-    /// # Safety
-    /// - `inner` must be valid to read and write from for `'static`.
-    /// - There must be no aliasing references or mutations on the same
-    ///   underlying object.
     #[doc(hidden)]
-    pub unsafe fn from_inner(_private: Private, inner: InnerMapMut<'static>) -> Self {
+    pub fn from_inner(_private: Private, inner: InnerMap) -> Self {
         Self { inner, _phantom: PhantomData }
     }
 
     pub fn as_raw(&self, _private: Private) -> RawMap {
-        self.inner.as_raw(Private)
-    }
-
-    pub fn inner(&self, _private: Private) -> InnerMapMut<'static> {
-        self.inner
+        self.inner.raw
     }
 }
 

rust/repeated.rs

@@ -17,7 +17,7 @@ use std::marker::PhantomData;
 use crate::{
     Mut, MutProxy, Proxied, SettableValue, View, ViewProxy,
     __internal::{Private, RawRepeatedField},
-    __runtime::InnerRepeatedMut,
+    __runtime::{InnerRepeated, InnerRepeatedMut},
 };
 
 /// Views the elements in a `repeated` field of `T`.
@@ -275,7 +275,7 @@ impl<'msg, T: ?Sized> Debug for RepeatedIter<'msg, T> {
 /// Users will generally write [`View<Repeated<T>>`](RepeatedView) or
 /// [`Mut<Repeated<T>>`](RepeatedMut) to access the repeated elements
 pub struct Repeated<T: ?Sized + ProxiedInRepeated> {
-    inner: InnerRepeatedMut<'static>,
+    inner: InnerRepeated,
     _phantom: PhantomData<T>,
 }
 
@@ -285,17 +285,12 @@ impl<T: ?Sized + ProxiedInRepeated> Repeated<T> {
         T::repeated_new(Private)
     }
 
-    pub(crate) unsafe fn from_inner(inner: InnerRepeatedMut<'static>) -> Self {
+    pub(crate) fn from_inner(inner: InnerRepeated) -> Self {
         Self { inner, _phantom: PhantomData }
     }
 
-    #[allow(dead_code)]
-    pub(crate) fn inner(&mut self) -> InnerRepeatedMut<'static> {
-        self.inner
-    }
-
     pub(crate) fn as_mut(&mut self) -> RepeatedMut<'_, T> {
-        RepeatedMut { inner: self.inner, _phantom: PhantomData }
+        RepeatedMut { inner: self.inner.as_mut(), _phantom: PhantomData }
     }
 }
 

rust/upb.rs

@@ -359,28 +359,30 @@ extern "C" {
     ) -> Option<RawMessage>;
 }
 
+/// The raw type-erased version of an owned `Repeated`.
+#[derive(Debug)]
+pub struct InnerRepeated {
+    raw: RawRepeatedField,
+    arena: Arena,
+}
+
+impl InnerRepeated {
+    pub fn as_mut(&mut self) -> InnerRepeatedMut<'_> {
+        InnerRepeatedMut::new(Private, self.raw, &self.arena)
+    }
+}
+
 /// The raw type-erased pointer version of `RepeatedMut`.
-///
-/// Contains a `upb_Array*` as well as `RawArena`, most likely that of the
-/// containing message. upb requires a `RawArena` to perform mutations on
-/// a repeated field.
-///
-/// An owned `Repeated` stores a `InnerRepeatedMut<'static>` and manages the
-/// contained `RawArena`.
 #[derive(Clone, Copy, Debug)]
 pub struct InnerRepeatedMut<'msg> {
     pub(crate) raw: RawRepeatedField,
-    // Storing a `RawArena` instead of `&Arena` allows this to be used for
-    // both `RepeatedMut<T>` and `Repeated<T>`.
-    arena: RawArena,
-    _phantom: PhantomData<&'msg Arena>,
+    arena: &'msg Arena,
 }
 
 impl<'msg> InnerRepeatedMut<'msg> {
     #[doc(hidden)]
-    #[allow(clippy::needless_pass_by_ref_mut)] // Sound construction requires mutable access.
     pub fn new(_private: Private, raw: RawRepeatedField, arena: &'msg Arena) -> Self {
-        InnerRepeatedMut { raw, arena: arena.raw(), _phantom: PhantomData }
+        InnerRepeatedMut { raw, arena }
     }
 }
 
@@ -445,25 +447,15 @@ macro_rules! impl_repeated_base {
     ($t:ty, $elem_t:ty, $ufield:ident, $upb_tag:expr) => {
         #[allow(dead_code)]
         fn repeated_new(_: Private) -> Repeated<$t> {
-            let arena = ManuallyDrop::new(Arena::new());
-            let raw_arena = arena.raw();
-            unsafe {
-                Repeated::from_inner(InnerRepeatedMut {
-                    raw: upb_Array_New(raw_arena, $upb_tag as c_int),
-                    arena: raw_arena,
-                    _phantom: PhantomData,
-                })
-            }
+            let arena = Arena::new();
+            Repeated::from_inner(InnerRepeated {
+                raw: unsafe { upb_Array_New(arena.raw(), $upb_tag as c_int) },
+                arena,
+            })
         }
         #[allow(dead_code)]
-        unsafe fn repeated_free(_: Private, f: &mut Repeated<$t>) {
-            // Freeing the array itself is handled by `Arena::Drop`
-            // SAFETY:
-            // - `f.raw_arena()` is a live `upb_Arena*` as
-            // - This function is only called once for `f`
-            unsafe {
-                upb_Arena_Free(f.inner().arena);
-            }
+        unsafe fn repeated_free(_: Private, _f: &mut Repeated<$t>) {
+            // No-op: the memory will be dropped by the arena.
         }
         fn repeated_len(f: View<Repeated<$t>>) -> usize {
             unsafe { upb_Array_Size(f.as_raw(Private)) }
@@ -571,7 +563,7 @@ impl<'msg, T: ?Sized> RepeatedMut<'msg, T> {
     // Returns a `RawArena` which is live for at least `'msg`
     #[doc(hidden)]
     pub fn raw_arena(&mut self, _private: Private) -> RawArena {
-        self.inner.arena
+        self.inner.arena.raw()
     }
 }
 
@@ -641,7 +633,7 @@ pub fn cast_enum_repeated_mut<E: Enum + ProxiedInRepeated>(
     // - No shared mutation is possible through the output.
     unsafe {
         let InnerRepeatedMut { arena, raw, .. } = repeated.inner;
-        RepeatedMut::from_inner(private, InnerRepeatedMut { arena, raw, _phantom: PhantomData })
+        RepeatedMut::from_inner(private, InnerRepeatedMut { arena, raw })
     }
 }
 
@@ -709,21 +701,36 @@ impl<'msg, K: ?Sized, V: ?Sized> MapMut<'msg, K, V> {
     // Returns a `RawArena` which is live for at least `'msg`
     #[doc(hidden)]
     pub fn raw_arena(&mut self, _private: Private) -> RawArena {
-        self.inner.raw_arena
+        self.inner.arena.raw()
+    }
+}
+
+#[derive(Debug)]
+pub struct InnerMap {
+    pub(crate) raw: RawMap,
+    arena: Arena,
+}
+
+impl InnerMap {
+    pub fn new(_private: Private, raw: RawMap, arena: Arena) -> Self {
+        Self { raw, arena }
+    }
+
+    pub fn as_mut(&mut self) -> InnerMapMut<'_> {
+        InnerMapMut { raw: self.raw, arena: &self.arena }
     }
 }
 
 #[derive(Clone, Copy, Debug)]
 pub struct InnerMapMut<'msg> {
     pub(crate) raw: RawMap,
-    raw_arena: RawArena,
-    _phantom: PhantomData<&'msg Arena>,
+    arena: &'msg Arena,
 }
 
 #[doc(hidden)]
 impl<'msg> InnerMapMut<'msg> {
-    pub fn new(_private: Private, raw: RawMap, raw_arena: RawArena) -> Self {
-        InnerMapMut { raw, raw_arena, _phantom: PhantomData }
+    pub fn new(_private: Private, raw: RawMap, arena: &'msg Arena) -> Self {
+        InnerMapMut { raw, arena }
     }
 
     #[doc(hidden)]
@@ -733,7 +740,7 @@ impl<'msg> InnerMapMut<'msg> {
 
     #[doc(hidden)]
     pub fn raw_arena(&self, _private: Private) -> RawArena {
-        self.raw_arena
+        self.arena.raw()
     }
 }
 
@@ -878,30 +885,16 @@ macro_rules! impl_ProxiedInMapValue_for_non_generated_value_types {
             impl ProxiedInMapValue<$key_t> for $t {
                 fn map_new(_private: Private) -> Map<$key_t, Self> {
                     let arena = Arena::new();
-                    let raw_arena = arena.raw();
-                    std::mem::forget(arena);
-
-                    unsafe {
-                        Map::from_inner(
-                            Private,
-                            InnerMapMut {
-                                raw: upb_Map_New(raw_arena,
-                                    <$key_t as UpbTypeConversions>::upb_type(),
-                                    <$t as UpbTypeConversions>::upb_type()),
-                                raw_arena,
-                                _phantom: PhantomData
-                            }
-                        )
-                    }
+                    let raw = unsafe {
+                        upb_Map_New(arena.raw(),
+                            <$key_t as UpbTypeConversions>::upb_type(),
+                            <$t as UpbTypeConversions>::upb_type())
+                    };
+                    Map::from_inner(Private, InnerMap { raw, arena })
                 }
 
-                unsafe fn map_free(_private: Private, map: &mut Map<$key_t, Self>) {
-                    // SAFETY:
-                    // - `map.inner.raw_arena` is a live `upb_Arena*`
-                    // - This function is only called once for `map` in `Drop`.
-                    unsafe {
-                        upb_Arena_Free(map.as_mut().raw_arena(Private));
-                    }
+                unsafe fn map_free(_private: Private, _map: &mut Map<$key_t, Self>) {
+                    // No-op: the memory will be dropped by the arena.
                 }
 
                 fn map_clear(mut map: Mut<'_, Map<$key_t, Self>>) {

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

@@ -86,7 +86,7 @@ void Map::InMsgImpl(Context& ctx, const FieldDescriptor& field,
                                            self.arena().raw())
                       };
                       let inner = $pbr$::InnerMapMut::new($pbi$::Private,
-                        raw, self.arena().raw());
+                        raw, self.arena());
                       unsafe { $pb$::MapMut::from_inner($pbi$::Private, inner) }
                     })rs");
                } else {

src/google/protobuf/compiler/rust/enum.cc

@@ -88,7 +88,7 @@ void EnumProxiedInMapValue(Context& ctx, const EnumDescriptor& desc) {
             unsafe {
                 $pb$::Map::from_inner(
                     $pbi$::Private,
-                    $pbr$::InnerMapMut::new($pbi$::Private, $map_new_thunk$())
+                    $pbr$::InnerMap::new($pbi$::Private, $map_new_thunk$())
                 )
             }
         }
@@ -166,29 +166,19 @@ void EnumProxiedInMapValue(Context& ctx, const EnumDescriptor& desc) {
       impl $pb$::ProxiedInMapValue<$key_t$> for $name$ {
           fn map_new(_private: $pbi$::Private) -> $pb$::Map<$key_t$, Self> {
               let arena = $pbr$::Arena::new();
-              let raw_arena = arena.raw();
-              std::mem::forget(arena);
-
-              unsafe {
-                  $pb$::Map::from_inner(
-                      $pbi$::Private,
-                      $pbr$::InnerMapMut::new(
-                          $pbi$::Private,
-                          $pbr$::upb_Map_New(
-                              raw_arena,
-                              <$key_t$ as $pbr$::UpbTypeConversions>::upb_type(),
-                              $pbr$::UpbCType::Enum),
-                          raw_arena))
-              }
+              let raw = unsafe {
+                  $pbr$::upb_Map_New(
+                      arena.raw(),
+                      <$key_t$ as $pbr$::UpbTypeConversions>::upb_type(),
+                      $pbr$::UpbCType::Enum)
+              };
+              $pb$::Map::from_inner(
+                  $pbi$::Private,
+                  $pbr$::InnerMap::new($pbi$::Private, raw, arena))
           }
 
           unsafe fn map_free(_private: $pbi$::Private, map: &mut $pb$::Map<$key_t$, Self>) {
-              // SAFETY:
-              // - `map.raw_arena($pbi$::Private)` is a live `upb_Arena*`
-              // - This function is only called once for `map` in `Drop`.
-              unsafe {
-                  $pbr$::upb_Arena_Free(map.inner($pbi$::Private).raw_arena($pbi$::Private));
-              }
+              // No-op: the memory will be dropped by the arena.
           }
 
           fn map_clear(mut map: $pb$::Mut<'_, $pb$::Map<$key_t$, Self>>) {

src/google/protobuf/compiler/rust/message.cc

@@ -470,7 +470,7 @@ void MessageProxiedInMapValue(Context& ctx, const Descriptor& msg) {
                     unsafe {
                         $pb$::Map::from_inner(
                             $pbi$::Private,
-                            $pbr$::InnerMapMut::new($pbi$::Private, $map_new_thunk$())
+                            $pbr$::InnerMap::new($pbi$::Private, $map_new_thunk$())
                         )
                     }
                 }
@@ -585,29 +585,20 @@ void MessageProxiedInMapValue(Context& ctx, const Descriptor& msg) {
             impl $pb$::ProxiedInMapValue<$key_t$> for $Msg$ {
                 fn map_new(_private: $pbi$::Private) -> $pb$::Map<$key_t$, Self> {
                     let arena = $pbr$::Arena::new();
-                    let raw_arena = arena.raw();
-                    std::mem::forget(arena);
+                    let raw = unsafe {
+                      $pbr$::upb_Map_New(
+                        arena.raw(),
+                        <$key_t$ as $pbr$::UpbTypeConversions>::upb_type(),
+                        <Self as $pbr$::UpbTypeConversions>::upb_type())
+                    };
 
-                    unsafe {
-                        $pb$::Map::from_inner(
-                            $pbi$::Private,
-                            $pbr$::InnerMapMut::new(
-                                $pbi$::Private,
-                                $pbr$::upb_Map_New(
-                                    raw_arena,
-                                    <$key_t$ as $pbr$::UpbTypeConversions>::upb_type(),
-                                    <Self as $pbr$::UpbTypeConversions>::upb_type()),
-                                raw_arena))
-                    }
+                    $pb$::Map::from_inner(
+                        $pbi$::Private,
+                        $pbr$::InnerMap::new($pbi$::Private, raw, arena))
                 }
 
                 unsafe fn map_free(_private: $pbi$::Private, map: &mut $pb$::Map<$key_t$, Self>) {
-                    // SAFETY:
-                    // - `map.raw_arena($pbi$::Private)` is a live `upb_Arena*`
-                    // - This function is only called once for `map` in `Drop`.
-                    unsafe {
-                        $pbr$::upb_Arena_Free(map.inner($pbi$::Private).raw_arena($pbi$::Private));
-                    }
+                    // No-op: the memory will be dropped by the arena.
                 }
 
                 fn map_clear(mut map: $pb$::Mut<'_, $pb$::Map<$key_t$, Self>>) {