|
| 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>; |
0 commit comments