File: arena.rs

package info (click to toggle)
chromium 139.0.7258.127-1
  • links: PTS, VCS
  • area: main
  • in suites:
  • size: 6,122,068 kB
  • sloc: cpp: 35,100,771; ansic: 7,163,530; javascript: 4,103,002; python: 1,436,920; asm: 946,517; xml: 746,709; pascal: 187,653; perl: 88,691; sh: 88,436; objc: 79,953; sql: 51,488; cs: 44,583; fortran: 24,137; makefile: 22,147; tcl: 15,277; php: 13,980; yacc: 8,984; ruby: 7,485; awk: 3,720; lisp: 3,096; lex: 1,327; ada: 727; jsp: 228; sed: 36
file content (234 lines) | stat: -rw-r--r-- 8,547 bytes parent folder | download | duplicates (6)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
// Protocol Buffers - Google's data interchange format
// Copyright 2024 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 super::opaque_pointee::opaque_pointee;
use core::cell::UnsafeCell;
use core::marker::PhantomData;
use core::mem::{align_of, align_of_val, size_of_val, MaybeUninit};
use core::ptr::{self, NonNull};
use core::slice;

opaque_pointee!(upb_Arena);
pub type RawArena = NonNull<upb_Arena>;

/// See `upb/port/def.inc`.
const UPB_MALLOC_ALIGN: usize = 8;
const _CHECK_UPB_MALLOC_ALIGN_AT_LEAST_POINTER_ALIGNED: () =
    assert!(UPB_MALLOC_ALIGN >= align_of::<*const ()>());

/// A wrapper over a `upb_Arena`.
///
/// This is not a safe wrapper per se, because the allocation functions still
/// have sharp edges (see their safety docs for more info).
///
/// This is an owning type and will automatically free the arena when
/// dropped.
///
/// Note that this type is not `Sync` as it implements unsynchronized interior
/// mutability. The upb_Arena C object could be understood as being Sync (at
/// least vacuously under current API since there are not any const upb_Arena*
/// API functions), but the Rust Arena is necessarily expressed as interior
/// mutability (&self rather than &mut self receivers) See https://doc.rust-lang.org/nomicon/lifetime-mismatch.html and
/// https://blog.reverberate.org/2021/12/19/arenas-and-rust.html, and the
/// 'known problems' section of https://rust-lang.github.io/rust-clippy/master/index.html#/mut_from_ref.
#[derive(Debug)]
pub struct Arena {
    // Safety invariant: this must always be a valid arena
    raw: RawArena,
    _not_sync: PhantomData<UnsafeCell<()>>,
}

// SAFETY: `Arena` uniquely holds the underlying RawArena and has no
// thread-local data.
unsafe impl Send for Arena {}

impl Arena {
    /// Allocates a fresh arena.
    #[inline]
    pub fn new() -> Self {
        #[inline(never)]
        #[cold]
        fn arena_new_failed() -> ! {
            panic!("Could not create a new UPB arena");
        }

        // SAFETY:
        // - `upb_Arena_New` is assumed to be implemented correctly and always sound to
        //   call; if it returned a non-null pointer, it is a valid arena.
        unsafe {
            let Some(raw) = upb_Arena_New() else { arena_new_failed() };
            Self { raw, _not_sync: PhantomData }
        }
    }

    /// # Safety
    /// - The `raw_arena` must point to a valid arena.
    /// - The caller must ensure that the Arena's destructor does not run.
    pub unsafe fn from_raw(raw_arena: RawArena) -> Self {
        Arena { raw: raw_arena, _not_sync: PhantomData }
    }

    /// Returns the raw, UPB-managed pointer to the arena.
    #[inline]
    pub fn raw(&self) -> RawArena {
        self.raw
    }

    /// Allocates some memory on the arena. Returns None if the allocation
    /// failed.
    ///
    /// # Safety
    ///
    /// - `align` must be less than `UPB_MALLOC_ALIGN`.
    #[allow(clippy::mut_from_ref)]
    #[inline]
    pub unsafe fn alloc(&self, size: usize, align: usize) -> Option<&mut [MaybeUninit<u8>]> {
        debug_assert!(align <= UPB_MALLOC_ALIGN);
        // SAFETY: `self.raw` is a valid UPB arena
        let ptr = unsafe { upb_Arena_Malloc(self.raw, size) };

        if ptr.is_null() {
            None
        } else {
            // SAFETY:
            // - `upb_Arena_Malloc` promises that if the return pointer is non-null, it is
            //   dereferencable for `size` bytes and has an alignment of `UPB_MALLOC_ALIGN`
            //   until the arena is destroyed.
            // - `[MaybeUninit<u8>]` has no alignment requirement, and `ptr` is aligned to a
            //   `UPB_MALLOC_ALIGN` boundary.
            Some(unsafe { slice::from_raw_parts_mut(ptr.cast(), size) })
        }
    }

