rust/upb/reflection.rs

// Protocol Buffers - Google's data interchange format
// Copyright 2026 Google LLC.  All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd

use core::ffi::{c_int, CStr};
use core::marker::PhantomData;

// In Bazel `upb` and `reflection` are dependency crates. In Cargo this file is compiled as
// <crate_root>::upb::reflection, so bring the parent and sibling modules into scope under the same
// names.
#[cfg(not(bzl))]
use super::super::upb;
#[cfg(not(bzl))]
use super::sys::reflection;

use reflection::def_pool::{
    upb_DefPool_FindMessageByNameWithSize, upb_DefPool_Free, upb_DefPool_Init_New,
    upb_DefPool_LoadDefInit, upb_DefPool_New, upb_MiniTableFile_New, RawDefPool,
};
use reflection::message_def::RawMessageDef;
use reflection::upb_TextEncode;

use upb::{
    Arena, MessagePtr, RawMiniTable, RawMiniTableEnum, RawMiniTableExtension, StringView,
    TextEncodeOptions,
};

pub use reflection::def_pool::upb_DefPool_Init;
pub use reflection::def_pool::RawDefPoolInit as DefPoolInitPtr;
pub use reflection::def_pool::RawMiniTableFile as MiniTableFilePtr;

/// A wrapper over a `upb_DefPool`.
///
/// This is an owning type and will automatically free the pool when
/// dropped.
pub struct DefPool {
    raw: RawDefPool,
}

impl Default for DefPool {
    fn default() -> Self {
        Self::new()
    }
}

impl DefPool {
    /// Allocates a new pool.
    #[inline]
    pub fn new() -> Self {
        // SAFETY: `upb_DefPool_New` has no preconditions.
        let raw = unsafe { upb_DefPool_New() }.expect("DefPool allocation failed");
        Self { raw }
    }

    /// Loads a generated descriptor, and everything it imports, into the pool.
    ///
    /// # Safety
    /// - `init` must point to a valid `upb_DefPool_Init`, and so must every init reachable from it
    ///   through its dependencies according to the correct orderings and counts.
    pub unsafe fn load_def_init(&mut self, init: *const upb_DefPool_Init) -> bool {
        // SAFETY:
        // - `self.raw` is a live pool; only `Drop` frees it.
        // - `init` is valid per the safety requirements of this function.
        unsafe { upb_DefPool_LoadDefInit(self.raw, init) }
    }

    /// Looks up a message by fully-qualified name in the pool.
    pub fn find_message_by_name(&self, full_name: &str) -> Option<MessageDef<'_>> {
        // SAFETY:
        // - `self.raw` is a live pool; only `Drop` frees it.
        // - `full_name.as_ptr()` is readable for `full_name.len()` bytes.
        let raw = unsafe {
            upb_DefPool_FindMessageByNameWithSize(
                self.raw,
                full_name.as_ptr().cast(),
                full_name.len(),
            )
        }?;

        // SAFETY: `raw` was allocated by `self`, and the returned `MessageDef` borrows `self`.
        Some(unsafe { MessageDef::from_raw(raw) })
    }

    /// Returns the internal NonNull representation of the pool.
    #[inline]
    pub(crate) fn raw(&self) -> RawDefPool {
        self.raw
    }
}

impl Drop for DefPool {
    fn drop(&mut self) {
        // SAFETY: The pointer was generated by upb_DefPool_New and hasn't been freed.
        unsafe { upb_DefPool_Free(self.raw) }
    }
}

/// A borrowed message descriptor tied to the lifetime of a DefPool. No drop implementation is
/// provided because this is not an owning type and cannot be individually freed.
#[derive(Copy, Clone)]
pub struct MessageDef<'pool> {
    raw: RawMessageDef,
    _marker: PhantomData<&'pool DefPool>,
}

impl<'pool> MessageDef<'pool> {
    /// # Safety
    /// The caller must guarantee `raw` was allocated by the pool providing the `'pool` lifetime.
    pub unsafe fn from_raw(raw: RawMessageDef) -> Self {
        Self { raw, _marker: PhantomData }
    }

    /// Returns the internal NonNull representation of the descriptor.
    #[inline]
    pub fn raw(&self) -> RawMessageDef {
        self.raw
    }
}

