Ubospica commented on code in PR #630: URL: https://github.com/apache/tvm-ffi/pull/630#discussion_r3521327042
########## rust/tvm-ffi/src/collections/optional.rs: ########## @@ -0,0 +1,328 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +//! In-place mirrors of C++ `ffi::Optional<T>` (`include/tvm/ffi/optional.h`). +//! +//! `ffi::Optional<T>` stores its value inline, in one of three ABI layouts +//! depending on `T`. The types here decode such a field's bytes directly — no +//! FFI call, no allocation, no reflection getter/setter. Pick the counterpart for +//! the field's `T`: +//! +//! - POD scalar (`i32`, `f64`, `bool`, …) → [`Optional<T>`](Optional) +//! - `String` → [`OptionalStr`] +//! - `ObjectRef` subtype → plain `Option<SomeRef>` (a single nullable pointer, +//! `nullptr` == `None`; no dedicated type needed) +//! +//! # `Optional<T>` — POD scalars +//! Mirrors the `std::optional<T>` fallback as `#[repr(C)] { value: T, engaged: +//! bool }` (payload at offset 0, flag at `size_of::<T>()`), byte-verified against +//! libstdc++/libc++. `T` must implement [`OptionalPod`] — the marker trait +//! carried by the fixed set of fixed-width scalars. Read with [`get`](Optional::get), +//! write with [`set`](Optional::set); `set` takes `&self` via interior mutability, +//! so a shared `&Optional<T>` aliasing a C++ field stays writable (hence `!Sync`). +//! +//! # `OptionalStr` — `String` +//! The C++ `String` specialization keeps the 16-byte string cell inline and marks +//! `nullopt` with the `type_index == kTVMFFINone` sentinel; [`OptionalStr`] wraps +//! [`String`] the same way and reuses its refcounting `Clone`/`Drop`. Borrow with +//! [`as_str`](OptionalStr::as_str), write with [`set`](OptionalStr::set) — which +//! takes `&mut self`, since a shared-ref setter could drop the backing string +//! under a live `&str`. (`ffi::Optional<Bytes>` would follow the same pattern.) + +use crate::String; +use std::cell::UnsafeCell; +use std::fmt::{self, Debug}; +use std::mem::MaybeUninit; + +//----------------------------------------------------- +// Optional<T> — POD scalars +//----------------------------------------------------- + +/// Marker for a POD scalar `T` that can back an [`Optional<T>`]; see the +/// [module docs](self). +/// +/// Unsafe: an implementor guarantees `T` is trivially copyable and its Rust +/// representation is byte-identical to the C++ field type (`i32` ↔ `int32_t`, +/// `f64` ↔ `double`, …), so the mirror can overlay the C++ `std::optional<T>`. +pub unsafe trait OptionalPod: Copy {} + +/// Layout-mirror of `std::optional<T>`: `{ T value @0; bool engaged @sizeof(T) }`. +#[repr(C)] +struct OptionalCell<T: OptionalPod> { + value: MaybeUninit<T>, + engaged: bool, +} Review Comment: We are currently assuming the memory layout from C++ side. Shall we add a static assert in C++ side, like ``` static_assert(sizeof(ffi::Optional<int32_t>) == 8) ``` To avoid potential compilation errors? ########## rust/tvm-ffi/src/collections/optional.rs: ########## @@ -0,0 +1,328 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +//! In-place mirrors of C++ `ffi::Optional<T>` (`include/tvm/ffi/optional.h`). +//! +//! `ffi::Optional<T>` stores its value inline, in one of three ABI layouts +//! depending on `T`. The types here decode such a field's bytes directly — no +//! FFI call, no allocation, no reflection getter/setter. Pick the counterpart for +//! the field's `T`: +//! +//! - POD scalar (`i32`, `f64`, `bool`, …) → [`Optional<T>`](Optional) +//! - `String` → [`OptionalStr`] +//! - `ObjectRef` subtype → plain `Option<SomeRef>` (a single nullable pointer, +//! `nullptr` == `None`; no dedicated type needed) +//! +//! # `Optional<T>` — POD scalars +//! Mirrors the `std::optional<T>` fallback as `#[repr(C)] { value: T, engaged: +//! bool }` (payload at offset 0, flag at `size_of::<T>()`), byte-verified against +//! libstdc++/libc++. `T` must implement [`OptionalPod`] — the marker trait +//! carried by the fixed set of fixed-width scalars. Read with [`get`](Optional::get), +//! write with [`set`](Optional::set); `set` takes `&self` via interior mutability, +//! so a shared `&Optional<T>` aliasing a C++ field stays writable (hence `!Sync`). +//! +//! # `OptionalStr` — `String` +//! The C++ `String` specialization keeps the 16-byte string cell inline and marks +//! `nullopt` with the `type_index == kTVMFFINone` sentinel; [`OptionalStr`] wraps +//! [`String`] the same way and reuses its refcounting `Clone`/`Drop`. Borrow with +//! [`as_str`](OptionalStr::as_str), write with [`set`](OptionalStr::set) — which +//! takes `&mut self`, since a shared-ref setter could drop the backing string +//! under a live `&str`. (`ffi::Optional<Bytes>` would follow the same pattern.) + +use crate::String; +use std::cell::UnsafeCell; +use std::fmt::{self, Debug}; +use std::mem::MaybeUninit; + +//----------------------------------------------------- +// Optional<T> — POD scalars +//----------------------------------------------------- + +/// Marker for a POD scalar `T` that can back an [`Optional<T>`]; see the +/// [module docs](self). +/// +/// Unsafe: an implementor guarantees `T` is trivially copyable and its Rust +/// representation is byte-identical to the C++ field type (`i32` ↔ `int32_t`, +/// `f64` ↔ `double`, …), so the mirror can overlay the C++ `std::optional<T>`. +pub unsafe trait OptionalPod: Copy {} + +/// Layout-mirror of `std::optional<T>`: `{ T value @0; bool engaged @sizeof(T) }`. +#[repr(C)] +struct OptionalCell<T: OptionalPod> { + value: MaybeUninit<T>, + engaged: bool, +} + +/// In-place mirror of C++ `ffi::Optional<T>` for POD `T`. +/// +/// Layout-compatible with the C++ type; see the [module docs](self). +#[repr(transparent)] +pub struct Optional<T: OptionalPod> { + cell: UnsafeCell<OptionalCell<T>>, +} + +impl<T: OptionalPod> Optional<T> { + /// Builds an engaged optional holding `value`. + #[inline] + pub fn some(value: T) -> Self { + // Only payload+flag are written; padding isn't part of the ABI. + Self { + cell: UnsafeCell::new(OptionalCell { + value: MaybeUninit::new(value), + engaged: true, + }), + } + } + + /// Builds a disengaged optional (`nullopt`). + #[inline] + pub fn none() -> Self { + // Zeroed (not `uninit`) payload keeps the byte-image tests reading init bytes. + Self { + cell: UnsafeCell::new(OptionalCell { + value: MaybeUninit::zeroed(), + engaged: false, + }), + } + } + + /// Decodes the value in place. No FFI call, no allocation. + #[inline] + pub fn get(&self) -> Option<T> { + // Read the payload only after confirming `engaged` (the cell is always initialized). + let cell = unsafe { &*self.cell.get() }; + if cell.engaged { + Some(unsafe { cell.value.assume_init() }) + } else { + None + } + } + + /// Returns whether a value is present. + #[inline] + pub fn has_value(&self) -> bool { + unsafe { (*self.cell.get()).engaged } + } + + /// Returns whether the optional is `nullopt`. + #[inline] + pub fn is_none(&self) -> bool { + !self.has_value() + } + + /// Overwrites the value in place through a shared reference. + /// + /// Mirrors C++ assignment: `Some(v)` engages and stores `v`; `None` + /// disengages without touching the payload bytes, as `std::optional::reset` + /// does for trivial `T`. + #[inline] + pub fn set(&self, value: Option<T>) { + // Interior mutation via `UnsafeCell`; caller must not race (`!Sync`). + let cell = unsafe { &mut *self.cell.get() }; + match value { + Some(v) => { + cell.value = MaybeUninit::new(v); + cell.engaged = true; + } + None => cell.engaged = false, + } + } +} Review Comment: pub fn set(&mut self, value: Option<T>) { match value { Some(v) => { self.cell.value = MaybeUninit::new(v); self.cell.engaged = true; } None => self.cell.engaged = false, } } Shall we just use a mutable set? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