    /// Same as alloc() but panics if `align > UPB_MALLOC_ALIGN`.
    #[allow(clippy::mut_from_ref)]
    #[inline]
    pub fn checked_alloc(&self, size: usize, align: usize) -> Option<&mut [MaybeUninit<u8>]> {
        assert!(align <= UPB_MALLOC_ALIGN);
        // SAFETY: align <= UPB_MALLOC_ALIGN asserted.
        unsafe { self.alloc(size, align) }
    }

    /// Copies the T into this arena and returns a pointer to the T data inside
    /// the arena. Returns None if the allocation failed.
    pub fn copy_in<'a, T: Copy>(&'a self, data: &T) -> Option<&'a T> {
        let size = size_of_val(data);
        let align = align_of_val(data);

        self.checked_alloc(size, align).map(|alloc| {
            // SAFETY:
            // - alloc is valid for `size` bytes and is the uninit bytes are written to not
            //   read from until written.
            // - T is copy so copying the bytes of the value is sound.
            unsafe {
                let alloc = alloc.as_mut_ptr().cast::<MaybeUninit<T>>();
                &*(*alloc).write(*data)
            }
        })
    }

    /// Copies the str into this arena and returns a pointer to the T data
    /// inside the arena. Returns None if the allocation failed.
    pub fn copy_str_in<'a>(&'a self, s: &str) -> Option<&'a str> {
        self.copy_slice_in(s.as_bytes()).map(|copied_bytes| {
            // SAFETY: `copied_bytes` has same contents as `s` and so must meet &str
            // criteria.
            unsafe { core::str::from_utf8_unchecked(copied_bytes) }
        })
    }

    /// Copies the slice into this arena and returns a pointer to the T data
    /// inside the arena. Returns None if the allocation failed.
    pub fn copy_slice_in<'a, T: Copy>(&'a self, data: &[T]) -> Option<&'a [T]> {
        let size = size_of_val(data);
        let align = align_of_val(data);
        self.checked_alloc(size, align).map(|alloc| {
            let alloc: *mut T = alloc.as_mut_ptr().cast();
            // SAFETY:
            // - uninit_alloc is valid for `layout.len()` bytes and is the uninit bytes are
            //   written to not read from until written.
            // - T is copy so copying the bytes of the values is sound.
            unsafe {
                ptr::copy_nonoverlapping(data.as_ptr(), alloc, data.len());
                slice::from_raw_parts(alloc, data.len())
            }
        })
    }

    /// Fuse two arenas so they share the same lifetime.
    ///
    /// `fuse` will make it so that the memory allocated by `self` or `other` is
    /// guaranteed to last until both `self` and `other` have been dropped.
    /// The pointers returned by `Arena::alloc` will continue to be valid so
    /// long as either `self` or `other` has not been dropped.
    pub fn fuse(&self, other: &Arena) {
        // SAFETY: `self.raw()` and `other.raw()` are both valid UPB arenas.
        let success = unsafe { upb_Arena_Fuse(self.raw(), other.raw()) };
        if !success {
            // Fusing can fail if any of the arenas has an initial block i.e. the arena is
            // backed by a preallocated chunk of memory that it doesn't own and thus cannot
            // lifetime extend. This function panics because this is typically not a
            // recoverable error but a logic bug in a program.
            panic!("Could not fuse two UPB arenas.");
        }
    }
}

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

impl Drop for Arena {
    #[inline]
    fn drop(&mut self) {
        unsafe {
            upb_Arena_Free(self.raw);
        }
    }
}

extern "C" {
    // `Option<NonNull<T: Sized>>` is ABI-compatible with `*mut T`
    fn upb_Arena_New() -> Option<RawArena>;
    fn upb_Arena_Free(arena: RawArena);
    fn upb_Arena_Malloc(arena: RawArena, size: usize) -> *mut u8;
    fn upb_Arena_Fuse(arena1: RawArena, arena2: RawArena) -> bool;
}

#[cfg(test)]
mod tests {
    use super::*;
    use googletest::gtest;

    #[gtest]
    fn assert_arena_linked() {
        use crate::assert_linked;
        assert_linked!(upb_Arena_New);
        assert_linked!(upb_Arena_Free);
        assert_linked!(upb_Arena_Malloc);
        assert_linked!(upb_Arena_Fuse);
    }

    #[gtest]
    fn raw_ffi_test() {
        // SAFETY: FFI unit test uses C API under expected patterns.
        unsafe {
            let arena = upb_Arena_New().unwrap();
            let bytes = upb_Arena_Malloc(arena, 3);
            *bytes.add(2) = 7;
            upb_Arena_Free(arena);
        }
    }

    #[gtest]
    fn test_arena_new_and_free() {
        let arena = Arena::new();
        drop(arena);
    }
}