/// Returns a string of field number to value entries of a message.
///
/// # Safety
/// - `msg` must be legally dereferenceable.
pub unsafe fn text_encode<'pool, T>(
    msg: MessagePtr<T>,
    def: MessageDef<'pool>,
    ext_pool: Option<&DefPool>,
    options: TextEncodeOptions,
) -> String {
    let msg = msg.raw();
    let def_raw = def.raw();
    let ext_pool_raw = ext_pool.map(|p| p.raw());
    let options = options.to_bits();

    // Only find out the length first to then allocate a buffer of the minimum size
    // needed.
    // SAFETY:
    // - `msg` is valid per safety requirements of this function.
    // - `def` is valid per safety requirements of this function.
    // - `ext_pool` is valid per safety requirements of this function.
    // - `buf` is null and `size` is 0.
    let min_len =
        unsafe { upb_TextEncode(msg, def_raw, ext_pool_raw, options, core::ptr::null_mut(), 0) };
    assert!(min_len < isize::MAX as usize);

    // +1 for the trailing NULL
    let mut buf = vec![0u8; min_len + 1];

    // SAFETY:
    // - `msg`, `def`, `ext_pool` are valid per safety requirements of this function.
    // - `buf` is legally writable for `buf.len()` bytes.
    let written_len =
        unsafe { upb_TextEncode(msg, def_raw, ext_pool_raw, options, buf.as_mut_ptr(), buf.len()) };
    assert_eq!(min_len, written_len);

    // Drop the trailing NULL written by `upb_TextEncode`.
    buf.truncate(written_len);
    String::from_utf8_lossy(buf.as_slice()).into_owned()
}

fn copy_slice_or_null<T: Copy>(arena: &Arena, slice: &[T]) -> *const T {
    if slice.is_empty() {
        return core::ptr::null();
    }
    arena.copy_slice_in(slice).expect("arena allocation failed").as_ptr()
}

/// Builds a `upb_MiniTableFile` in `arena` from the MiniTables of one .proto file.
pub unsafe fn build_mini_table_file(
    arena: &Arena,
    msgs: &[RawMiniTable],
    enums: &[RawMiniTableEnum],
    exts: &[RawMiniTableExtension],
) -> MiniTableFilePtr {
    // Copy the minitables that msgs, enums and exts point to into the arena.
    let msgs_ptr = copy_slice_or_null(arena, msgs);
    let enums_ptr = copy_slice_or_null(arena, enums);
    let exts_ptr = copy_slice_or_null(arena, exts);

    // SAFETY:
    // - `arena` is live for the duration of the call.
    // - each array is readable for the count passed alongside it.
    unsafe {
        upb_MiniTableFile_New(
            arena.raw(),
            msgs_ptr,
            msgs.len() as c_int,
            enums_ptr,
            enums.len() as c_int,
            exts_ptr,
            exts.len() as c_int,
        )
    }
    .expect("arena allocation failed")
}

/// Builds a upb_DefPool_Init
///
/// # Safety
/// - `layout` and every init in `deps` must outlive any `DefPool` this init is loaded into, as must
///   `arena`: upb keeps pointing at them rather than copying them.
pub unsafe fn build_def_pool_init(
    arena: &Arena,
    filename: &'static CStr,
    descriptor: &'static [u8],
    deps: &[DefPoolInitPtr],
    layout: MiniTableFilePtr,
) -> DefPoolInitPtr {
    // Add a `None` to the end of the deps to terminate the list.
    let mut deps_terminated: Vec<Option<DefPoolInitPtr>> = Vec::with_capacity(deps.len() + 1);
    for dep in deps {
        deps_terminated.push(Some(*dep));
    }

    deps_terminated.push(None);
    let deps_ptr = arena.copy_slice_in(&deps_terminated).expect("arena allocation failed").as_ptr();

    // SAFETY:
    // - `arena` is live for the duration of the call.
    // - `filename` and `descriptor` are readable for 'static.
    // - `deps_ptr` points at a null-terminated array. It is cast to a mutable pointer only to match
    //   the C signature; upb never writes through it.
    unsafe {
        upb_DefPool_Init_New(
            arena.raw(),
            filename.as_ptr(),
            StringView::from(descriptor),
            deps_ptr as *mut Option<DefPoolInitPtr>,
            layout.as_ptr(),
        )
    }
    .expect("arena allocation failed")
}