From c3df35950b147ea10d4a97de449ac36869b11cfc Mon Sep 17 00:00:00 2001 From: Adam Cozzette Date: Thu, 5 Sep 2024 11:09:02 -0700 Subject: [PATCH] Rust: create blanket implementation of `ProxiedInMapValue` for upb-backed messages This turned out to be much easier for upb than for C++. This CL pretty much just does a cut-and-paste of the `ProxiedInMapValue` implementation from the upb code generator to the runtime. This should help a bit with code size since it removes the need to generate six `ProxiedInMapValue` implementations per message. PiperOrigin-RevId: 671438289 --- rust/upb.rs | 106 ++++++++++++++++++- src/google/protobuf/compiler/rust/message.cc | 91 ---------------- 2 files changed, 101 insertions(+), 96 deletions(-) diff --git a/rust/upb.rs b/rust/upb.rs index 37d9158cf9..7ace10e323 100644 --- a/rust/upb.rs +++ b/rust/upb.rs @@ -9,15 +9,15 @@ use crate::__internal::{Enum, Private, SealedInternal}; use crate::{ - IntoProxied, Map, MapIter, MapMut, MapView, Mut, ProtoBytes, ProtoStr, ProtoString, Proxied, - ProxiedInMapValue, ProxiedInRepeated, Repeated, RepeatedMut, RepeatedView, View, + IntoProxied, Map, MapIter, MapMut, MapView, Message, Mut, ProtoBytes, ProtoStr, ProtoString, + Proxied, ProxiedInMapValue, ProxiedInRepeated, Repeated, RepeatedMut, RepeatedView, View, }; -use std::slice; use core::fmt::Debug; use std::alloc::Layout; -use std::sync::OnceLock; -use std::ptr::{self, NonNull}; use std::mem::{size_of, ManuallyDrop, MaybeUninit}; +use std::ptr::{self, NonNull}; +use std::slice; +use std::sync::OnceLock; #[cfg(bzl)] extern crate upb; @@ -797,6 +797,102 @@ macro_rules! impl_ProxiedInMapValue_for_key_types { impl_ProxiedInMapValue_for_key_types!(i32, u32, i64, u64, bool, ProtoString); +impl ProxiedInMapValue for MessageType +where + Key: Proxied + UpbTypeConversions, + MessageType: Proxied + UpbTypeConversions + Message, +{ + fn map_new(_private: Private) -> Map { + let arena = Arena::new(); + let raw = unsafe { + upb_Map_New( + arena.raw(), + ::upb_type(), + ::upb_type(), + ) + }; + + Map::from_inner(Private, InnerMap::new(raw, arena)) + } + + unsafe fn map_free(_private: Private, _map: &mut Map) { + // No-op: the memory will be dropped by the arena. + } + + fn map_clear(mut map: MapMut) { + unsafe { + upb_Map_Clear(map.as_raw(Private)); + } + } + + fn map_len(map: MapView) -> usize { + unsafe { upb_Map_Size(map.as_raw(Private)) } + } + + fn map_insert( + mut map: MapMut, + key: View<'_, Key>, + value: impl IntoProxied, + ) -> bool { + let arena = map.inner(Private).raw_arena(); + unsafe { + upb_Map_InsertAndReturnIfInserted( + map.as_raw(Private), + ::to_message_value(key), + ::into_message_value_fuse_if_required( + arena, + value.into_proxied(Private), + ), + arena, + ) + } + } + + fn map_get<'a>(map: MapView<'a, Key, Self>, key: View<'_, Key>) -> Option> { + let mut val = MaybeUninit::uninit(); + let found = unsafe { + upb_Map_Get( + map.as_raw(Private), + ::to_message_value(key), + val.as_mut_ptr(), + ) + }; + if !found { + return None; + } + Some(unsafe { ::from_message_value(val.assume_init()) }) + } + + fn map_remove(mut map: MapMut, key: View<'_, Key>) -> bool { + unsafe { + upb_Map_Delete( + map.as_raw(Private), + ::to_message_value(key), + ptr::null_mut(), + ) + } + } + fn map_iter(map: MapView) -> MapIter { + // SAFETY: MapView<'_,..>> guarantees its RawMap outlives '_. + unsafe { MapIter::from_raw(Private, RawMapIter::new(map.as_raw(Private))) } + } + + fn map_iter_next<'a>( + iter: &mut MapIter<'a, Key, Self>, + ) -> Option<(View<'a, Key>, View<'a, Self>)> { + // SAFETY: MapIter<'a, ..> guarantees its RawMapIter outlives 'a. + unsafe { iter.as_raw_mut(Private).next_unchecked() } + // SAFETY: MapIter returns key and values message values + // with the variants for K and V active. + .map(|(k, v)| unsafe { + ( + ::from_message_value(k), + ::from_message_value(v), + ) + }) + } +} + /// `upb_Map_Insert`, but returns a `bool` for whether insert occurred. /// /// Returns `true` if the entry was newly inserted. diff --git a/src/google/protobuf/compiler/rust/message.cc b/src/google/protobuf/compiler/rust/message.cc index eccd575a86..cacaa18302 100644 --- a/src/google/protobuf/compiler/rust/message.cc +++ b/src/google/protobuf/compiler/rust/message.cc @@ -718,97 +718,6 @@ void MessageProxiedInMapValue(Context& ctx, const Descriptor& msg) { } } )rs"); - for (const auto& t : kMapKeyTypes) { - ctx.Emit({io::Printer::Sub("key_t", [&] { ctx.Emit(t.rs_key_t); }) - .WithSuffix("")}, - R"rs( - 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 = unsafe { - $pbr$::upb_Map_New( - arena.raw(), - <$key_t$ as $pbr$::UpbTypeConversions>::upb_type(), - ::upb_type()) - }; - - $pb$::Map::from_inner( - $pbi$::Private, - $pbr$::InnerMap::new(raw, arena)) - } - - unsafe fn map_free(_private: $pbi$::Private, _map: &mut $pb$::Map<$key_t$, Self>) { - // No-op: the memory will be dropped by the arena. - } - - fn map_clear(mut map: $pb$::MapMut<$key_t$, Self>) { - unsafe { - $pbr$::upb_Map_Clear(map.as_raw($pbi$::Private)); - } - } - - fn map_len(map: $pb$::MapView<$key_t$, Self>) -> usize { - unsafe { - $pbr$::upb_Map_Size(map.as_raw($pbi$::Private)) - } - } - - fn map_insert(mut map: $pb$::MapMut<$key_t$, Self>, key: $pb$::View<'_, $key_t$>, value: impl $pb$::IntoProxied) -> bool { - let arena = map.inner($pbi$::Private).raw_arena(); - unsafe { - $pbr$::upb_Map_InsertAndReturnIfInserted( - map.as_raw($pbi$::Private), - <$key_t$ as $pbr$::UpbTypeConversions>::to_message_value(key), - ::into_message_value_fuse_if_required(arena, value.into_proxied($pbi$::Private)), - arena - ) - } - } - - fn map_get<'a>(map: $pb$::MapView<'a, $key_t$, Self>, key: $pb$::View<'_, $key_t$>) -> Option<$pb$::View<'a, Self>> { - let mut val = $std$::mem::MaybeUninit::uninit(); - let found = unsafe { - $pbr$::upb_Map_Get( - map.as_raw($pbi$::Private), - <$key_t$ as $pbr$::UpbTypeConversions>::to_message_value(key), - val.as_mut_ptr()) - }; - if !found { - return None; - } - Some(unsafe { ::from_message_value(val.assume_init()) }) - } - - fn map_remove(mut map: $pb$::MapMut<$key_t$, Self>, key: $pb$::View<'_, $key_t$>) -> bool { - unsafe { - $pbr$::upb_Map_Delete( - map.as_raw($pbi$::Private), - <$key_t$ as $pbr$::UpbTypeConversions>::to_message_value(key), - $std$::ptr::null_mut()) - } - } - fn map_iter(map: $pb$::MapView<$key_t$, Self>) -> $pb$::MapIter<$key_t$, Self> { - // SAFETY: MapView<'_,..>> guarantees its RawMap outlives '_. - unsafe { - $pb$::MapIter::from_raw($pbi$::Private, $pbr$::RawMapIter::new(map.as_raw($pbi$::Private))) - } - } - - fn map_iter_next<'a>( - iter: &mut $pb$::MapIter<'a, $key_t$, Self> - ) -> Option<($pb$::View<'a, $key_t$>, $pb$::View<'a, Self>)> { - // SAFETY: MapIter<'a, ..> guarantees its RawMapIter outlives 'a. - unsafe { iter.as_raw_mut($pbi$::Private).next_unchecked() } - // SAFETY: MapIter returns key and values message values - // with the variants for K and V active. - .map(|(k, v)| unsafe {( - <$key_t$ as $pbr$::UpbTypeConversions>::from_message_value(k), - ::from_message_value(v), - )}) - } - } - )rs"); - } } }