Skip to content

Commit 458aaf7

Browse files
[FEAT][RUST]Add tvm_ffi::optional in Rust (#630)
This PR supports `ffi::optionalPod`, `ffi::optionalStr` in `Rust`. It keeps the same memory layout as C++'s ffi::optional. --------- Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>
1 parent 331664c commit 458aaf7

5 files changed

Lines changed: 634 additions & 0 deletions

File tree

rust/tvm-ffi/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ pub mod function;
2727
pub mod function_internal;
2828
pub mod macros;
2929
pub mod object;
30+
pub mod option;
3031
pub mod string;
3132
pub mod type_traits;
3233
pub use tvm_ffi_sys;

rust/tvm-ffi/src/option.rs

Lines changed: 350 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,350 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
//! In-place mirrors of C++ `ffi::Optional<T>` (`include/tvm/ffi/optional.h`).
20+
//!
21+
//! `ffi::Optional<T>` stores its value inline, in one of three ABI layouts
22+
//! depending on `T`. The types here decode such a field's bytes directly — no
23+
//! FFI call, no allocation, no reflection getter/setter. Pick the counterpart for
24+
//! the field's `T`:
25+
//!
26+
//! - POD scalar (`i32`, `f64`, `bool`, …) → [`OptionPod<T>`](OptionPod)
27+
//! - `String` → [`OptionStr`]
28+
//! - `ObjectRef` subtype → [`OptionObjRef<T>`](OptionObjRef), an alias of plain
29+
//! `Option<T>` (a single nullable pointer, `nullptr` == `None`)
30+
//!
31+
//! # `OptionPod<T>` — POD scalars
32+
//! Mirrors the `std::optional<T>` fallback as `#[repr(C)] { value: T, engaged:
33+
//! bool }` (payload at offset 0, flag at `size_of::<T>()`), byte-verified against
34+
//! libstdc++/libc++. `T` must implement [`OptionalCompatiblePod`] — the marker trait
35+
//! carried by the fixed set of fixed-width scalars. Read with [`get`](OptionPod::get),
36+
//! write with [`set`](OptionPod::set).
37+
//!
38+
//! # `OptionStr` — `String`
39+
//! The C++ `String` specialization keeps the 16-byte string cell inline and marks
40+
//! `nullopt` with the `type_index == kTVMFFINone` sentinel; [`OptionStr`] wraps
41+
//! [`String`] the same way and reuses its refcounting `Clone`/`Drop`. Borrow with
42+
//! [`as_str`](OptionStr::as_str), write with [`set`](OptionStr::set).
43+
//! (`ffi::Optional<Bytes>` would follow the same pattern.)
44+
45+
use crate::String;
46+
use std::fmt::{self, Debug};
47+
use std::mem::MaybeUninit;
48+
49+
//-----------------------------------------------------
50+
// OptionPod<T> — POD scalars
51+
//-----------------------------------------------------
52+
53+
/// Marker for a POD scalar `T` that can back an [`OptionPod<T>`]; see the
54+
/// [module docs](self).
55+
///
56+
/// Unsafe: an implementor guarantees `T` is trivially copyable and its Rust
57+
/// representation is byte-identical to the C++ field type (`i32` ↔ `int32_t`,
58+
/// `f64` ↔ `double`, …), so the mirror can overlay the C++ `std::optional<T>`.
59+
///
60+
/// Non-scalar payloads are rejected at compile time with a pointer to the
61+
/// right counterpart:
62+
///
63+
/// ```compile_fail,E0277
64+
/// // `Array` is an `ObjectRef` subtype → use `Option<Array<i64>>` (`OptionObjRef`).
65+
/// let _ = tvm_ffi::option::OptionPod::<tvm_ffi::Array<i64>>::none();
66+
/// ```
67+
#[diagnostic::on_unimplemented(
68+
message = "`OptionPod<{Self}>` only mirrors `ffi::Optional` of fixed-width POD scalars",
69+
label = "`{Self}` is not a fixed-width POD scalar",
70+
note = "for an `ObjectRef` subtype use plain `Option<{Self}>` (alias `tvm_ffi::option::OptionObjRef`): the C++ side is a single nullable pointer",
71+
note = "for `String` use `tvm_ffi::option::OptionStr`"
72+
)]
73+
pub unsafe trait OptionalCompatiblePod: Copy {}
74+
75+
/// In-place mirror of C++ `ffi::Optional<T>` for POD `T`, laid out as
76+
/// `std::optional<T>`: `{ T value @0; bool engaged @sizeof(T) }`.
77+
///
78+
/// Layout-compatible with the C++ type; see the [module docs](self).
79+
#[repr(C)]
80+
#[derive(Clone, Copy)]
81+
pub struct OptionPod<T: OptionalCompatiblePod> {
82+
value: MaybeUninit<T>,
83+
engaged: bool,
84+
}
85+
86+
impl<T: OptionalCompatiblePod> OptionPod<T> {
87+
/// Builds an engaged optional holding `value`.
88+
#[inline]
89+
pub fn some(value: T) -> Self {
90+
// Only payload+flag are written; padding isn't part of the ABI.
91+
Self {
92+
value: MaybeUninit::new(value),
93+
engaged: true,
94+
}
95+
}
96+
97+
/// Builds a disengaged optional (`nullopt`).
98+
#[inline]
99+
pub fn none() -> Self {
100+
// Zeroed (not `uninit`) payload keeps the byte-image tests reading init bytes.
101+
Self {
102+
value: MaybeUninit::zeroed(),
103+
engaged: false,
104+
}
105+
}
106+
107+
/// Decodes the value in place. No FFI call, no allocation.
108+
#[inline]
109+
pub fn get(&self) -> Option<T> {
110+
if self.engaged {
111+
// The payload is written whenever `engaged` is set.
112+
Some(unsafe { self.value.assume_init() })
113+
} else {
114+
None
115+
}
116+
}
117+
118+
/// Returns whether a value is present.
119+
#[inline]
120+
pub fn has_value(&self) -> bool {
121+
self.engaged
122+
}
123+
124+
/// Returns whether the optional is `nullopt`.
125+
#[inline]
126+
pub fn is_none(&self) -> bool {
127+
!self.has_value()
128+
}
129+
130+
/// Overwrites the value in place.
131+
///
132+
/// Mirrors C++ assignment: `Some(v)` engages and stores `v`; `None`
133+
/// disengages without touching the payload bytes, as `std::optional::reset`
134+
/// does for trivial `T`.
135+
#[inline]
136+
pub fn set(&mut self, value: Option<T>) {
137+
match value {
138+
Some(v) => {
139+
self.value = MaybeUninit::new(v);
140+
self.engaged = true;
141+
}
142+
None => self.engaged = false,
143+
}
144+
}
145+
}
146+
147+
impl<T: OptionalCompatiblePod> Default for OptionPod<T> {
148+
/// `nullopt`, matching the C++ default constructor.
149+
#[inline]
150+
fn default() -> Self {
151+
Self::none()
152+
}
153+
}
154+
155+
impl<T: OptionalCompatiblePod + PartialEq> PartialEq for OptionPod<T> {
156+
#[inline]
157+
fn eq(&self, other: &Self) -> bool {
158+
self.get() == other.get()
159+
}
160+
}
161+
162+
impl<T: OptionalCompatiblePod + Eq> Eq for OptionPod<T> {}
163+
164+
impl<T: OptionalCompatiblePod> From<Option<T>> for OptionPod<T> {
165+
#[inline]
166+
fn from(value: Option<T>) -> Self {
167+
match value {
168+
Some(v) => Self::some(v),
169+
None => Self::none(),
170+
}
171+
}
172+
}
173+
174+
impl<T: OptionalCompatiblePod> From<OptionPod<T>> for Option<T> {
175+
#[inline]
176+
fn from(value: OptionPod<T>) -> Self {
177+
value.get()
178+
}
179+
}
180+
181+
impl<T: OptionalCompatiblePod + Debug> Debug for OptionPod<T> {
182+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
183+
match self.get() {
184+
Some(v) => write!(f, "OptionPod::Some({v:?})"),
185+
None => f.write_str("OptionPod::None"),
186+
}
187+
}
188+
}
189+
190+
// Registers each supported scalar from one list: the `OptionalCompatiblePod` impl and a
191+
// compile-time guard that `OptionPod<T>` matches the `std::optional<T>` footprint
192+
// (`size == round_up(size_of::<T>()+1, align)`). One list keeps the impl and its
193+
// layout check from drifting.
194+
macro_rules! impl_optional_compatible_pod {
195+
($($t:ty),* $(,)?) => { $(
196+
// Fixed-width scalar; repr matches the C++ field's `std::optional`
197+
// fallback (layout proven by the `const` block below).
198+
unsafe impl OptionalCompatiblePod for $t {}
199+
const _: () = {
200+
let tsz = core::mem::size_of::<$t>();
201+
let tal = core::mem::align_of::<$t>();
202+
let expect = (tsz + 1).div_ceil(tal) * tal;
203+
assert!(core::mem::align_of::<OptionPod<$t>>() == tal);
204+
assert!(core::mem::size_of::<OptionPod<$t>>() == expect);
205+
};
206+
)* };
207+
}
208+
// Keep in sync with the static_assert guard at the end of include/tvm/ffi/optional.h.
209+
impl_optional_compatible_pod!(bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64);
210+
211+
//-----------------------------------------------------
212+
// OptionStr — String
213+
//-----------------------------------------------------
214+
215+
/// In-place mirror of C++ `ffi::Optional<String>`: the 16-byte string cell
216+
/// itself, with `type_index == kTVMFFINone` meaning `nullopt` (the C++
217+
/// String/Bytes spec stores the sentinel in-cell, not as a separate flag).
218+
/// Reuses [`String`]'s `Clone`/`Drop`, whose refcounting is a no-op on the
219+
/// `nullopt` cell (`type_index` below `kTVMFFIStaticObjectBegin`).
220+
#[repr(transparent)]
221+
#[derive(Clone)]
222+
pub struct OptionStr {
223+
// Never handed out or accessed while disengaged (a `nullopt` cell is not a
224+
// valid string).
225+
inner: String,
226+
}
227+
228+
// Must stay 16 bytes to overlay C++ `ffi::Optional<String>` (parity with the POD
229+
// guard in `impl_optional_compatible_pod!`).
230+
const _: () = assert!(
231+
std::mem::size_of::<OptionStr>() == 16
232+
&& std::mem::align_of::<OptionStr>() == std::mem::align_of::<crate::TVMFFIAny>()
233+
);
234+
235+
impl OptionStr {
236+
/// An engaged optional holding `value`.
237+
#[inline]
238+
pub fn some(value: String) -> Self {
239+
Self { inner: value }
240+
}
241+
242+
/// A disengaged optional (`nullopt`).
243+
#[inline]
244+
pub fn none() -> Self {
245+
Self {
246+
inner: String::none_cell(),
247+
}
248+
}
249+
250+
/// Whether a value is present.
251+
#[inline]
252+
pub fn has_value(&self) -> bool {
253+
!self.inner.is_none_cell()
254+
}
255+
256+
/// Whether the optional is `nullopt`.
257+
#[inline]
258+
pub fn is_none(&self) -> bool {
259+
self.inner.is_none_cell()
260+
}
261+
262+
/// Borrows the engaged string as `&str`, or `None` when `nullopt`.
263+
#[inline]
264+
pub fn as_str(&self) -> Option<&str> {
265+
if self.has_value() {
266+
Some(self.inner.as_str())
267+
} else {
268+
None
269+
}
270+
}
271+
272+
/// Takes the value out, consuming self.
273+
#[inline]
274+
pub fn get(self) -> Option<String> {
275+
let OptionStr { inner } = self; // no `Drop` impl, so the move is allowed
276+
if inner.is_none_cell() {
277+
None
278+
} else {
279+
Some(inner)
280+
}
281+
}
282+
283+
/// Overwrites the value in place, dropping the previous one first (dec-ref'd
284+
/// if it was a heap string).
285+
#[inline]
286+
pub fn set(&mut self, value: Option<String>) {
287+
// Assignment drops the old `String` (dec_ref if heap) before moving in the new.
288+
self.inner = match value {
289+
Some(s) => s,
290+
None => String::none_cell(),
291+
};
292+
}
293+
}
294+
295+
impl Default for OptionStr {
296+
/// `nullopt`, matching the C++ default constructor.
297+
#[inline]
298+
fn default() -> Self {
299+
Self::none()
300+
}
301+
}
302+
303+
impl From<Option<String>> for OptionStr {
304+
#[inline]
305+
fn from(value: Option<String>) -> Self {
306+
match value {
307+
Some(s) => Self::some(s),
308+
None => Self::none(),
309+
}
310+
}
311+
}
312+
313+
impl From<OptionStr> for Option<String> {
314+
#[inline]
315+
fn from(value: OptionStr) -> Self {
316+
value.get()
317+
}
318+
}
319+
320+
impl PartialEq for OptionStr {
321+
// NOT derivable: derived eq would run `String::eq` on a `nullopt` cell and
322+
// dereference its null object pointer; `as_str` checks the sentinel first.
323+
#[inline]
324+
fn eq(&self, other: &Self) -> bool {
325+
self.as_str() == other.as_str()
326+
}
327+
}
328+
329+
impl Eq for OptionStr {}
330+
331+
impl Debug for OptionStr {
332+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
333+
match self.as_str() {
334+
Some(s) => write!(f, "OptionStr::Some({s:?})"),
335+
None => f.write_str("OptionStr::None"),
336+
}
337+
}
338+
}
339+
340+
//-----------------------------------------------------
341+
// OptionObjRef — ObjectRef subtypes
342+
//-----------------------------------------------------
343+
344+
/// Alias of [`Option`] for `ObjectRef`-subtype fields.
345+
///
346+
/// C++ `ffi::Optional<SomeRef>` is a single nullable pointer, so `Option<SomeRef>`
347+
/// (niche-optimized over the ref's non-null [`ObjectArc`](crate::ObjectArc)
348+
/// pointer) already mirrors it in place: `None` == `nullptr`. The alias only
349+
/// names that contract for consistency.
350+
pub type OptionObjRef<T> = Option<T>;

rust/tvm-ffi/src/string.rs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,23 @@ impl String {
284284
pub fn as_str(&self) -> &str {
285285
unsafe { std::str::from_utf8_unchecked(self.as_bytes()) }
286286
}
287+
288+
/// Whether the cell is the `nullopt` sentinel (`type_index == kTVMFFINone`).
289+
/// Such a cell is NOT a valid string; only [`OptionStr`](crate::option::OptionStr)
290+
/// holds one, and never calls string accessors on it.
291+
#[inline]
292+
pub(crate) fn is_none_cell(&self) -> bool {
293+
self.data.type_index == TypeIndex::kTVMFFINone as i32
294+
}
295+
296+
/// A `nullopt` sentinel cell for [`OptionStr`](crate::option::OptionStr).
297+
/// `TVMFFIAny::new()` is exactly the all-zero `kTVMFFINone` cell.
298+
#[inline]
299+
pub(crate) fn none_cell() -> Self {
300+
Self {
301+
data: TVMFFIAny::new(),
302+
}
303+
}
287304
}
288305

289306
impl<T> From<T> for String

0 commit comments

Comments
 (0)