Skip to content

Commit

Permalink
Rust: create blanket implementation of ProxiedInMapValue for upb-ba…
Browse files Browse the repository at this point in the history 
…cked 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
  • Loading branch information
@acozzette @copybara-github
acozzette authored and copybara-github committed Sep 5, 2024
1 parent 37f3ced commit c3df359
Show file tree
Hide file tree
Showing 2 changed files with 101 additions and 96 deletions.
106 changes: 101 additions & 5 deletions rust/upb.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -797,6 +797,102 @@ macro_rules! impl_ProxiedInMapValue_for_key_types {

impl_ProxiedInMapValue_for_key_types!(i32, u32, i64, u64, bool, ProtoString);

impl<Key, MessageType> ProxiedInMapValue<Key> for MessageType
where
Key: Proxied + UpbTypeConversions,
MessageType: Proxied + UpbTypeConversions + Message,
{
fn map_new(_private: Private) -> Map<Key, Self> {
let arena = Arena::new();
let raw = unsafe {
upb_Map_New(
arena.raw(),
<Key as UpbTypeConversions>::upb_type(),
<Self as UpbTypeConversions>::upb_type(),
)
};

Map::from_inner(Private, InnerMap::new(raw, arena))
}

unsafe fn map_free(_private: Private, _map: &mut Map<Key, Self>) {
// No-op: the memory will be dropped by the arena.
}

fn map_clear(mut map: MapMut<Key, Self>) {
unsafe {
upb_Map_Clear(map.as_raw(Private));
}
}

fn map_len(map: MapView<Key, Self>) -> usize {
unsafe { upb_Map_Size(map.as_raw(Private)) }
}

fn map_insert(
mut map: MapMut<Key, Self>,
key: View<'_, Key>,
value: impl IntoProxied<Self>,
) -> bool {
let arena = map.inner(Private).raw_arena();
unsafe {
upb_Map_InsertAndReturnIfInserted(
map.as_raw(Private),
<Key as UpbTypeConversions>::to_message_value(key),
<Self as UpbTypeConversions>::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<View<'a, Self>> {
let mut val = MaybeUninit::uninit();
let found = unsafe {
upb_Map_Get(
map.as_raw(Private),
<Key as UpbTypeConversions>::to_message_value(key),
val.as_mut_ptr(),
)
};
if !found {
return None;
}
Some(unsafe { <Self as UpbTypeConversions>::from_message_value(val.assume_init()) })
}

fn map_remove(mut map: MapMut<Key, Self>, key: View<'_, Key>) -> bool {
unsafe {
upb_Map_Delete(
map.as_raw(Private),
<Key as UpbTypeConversions>::to_message_value(key),
ptr::null_mut(),
)
}
}
fn map_iter(map: MapView<Key, Self>) -> MapIter<Key, Self> {
// 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<K, V> returns key and values message values
// with the variants for K and V active.
.map(|(k, v)| unsafe {
(
<Key as UpbTypeConversions>::from_message_value(k),
<Self as UpbTypeConversions>::from_message_value(v),
)
})
}
}

/// `upb_Map_Insert`, but returns a `bool` for whether insert occurred.
///
/// Returns `true` if the entry was newly inserted.
Expand Down
91 changes: 0 additions & 91 deletions src/google/protobuf/compiler/rust/message.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
<Self as $pbr$::UpbTypeConversions>::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<Self>) -> 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),
<Self as $pbr$::UpbTypeConversions>::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 { <Self as $pbr$::UpbTypeConversions>::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<K, V> 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),
<Self as $pbr$::UpbTypeConversions>::from_message_value(v),
)})
}
}
)rs");
}
}
}

Expand Down

0 comments on commit c3df359

Please sign in to comment.