This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new ab87e916 [Rust] Data and ABI soundness fixes: open DLDeviceType,
checked tensor sizes, unsafe AsDLDataType (#957)
ab87e916 is described below
commit ab87e91655f172b16d888f84b2c16c40fddbab7d
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Nov 7 08:04:37 2026 -0500
[Rust] Data and ABI soundness fixes: open DLDeviceType, checked tensor
sizes, unsafe AsDLDataType (#957)
This PR fixes data and ABI soundness issues in the Rust bindings, where
safe code could create invalid values or read out of bounds.
- **`DLDevice::device_type` is an `i32`.** Device type codes are an open
set, so a code read from a foreign `DLDevice` may not be a
`DLDeviceType` variant. The field is now an `i32`, as `DLDataType::code`
is a `u8`, and `DLDeviceType` stays the enum of known codes.
- **Checked tensor sizes.** The element count and byte size of a tensor
wrapped on overflow, and a negative extent became a huge count, so
`Tensor::from_nd_alloc(CPUNDAlloc, &[(1 << 61) + 1], f64, cpu)`
allocated 8 bytes behind a 2^61-element slice. Both now use checked
arithmetic: `from_nd_alloc` panics on such a shape before allocating, as
`Vec::with_capacity` does on capacity overflow, and `Tensor::from_slice`
returns a `ValueError`.
- **`AsDLDataType` is an `unsafe trait`.** `Tensor::data_as_slice::<T>`
reads data of dtype `T::DL_DATA_TYPE` as `T`, so a safe implementation
naming the dtype of an unrelated type let safe code read arbitrary bits
as `T`. The trait now states the layout contract that `data_as_slice`
relies on.
- **`TVMFFIByteArray::as_str` is `unsafe`.** It views a raw pointer and
length as text. Converting a raw C string argument that is not UTF-8 to
a `String` now fails as a type mismatch instead of panicking.
`String::as_str` stays unchecked, since an `ffi.String` holds UTF-8 text
by the ABI.
- **Null pointers with zero length.** C passes an empty array as a null
pointer with length zero, which a slice cannot hold. A crate-internal
`util::slice_from_nullable_raw_parts` allows that case and replaces the
hand-written checks for packed call arguments, tensor shapes and data,
`Shape`, `List` cells, `Bytes` from a byte array, and type field tables.
- **`Map` accepts an `ffi.Dict`.** As C++ `Map<K, V>` does, converting
an `ffi.Dict` to a `Map` copies its entries into a new map, converting
each to `K` and `V`.
Source-breaking changes:
- `DLDevice::device_type` is an `i32`: compare it with
`DLDeviceType::kDLCPU as i32`, and drop `as i32` where it is passed on
as a code.
- Implementations of `AsDLDataType` need `unsafe impl`.
- Calls to `TVMFFIByteArray::as_str` need an `unsafe` block.
---
rust/tvm-ffi-sys/src/c_api.rs | 6 +-
rust/tvm-ffi-sys/src/dlpack.rs | 5 +-
rust/tvm-ffi/src/collections/list.rs | 5 +-
rust/tvm-ffi/src/collections/map.rs | 74 +++++++++++++------
rust/tvm-ffi/src/collections/shape.rs | 6 +-
rust/tvm-ffi/src/collections/tensor.rs | 102 +++++++++++++++------------
rust/tvm-ffi/src/device.rs | 10 +--
rust/tvm-ffi/src/dtype.rs | 9 ++-
rust/tvm-ffi/src/error.rs | 9 ++-
rust/tvm-ffi/src/extra/structural_common.rs | 5 +-
rust/tvm-ffi/src/extra/structural_mutate.rs | 3 +-
rust/tvm-ffi/src/function.rs | 2 +-
rust/tvm-ffi/src/function_internal.rs | 11 +++
rust/tvm-ffi/src/lib.rs | 1 +
rust/tvm-ffi/src/macros.rs | 3 +-
rust/tvm-ffi/src/reflection/accessor.rs | 19 +++--
rust/tvm-ffi/src/reflection/unsafe_init.rs | 4 +-
rust/tvm-ffi/src/string.rs | 5 +-
rust/tvm-ffi/src/type_traits.rs | 4 +-
rust/tvm-ffi/src/util.rs | 47 ++++++++++++
rust/tvm-ffi/tests/test_any.rs | 11 +--
rust/tvm-ffi/tests/test_map.rs | 12 ++++
rust/tvm-ffi/tests/test_object_reflection.rs | 37 ++++++----
rust/tvm-ffi/tests/test_rs_object.rs | 13 ++--
rust/tvm-ffi/tests/test_tensor.rs | 2 +-
25 files changed, 273 insertions(+), 132 deletions(-)
diff --git a/rust/tvm-ffi-sys/src/c_api.rs b/rust/tvm-ffi-sys/src/c_api.rs
index 199ee58c..0eebb92f 100644
--- a/rust/tvm-ffi-sys/src/c_api.rs
+++ b/rust/tvm-ffi-sys/src/c_api.rs
@@ -246,7 +246,11 @@ impl TVMFFIByteArray {
///
/// # Returns
/// * `&str` - The converted str view.
- pub fn as_str(&self) -> &str {
+ ///
+ /// # Safety
+ /// `data` must be null or point to `size` bytes of UTF-8 text that stay
+ /// valid and unchanged while the result is used.
+ pub unsafe fn as_str(&self) -> &str {
// An empty array may hold a null pointer, as the type table stores an
// empty doc string.
if self.data.is_null() {
diff --git a/rust/tvm-ffi-sys/src/dlpack.rs b/rust/tvm-ffi-sys/src/dlpack.rs
index 2037b43f..96d0b71f 100644
--- a/rust/tvm-ffi-sys/src/dlpack.rs
+++ b/rust/tvm-ffi-sys/src/dlpack.rs
@@ -49,7 +49,8 @@ pub enum DLDeviceType {
#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct DLDevice {
- pub device_type: DLDeviceType,
+ /// A [`DLDeviceType`] code, an `i32` so the codes stay an open set.
+ pub device_type: i32,
pub device_id: i32,
}
@@ -111,7 +112,7 @@ pub struct DLTensor {
impl DLDevice {
pub fn new(device_type: DLDeviceType, device_id: i32) -> Self {
Self {
- device_type,
+ device_type: device_type as i32,
device_id,
}
}
diff --git a/rust/tvm-ffi/src/collections/list.rs
b/rust/tvm-ffi/src/collections/list.rs
index e1650779..799c725f 100644
--- a/rust/tvm-ffi/src/collections/list.rs
+++ b/rust/tvm-ffi/src/collections/list.rs
@@ -231,10 +231,7 @@ unsafe impl<T: ContainerElement> ObjectRefCore for List<T>
{
/// mutated while the slice is alive.
pub(crate) unsafe fn seq_cells(data: &TVMFFIAny) -> &[TVMFFIAny] {
let seq = &*(data.data_union.v_obj as *const ListObj);
- if *seq.size.get() == 0 {
- return &[];
- }
- std::slice::from_raw_parts(
+ crate::util::slice_from_nullable_raw_parts(
*seq.data.get() as *const TVMFFIAny,
*seq.size.get() as usize,
)
diff --git a/rust/tvm-ffi/src/collections/map.rs
b/rust/tvm-ffi/src/collections/map.rs
index 94278fdb..2ba7c449 100644
--- a/rust/tvm-ffi/src/collections/map.rs
+++ b/rust/tvm-ffi/src/collections/map.rs
@@ -350,24 +350,47 @@ where
/// `try_cast_from_any_view`, which run during argument decoding where a
/// panic could unwind across the C ABI.
fn try_raw_entries(&self) -> Result<Vec<(Any, Any)>> {
- let mut entries = Vec::with_capacity(self.len());
- let mut remaining = self.len();
+ forward_entries(
+ crate::cached_global_func!("ffi.MapForwardIterFunctor"),
+ AnyView::from(self),
+ self.len(),
+ )
+ }
+
+ /// Converts raw entries one by one into a new map.
+ fn from_raw_entries(entries: Vec<(Any, Any)>) -> Result<Self> {
+ let mut pairs = Vec::with_capacity(entries.len());
+ for (k, v) in entries {
+ pairs.push((element_from_any::<K>(k)?, element_from_any::<V>(v)?));
+ }
+ Self::from_pairs(&pairs)
+ }
+}
+
+/// The `len` entries of the map or dict `container`, read with the iterator
+/// functor that `iter_functor` makes for it.
+fn forward_entries(
+ iter_functor: &Function,
+ container: AnyView,
+ len: usize,
+) -> Result<Vec<(Any, Any)>> {
+ let mut entries = Vec::with_capacity(len);
+ let mut remaining = len;
+ if remaining == 0 {
+ return Ok(entries);
+ }
+ let functor = iter_functor
+ .call_packed(&[container])
+ .and_then(Function::try_from)?;
+ loop {
+ let k = functor.call_packed(&[AnyView::from(&0i64)])?;
+ let v = functor.call_packed(&[AnyView::from(&1i64)])?;
+ entries.push((k, v));
+ remaining -= 1;
if remaining == 0 {
return Ok(entries);
}
- let functor = crate::cached_global_func!("ffi.MapForwardIterFunctor")
- .call_packed(&[AnyView::from(self)])
- .and_then(Function::try_from)?;
- loop {
- let k = functor.call_packed(&[AnyView::from(&0i64)])?;
- let v = functor.call_packed(&[AnyView::from(&1i64)])?;
- entries.push((k, v));
- remaining -= 1;
- if remaining == 0 {
- return Ok(entries);
- }
- functor.call_packed(&[AnyView::from(&2i64)])?;
- }
+ functor.call_packed(&[AnyView::from(&2i64)])?;
}
}
@@ -649,6 +672,18 @@ where
}
unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFIDict as i32 {
+ // A dict is mutable, so it converts by copy, as C++ `Map` accepts
it.
+ // `ffi.Dict` shares the `MapBaseObj` header.
+ let len = (*(data.data_union.v_obj as *const MapObj)).size as
usize;
+ let entries = forward_entries(
+ crate::cached_global_func!("ffi.DictForwardIterFunctor"),
+ AnyView::from_raw_ffi_any(*data),
+ len,
+ )
+ .map_err(|_| ())?;
+ return Self::from_raw_entries(entries).map_err(|_| ());
+ }
if data.type_index != TypeIndex::kTVMFFIMap as i32 {
return Err(());
}
@@ -663,13 +698,8 @@ where
// Slow path: try to convert entry by entry into a new map, as C++
// `TryCastFromAnyView` does.
let src = <Self as
AnyCompatible>::copy_from_any_view_after_check(data);
- let mut pairs = Vec::with_capacity(src.len());
- for (k, v) in src.try_raw_entries().map_err(|_| ())? {
- let k = element_from_any::<K>(k).map_err(|_| ())?;
- let v = element_from_any::<V>(v).map_err(|_| ())?;
- pairs.push((k, v));
- }
- Self::from_pairs(&pairs).map_err(|_| ())
+ let entries = src.try_raw_entries().map_err(|_| ())?;
+ Self::from_raw_entries(entries).map_err(|_| ())
}
}
diff --git a/rust/tvm-ffi/src/collections/shape.rs
b/rust/tvm-ffi/src/collections/shape.rs
index ef0eb3b0..f7214c0c 100644
--- a/rust/tvm-ffi/src/collections/shape.rs
+++ b/rust/tvm-ffi/src/collections/shape.rs
@@ -86,11 +86,9 @@ impl Shape {
/// The dimensions as a slice.
pub fn as_slice(&self) -> &[i64] {
- // An empty shape may have a null data pointer, which a slice cannot.
- if self.data.data.size == 0 {
- return &[];
+ unsafe {
+ crate::util::slice_from_nullable_raw_parts(self.data.data.data,
self.data.data.size)
}
- unsafe { std::slice::from_raw_parts(self.data.data.data,
self.data.data.size) }
}
/// Fill `strides` with the row-major strides of `shape`.
diff --git a/rust/tvm-ffi/src/collections/tensor.rs
b/rust/tvm-ffi/src/collections/tensor.rs
index 1dd1b2da..83b82dfd 100644
--- a/rust/tvm-ffi/src/collections/tensor.rs
+++ b/rust/tvm-ffi/src/collections/tensor.rs
@@ -68,6 +68,9 @@ pub unsafe trait NDAllocator: 'static {
/// Size queries on a `DLTensor`.
pub trait DLTensorExt {
/// The number of elements: the product of the shape, 1 for a scalar.
+ ///
+ /// # Panics
+ /// If an extent is negative or the product overflows `usize`.
fn numel(&self) -> usize;
/// The bytes of one element: `bits * lanes`, rounded up to whole bytes.
fn item_size(&self) -> usize;
@@ -76,14 +79,9 @@ pub trait DLTensorExt {
impl DLTensorExt for DLTensor {
fn numel(&self) -> usize {
// A scalar (`ndim == 0`) has one element and may have a null shape.
- if self.ndim <= 0 {
- return 1;
- }
- unsafe {
- std::slice::from_raw_parts(self.shape, self.ndim as usize)
- .iter()
- .product::<i64>() as usize
- }
+ let shape =
+ unsafe { crate::util::slice_from_nullable_raw_parts(self.shape,
self.ndim as usize) };
+ checked_numel(shape).expect("tensor shape has a negative extent or
overflows usize")
}
fn item_size(&self) -> usize {
@@ -91,6 +89,22 @@ impl DLTensorExt for DLTensor {
}
}
+/// The product of `shape`; `None` if an extent is negative or it overflows.
+fn checked_numel(shape: &[i64]) -> Option<usize> {
+ shape.iter().try_fold(1usize, |numel, &extent| {
+ numel.checked_mul(usize::try_from(extent).ok()?)
+ })
+}
+
+/// The bytes of `tensor`'s data; it fits in `isize`, as an allocation must.
+fn data_size(tensor: &DLTensor) -> usize {
+ tensor
+ .numel()
+ .checked_mul(tensor.item_size())
+ .filter(|&size| size <= isize::MAX as usize)
+ .expect("tensor byte size overflows isize")
+}
+
//-----------------------------------------------------
// Tensor
//-----------------------------------------------------
@@ -161,7 +175,7 @@ impl Tensor {
dtype.to_string()
);
}
- if self.device().device_type != DLDeviceType::kDLCPU {
+ if self.device().device_type != DLDeviceType::kDLCPU as i32 {
crate::bail!(crate::error::RUNTIME_ERROR, "Tensor is not on CPU");
}
crate::ensure!(
@@ -171,11 +185,7 @@ impl Tensor {
);
unsafe {
- // An empty tensor may have a null data pointer, which a slice
cannot.
- if self.numel() == 0 {
- return Ok(&[]);
- }
- Ok(std::slice::from_raw_parts(
+ Ok(crate::util::slice_from_nullable_raw_parts(
self.data.dltensor.data as *const T,
self.numel(),
))
@@ -208,7 +218,7 @@ impl Tensor {
self.dtype().to_string()
);
}
- if self.device().device_type != DLDeviceType::kDLCPU {
+ if self.device().device_type != DLDeviceType::kDLCPU as i32 {
crate::bail!(crate::error::RUNTIME_ERROR, "Tensor is not on CPU");
}
crate::ensure!(
@@ -217,10 +227,7 @@ impl Tensor {
"Tensor is not contiguous"
);
unsafe {
- if self.numel() == 0 {
- return Ok(&mut []);
- }
- Ok(std::slice::from_raw_parts_mut(
+ Ok(crate::util::slice_from_nullable_raw_parts_mut(
self.data.dltensor.data as *mut T,
self.numel(),
))
@@ -230,11 +237,7 @@ impl Tensor {
/// The extent of each dimension.
pub fn shape(&self) -> &[i64] {
let t = &self.data.dltensor;
- // A scalar may have a null shape, which a slice cannot.
- if t.ndim == 0 {
- return &[];
- }
- unsafe { std::slice::from_raw_parts(t.shape, t.ndim as usize) }
+ unsafe { crate::util::slice_from_nullable_raw_parts(t.shape, t.ndim as
usize) }
}
/// The number of dimensions.
@@ -243,6 +246,9 @@ impl Tensor {
}
/// The number of elements.
+ ///
+ /// # Panics
+ /// As [`DLTensorExt::numel`].
pub fn numel(&self) -> usize {
self.data.dltensor.numel()
}
@@ -320,6 +326,9 @@ impl Tensor {
///
/// # Returns
/// The new tensor, with row-major strides.
+ ///
+ /// # Panics
+ /// If an extent of `shape` is negative or the byte size overflows `isize`.
pub fn from_nd_alloc<TNDAlloc>(
alloc: TNDAlloc,
shape: &[i64],
@@ -329,17 +338,23 @@ impl Tensor {
where
TNDAlloc: NDAllocator,
{
+ let dltensor = DLTensor {
+ data: std::ptr::null_mut(),
+ device,
+ ndim: shape.len() as i32,
+ dtype,
+ shape: shape.as_ptr() as *mut i64,
+ strides: std::ptr::null_mut(),
+ byte_offset: 0,
+ };
+ // Allocators size the data from the shape, so check it before
allocating.
+ data_size(&dltensor);
let tensor_obj = TensorObjFromNDAlloc {
base_: TensorObj {
object: Object::new(),
dltensor: DLTensor {
- data: std::ptr::null_mut(),
- device,
- ndim: shape.len() as i32,
- dtype,
shape: std::ptr::null_mut(),
- strides: std::ptr::null_mut(),
- byte_offset: 0,
+ ..dltensor
},
},
alloc,
@@ -409,14 +424,15 @@ impl Tensor {
/// The new tensor, allocated with [`CPUNDAlloc`].
///
/// # Errors
- /// A `ValueError` if `slice` does not have as many elements as `shape`.
+ /// A `ValueError` if `slice` does not have as many elements as `shape`,
+ /// or `shape` has a negative extent.
pub fn from_slice<T: AsDLDataType>(slice: &[T], shape: &[i64]) ->
Result<Self> {
+ if checked_numel(shape) != Some(slice.len()) {
+ crate::bail!(crate::error::VALUE_ERROR, "Slice length mismatch");
+ }
let dtype = T::DL_DATA_TYPE;
let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
let tensor = Tensor::from_nd_alloc(CPUNDAlloc, shape, dtype, device);
- if tensor.numel() != slice.len() {
- crate::bail!(crate::error::VALUE_ERROR, "Slice length mismatch");
- }
// SAFETY: the tensor was just allocated, so nothing else can access
it.
unsafe { tensor.data_as_slice_mut::<T>()? }.copy_from_slice(slice);
Ok(tensor)
@@ -432,9 +448,7 @@ unsafe impl NDAllocator for CPUNDAlloc {
const MIN_ALIGN: usize = 64;
unsafe fn alloc_data(&mut self, prototype: &DLTensor) -> *mut
core::ffi::c_void {
- let numel = prototype.numel();
- let item_size = prototype.item_size();
- let size = numel * item_size;
+ let size = data_size(prototype);
// A zero-size allocation is undefined behavior; an empty tensor gets
an
// aligned dangling pointer, which `free_data` does not release.
if size == 0 {
@@ -449,9 +463,7 @@ unsafe impl NDAllocator for CPUNDAlloc {
}
unsafe fn free_data(&mut self, tensor: &DLTensor) {
- let numel = tensor.numel();
- let item_size = tensor.item_size();
- let size = numel * item_size;
+ let size = data_size(tensor);
if size == 0 {
return;
}
@@ -539,15 +551,15 @@ impl<'a> TensorView<'a> {
/// The extent of each dimension.
pub fn shape(&self) -> &'a [i64] {
let t = self.dltensor();
- if t.ndim == 0 {
- return &[];
- }
- unsafe { std::slice::from_raw_parts(t.shape, t.ndim as usize) }
+ unsafe { crate::util::slice_from_nullable_raw_parts(t.shape, t.ndim as
usize) }
}
/// The number of elements.
+ ///
+ /// # Panics
+ /// As [`DLTensorExt::numel`].
pub fn numel(&self) -> usize {
- self.shape().iter().product::<i64>() as usize
+ self.dltensor().numel()
}
/// The strides, in elements.
diff --git a/rust/tvm-ffi/src/device.rs b/rust/tvm-ffi/src/device.rs
index a7513d23..cd39363a 100644
--- a/rust/tvm-ffi/src/device.rs
+++ b/rust/tvm-ffi/src/device.rs
@@ -27,7 +27,7 @@ use tvm_ffi_sys::{TVMFFIEnvGetStream, TVMFFIEnvSetStream,
TVMFFIStreamHandle};
/// The current stream of `device`, as the caller or [`with_stream`] set it;
/// null if none is set.
pub fn current_stream(device: &DLDevice) -> TVMFFIStreamHandle {
- unsafe { TVMFFIEnvGetStream(device.device_type as i32, device.device_id) }
+ unsafe { TVMFFIEnvGetStream(device.device_type, device.device_id) }
}
/// Call `f` with the device stream temporarily set to `stream`.
@@ -46,7 +46,7 @@ pub unsafe fn with_stream<T>(
let mut prev_stream: TVMFFIStreamHandle = std::ptr::null_mut();
unsafe {
crate::check_safe_call!(TVMFFIEnvSetStream(
- device.device_type as i32,
+ device.device_type,
device.device_id,
stream,
&mut prev_stream as *mut TVMFFIStreamHandle
@@ -74,7 +74,7 @@ impl RestoreStream {
fn set(&self) -> i32 {
unsafe {
TVMFFIEnvSetStream(
- self.device.device_type as i32,
+ self.device.device_type,
self.device.device_id,
self.stream,
std::ptr::null_mut(),
@@ -99,8 +99,8 @@ impl Drop for RestoreStream {
unsafe impl AnyCompatible for DLDevice {
#[inline]
unsafe fn is_undefined(field: *const Self) -> bool {
- // No device type is zero; read as an integer, since zero is no
`DLDeviceType`.
- *std::ptr::addr_of!((*field).device_type).cast::<u32>() == 0
+ // No device type is zero.
+ *std::ptr::addr_of!((*field).device_type) == 0
}
fn type_str() -> String {
diff --git a/rust/tvm-ffi/src/dtype.rs b/rust/tvm-ffi/src/dtype.rs
index 8aa77ca3..d926fd6e 100644
--- a/rust/tvm-ffi/src/dtype.rs
+++ b/rust/tvm-ffi/src/dtype.rs
@@ -122,14 +122,19 @@ unsafe impl AnyCompatible for DLDataType {
}
/// A Rust element type of a tensor, such as `f32`, and its `DLDataType`.
-pub trait AsDLDataType: Copy {
+///
+/// # Safety
+/// Tensor data of `DL_DATA_TYPE` is read as `Self` without checks, so `Self`
+/// must have that data type's item size, be no more aligned than the item
+/// size, and accept every bit pattern of it.
+pub unsafe trait AsDLDataType: Copy {
/// The data type of `Self`.
const DL_DATA_TYPE: DLDataType;
}
macro_rules! impl_as_dl_data_type {
($type: ty, $code: expr, $bits: expr) => {
- impl AsDLDataType for $type {
+ unsafe impl AsDLDataType for $type {
const DL_DATA_TYPE: DLDataType = DLDataType {
code: $code as u8,
bits: $bits as u8,
diff --git a/rust/tvm-ffi/src/error.rs b/rust/tvm-ffi/src/error.rs
index 48f6cf58..2cf64ae5 100644
--- a/rust/tvm-ffi/src/error.rs
+++ b/rust/tvm-ffi/src/error.rs
@@ -281,7 +281,8 @@ impl Error {
///
/// The kind of the error, such as [`VALUE_ERROR`].
pub fn kind(&self) -> ErrorKind<'_> {
- ErrorKind(self.data.cell.kind.as_str())
+ // SAFETY: the error cell holds UTF-8 text that lives as long as the
error.
+ ErrorKind(unsafe { self.data.cell.kind.as_str() })
}
/// Get the message of the error.
@@ -290,7 +291,8 @@ impl Error {
///
/// The message of the error.
pub fn message(&self) -> &str {
- self.data.cell.message.as_str()
+ // SAFETY: the error cell holds UTF-8 text that lives as long as the
error.
+ unsafe { self.data.cell.message.as_str() }
}
/// Get the backtrace of the error.
@@ -299,7 +301,8 @@ impl Error {
///
/// The backtrace, most recent call first.
pub fn backtrace(&self) -> &str {
- self.data.cell.backtrace.as_str()
+ // SAFETY: the error cell holds UTF-8 text that lives as long as the
error.
+ unsafe { self.data.cell.backtrace.as_str() }
}
/// The error that caused this one, the next link of the cause chain; the
diff --git a/rust/tvm-ffi/src/extra/structural_common.rs
b/rust/tvm-ffi/src/extra/structural_common.rs
index 441fe017..79afcae5 100644
--- a/rust/tvm-ffi/src/extra/structural_common.rs
+++ b/rust/tvm-ffi/src/extra/structural_common.rs
@@ -225,10 +225,11 @@ unsafe fn visit_field_level<B>(
info: *const TVMFFITypeInfo,
callback: &mut impl FnMut(&'static TVMFFIFieldInfo) -> ControlFlow<B>,
) -> Option<B> {
- if info.is_null() || (*info).fields.is_null() {
+ if info.is_null() {
return None;
}
- let fields = std::slice::from_raw_parts((*info).fields, (*info).num_fields
as usize);
+ let fields =
+ crate::util::slice_from_nullable_raw_parts((*info).fields,
(*info).num_fields as usize);
for field in fields {
// C reflection tables are immortal once registered.
let field: &'static TVMFFIFieldInfo = &*(field as *const
TVMFFIFieldInfo);
diff --git a/rust/tvm-ffi/src/extra/structural_mutate.rs
b/rust/tvm-ffi/src/extra/structural_mutate.rs
index 2d205453..ec207551 100644
--- a/rust/tvm-ffi/src/extra/structural_mutate.rs
+++ b/rust/tvm-ffi/src/extra/structural_mutate.rs
@@ -1583,7 +1583,8 @@ fn call_field_setter(
value: &TVMFFIAny,
) -> Result<()> {
if field.setter.is_null() {
- return Err(missing_field_setter_error(field.name.as_str()));
+ // SAFETY: registered field names are UTF-8 text that lives as long as
the process.
+ return Err(missing_field_setter_error(unsafe { field.name.as_str() }));
}
let return_code = unsafe {
if field.flags & FLAG_SETTER_IS_FUNCTION == 0 {
diff --git a/rust/tvm-ffi/src/function.rs b/rust/tvm-ffi/src/function.rs
index b8dd4e0b..009f58a6 100644
--- a/rust/tvm-ffi/src/function.rs
+++ b/rust/tvm-ffi/src/function.rs
@@ -116,7 +116,7 @@ impl<F: Fn(&[AnyView]) -> Result<Any> + Send + Sync +
'static> CallbackFunctionO
result: *mut TVMFFIAny,
) -> i32 {
let this = &*(handle as *mut Self);
- let packed_args = std::slice::from_raw_parts(args as *const AnyView,
num_args as usize);
+ let packed_args = crate::function_internal::packed_args(args,
num_args);
let ret_value = match
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
(this.callback)(packed_args)
})) {
diff --git a/rust/tvm-ffi/src/function_internal.rs
b/rust/tvm-ffi/src/function_internal.rs
index c992d7ea..f9de7fd6 100644
--- a/rust/tvm-ffi/src/function_internal.rs
+++ b/rust/tvm-ffi/src/function_internal.rs
@@ -35,6 +35,17 @@ pub fn panic_to_error(payload: Box<dyn std::any::Any +
Send>) -> Error {
Error::from_panic(payload)
}
+/// The arguments of a packed call from C.
+///
+/// # Safety
+/// If `num_args` is positive, `args` must point to `num_args` valid
+/// arguments that outlive `'a`.
+#[doc(hidden)]
+#[inline(always)]
+pub unsafe fn packed_args<'a>(args: *const TVMFFIAny, num_args: i32) -> &'a
[AnyView<'a>] {
+ crate::util::slice_from_nullable_raw_parts(args as *const AnyView,
num_args as usize)
+}
+
pub trait AsPackedCallable<I, O> {
// Call the function in packed convention
fn call_packed(&self, packed_args: &[AnyView]) -> Result<Any>;
diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs
index 7012023d..2e027dd8 100644
--- a/rust/tvm-ffi/src/lib.rs
+++ b/rust/tvm-ffi/src/lib.rs
@@ -66,6 +66,7 @@ pub mod rvalue_ref;
pub mod static_init;
pub mod string;
pub mod type_traits;
+mod util;
pub use tvm_ffi_sys;
pub use crate::any::{Any, AnyView};
diff --git a/rust/tvm-ffi/src/macros.rs b/rust/tvm-ffi/src/macros.rs
index 4b59ecb4..a37ece0d 100644
--- a/rust/tvm-ffi/src/macros.rs
+++ b/rust/tvm-ffi/src/macros.rs
@@ -641,8 +641,7 @@ macro_rules! tvm_ffi_dll_export_typed_func {
num_args: i32,
result: *mut $crate::tvm_ffi_sys::TVMFFIAny,
) -> i32 {
- let packed_args =
- std::slice::from_raw_parts(args as *const
$crate::any::AnyView, num_args as usize);
+ let packed_args = $crate::function_internal::packed_args(args,
num_args);
let ret_value = match
::std::panic::catch_unwind(::std::panic::AssertUnwindSafe(|| {
$crate::function_internal::call_packed_callable($func,
packed_args)
})) {
diff --git a/rust/tvm-ffi/src/reflection/accessor.rs
b/rust/tvm-ffi/src/reflection/accessor.rs
index 2b1f111b..40d10a7d 100644
--- a/rust/tvm-ffi/src/reflection/accessor.rs
+++ b/rust/tvm-ffi/src/reflection/accessor.rs
@@ -101,7 +101,8 @@ pub fn type_key(type_index: i32) -> Option<&'static str> {
// SAFETY: the index is registered, and type information lives as long as
// the process.
let info: &'static TVMFFITypeInfo = unsafe {
&*TVMFFIGetTypeInfo(type_index) };
- Some(info.type_key.as_str())
+ // SAFETY: a type key is UTF-8 text.
+ Some(unsafe { info.type_key.as_str() })
}
/// The default value `T` registers for its field `name`, or `None` if `T`
@@ -113,12 +114,14 @@ pub fn field_default<T: ObjectCore>(name: &str) ->
Option<Any> {
// SAFETY: `ObjectCore` types are registered, and type information lives
// as long as the process.
let info: &'static TVMFFITypeInfo = unsafe {
&*TVMFFIGetTypeInfo(T::type_index()) };
- if info.fields.is_null() {
- return None;
- }
// SAFETY: a registered type lists `num_fields` initialized fields.
- let fields = unsafe { std::slice::from_raw_parts(info.fields,
info.num_fields as usize) };
- let field = fields.iter().find(|field| field.name.as_str() == name)?;
+ let fields = unsafe {
+ crate::util::slice_from_nullable_raw_parts(info.fields,
info.num_fields as usize)
+ };
+ // SAFETY: a field name is UTF-8 text.
+ let field = fields
+ .iter()
+ .find(|field| unsafe { field.name.as_str() } == name)?;
let has_default =
TVMFFIFieldFlagBitMask::kTVMFFIFieldFlagBitMaskHasDefault as i64;
let from_factory =
TVMFFIFieldFlagBitMask::kTVMFFIFieldFlagBitMaskDefaultFromFactory as i64;
if field.flags & has_default == 0 || field.flags & from_factory != 0 {
@@ -285,9 +288,11 @@ impl FieldGetter {
})?;
let field = unsafe { field.as_ref() };
let getter = field.getter.ok_or_else(|| {
+ // SAFETY: a field name is UTF-8 text.
+ let name = unsafe { field.name.as_str() };
Error::new(
TYPE_ERROR,
- &format!("Reflected field `{}` has no getter",
field.name.as_str()),
+ &format!("Reflected field `{name}` has no getter"),
"",
)
})?;
diff --git a/rust/tvm-ffi/src/reflection/unsafe_init.rs
b/rust/tvm-ffi/src/reflection/unsafe_init.rs
index ba4493c7..31b3fce4 100644
--- a/rust/tvm-ffi/src/reflection/unsafe_init.rs
+++ b/rust/tvm-ffi/src/reflection/unsafe_init.rs
@@ -188,7 +188,7 @@ mod tests {
assert!(!zero_is_undefined::<i64>());
assert!(!zero_is_undefined::<bool>());
- // `None`, through the niche of a reference or a device, or through a
tag.
+ // `None`, through the niche of a reference or through a tag.
assert!(!zero_is_undefined::<Option<Array<i64>>>());
assert!(!zero_is_undefined::<Option<DLDevice>>());
assert!(!zero_is_undefined::<Option<String>>());
@@ -197,7 +197,7 @@ mod tests {
#[test]
fn values_are_defined() {
let cpu = DLDevice {
- device_type: DLDeviceType::kDLCPU,
+ device_type: DLDeviceType::kDLCPU as i32,
device_id: 0,
};
assert!(!is_undefined(Array::<i64>::new()));
diff --git a/rust/tvm-ffi/src/string.rs b/rust/tvm-ffi/src/string.rs
index 247f1649..c91b33a3 100644
--- a/rust/tvm-ffi/src/string.rs
+++ b/rust/tvm-ffi/src/string.rs
@@ -524,7 +524,7 @@ unsafe impl AnyCompatible for Bytes {
if data.type_index == TypeIndex::kTVMFFIByteArrayPtr as i32 {
// deep copy from bytearray ptr
let bytes = &*(data.data_union.v_ptr as *const TVMFFIByteArray);
- Ok(Self::from(std::slice::from_raw_parts(
+ Ok(Self::from(crate::util::slice_from_nullable_raw_parts(
bytes.data, bytes.size,
)))
} else if data.type_index == TypeIndex::kTVMFFISmallBytes as i32 {
@@ -576,10 +576,9 @@ unsafe impl AnyCompatible for String {
unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
if data.type_index == TypeIndex::kTVMFFIRawStr as i32 {
- // 1. Create a CStr wrapper from the raw pointer.
let c_str =
std::ffi::CStr::from_ptr(data.data_union.v_c_str as *const
std::os::raw::c_char);
- Ok(Self::from(c_str.to_str().expect("Invalid UTF-8")))
+ c_str.to_str().map(Self::from).map_err(|_| ())
} else if data.type_index == TypeIndex::kTVMFFISmallStr as i32 {
Ok(Self { data: *data })
} else if data.type_index == TypeIndex::kTVMFFIStr as i32 {
diff --git a/rust/tvm-ffi/src/type_traits.rs b/rust/tvm-ffi/src/type_traits.rs
index 4495a103..f6aba47c 100644
--- a/rust/tvm-ffi/src/type_traits.rs
+++ b/rust/tvm-ffi/src/type_traits.rs
@@ -785,8 +785,8 @@ unsafe impl<T: AnyCompatible> AnyCompatible for Option<T> {
unsafe fn is_undefined(field: *const Self) -> bool {
// With a tag, all-zero bytes are `None`. A niche-optimized `Option<T>`
// keeps `T`'s bytes, where all-zero bytes are `None` if `None` looks
- // undefined to `T` (a reference's null pointer, a device's zero type),
- // and undefined if `None` sits elsewhere in the niche.
+ // undefined to `T` (a reference's null pointer), and undefined if
+ // `None` sits elsewhere in the niche.
if std::mem::size_of::<Self>() != std::mem::size_of::<T>() {
return false;
}
diff --git a/rust/tvm-ffi/src/util.rs b/rust/tvm-ffi/src/util.rs
new file mode 100644
index 00000000..515e2b20
--- /dev/null
+++ b/rust/tvm-ffi/src/util.rs
@@ -0,0 +1,47 @@
+/*
+ * 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.
+ */
+
+//! Crate-internal helpers.
+
+/// [`std::slice::from_raw_parts`] that allows a null `data` when `len` is 0.
+///
+/// # Safety
+/// As for [`std::slice::from_raw_parts`] when `len` is not 0.
+#[inline(always)]
+pub(crate) unsafe fn slice_from_nullable_raw_parts<'a, T>(data: *const T, len:
usize) -> &'a [T] {
+ if len == 0 {
+ return &[];
+ }
+ std::slice::from_raw_parts(data, len)
+}
+
+/// [`std::slice::from_raw_parts_mut`] that allows a null `data` when `len` is
0.
+///
+/// # Safety
+/// As for [`std::slice::from_raw_parts_mut`] when `len` is not 0.
+#[inline(always)]
+pub(crate) unsafe fn slice_from_nullable_raw_parts_mut<'a, T>(
+ data: *mut T,
+ len: usize,
+) -> &'a mut [T] {
+ if len == 0 {
+ return &mut [];
+ }
+ std::slice::from_raw_parts_mut(data, len)
+}
diff --git a/rust/tvm-ffi/tests/test_any.rs b/rust/tvm-ffi/tests/test_any.rs
index fa91998b..c73bed03 100644
--- a/rust/tvm-ffi/tests/test_any.rs
+++ b/rust/tvm-ffi/tests/test_any.rs
@@ -326,13 +326,13 @@ fn test_any_dl_device() {
// Test Any with CPU device
assert_eq!(any_cpu.type_index(), TypeIndex::kTVMFFIDevice as i32);
let converted_cpu = DLDevice::try_from(any_cpu).unwrap();
- assert_eq!(converted_cpu.device_type, DLDeviceType::kDLCPU);
+ assert_eq!(converted_cpu.device_type, DLDeviceType::kDLCPU as i32);
assert_eq!(converted_cpu.device_id, 0);
// Test AnyView with CPU device
assert_eq!(any_view_cpu.type_index(), TypeIndex::kTVMFFIDevice as i32);
let converted_cpu_view = DLDevice::try_from(any_view_cpu).unwrap();
- assert_eq!(converted_cpu_view.device_type, DLDeviceType::kDLCPU);
+ assert_eq!(converted_cpu_view.device_type, DLDeviceType::kDLCPU as i32);
assert_eq!(converted_cpu_view.device_id, 0);
// Test DLDevice with CUDA device
@@ -343,13 +343,16 @@ fn test_any_dl_device() {
// Test Any with CUDA device
assert_eq!(any_cuda.type_index(), TypeIndex::kTVMFFIDevice as i32);
let converted_cuda = DLDevice::try_from(any_cuda).unwrap();
- assert_eq!(converted_cuda.device_type, DLDeviceType::kDLCUDA);
+ assert_eq!(converted_cuda.device_type, DLDeviceType::kDLCUDA as i32);
assert_eq!(converted_cuda.device_id, 1);
// Test AnyView with CUDA device
assert_eq!(any_view_cuda.type_index(), TypeIndex::kTVMFFIDevice as i32);
let converted_cuda_view = DLDevice::try_from(any_view_cuda).unwrap();
- assert_eq!(converted_cuda_view.device_type, DLDeviceType::kDLCUDA);
+ assert_eq!(
+ converted_cuda_view.device_type,
+ DLDeviceType::kDLCUDA as i32
+ );
assert_eq!(converted_cuda_view.device_id, 1);
}
diff --git a/rust/tvm-ffi/tests/test_map.rs b/rust/tvm-ffi/tests/test_map.rs
index 950c8f64..dcc37344 100644
--- a/rust/tvm-ffi/tests/test_map.rs
+++ b/rust/tvm-ffi/tests/test_map.rs
@@ -382,3 +382,15 @@ fn test_map_any_cast_slow_path_converts() {
assert_eq!(AnyView::from(&real).debug_strong_count(), Some(1));
assert_eq!(AnyView::from(&converted).debug_strong_count(), Some(1));
}
+
+#[test]
+fn test_map_from_dict_copies() {
+ let dict = Function::get_global("ffi.Dict")
+ .unwrap()
+ .call_tuple((String::from("a"), 1i64, String::from("b"), 2i64))
+ .unwrap();
+ let map = Map::<String, i64>::try_from(dict.clone()).unwrap();
+ assert_eq!(map.len(), 2);
+ assert_eq!(map.get(&String::from("b")).unwrap(), Some(2));
+ assert_eq!(dict.debug_strong_count(), Some(1));
+}
diff --git a/rust/tvm-ffi/tests/test_object_reflection.rs
b/rust/tvm-ffi/tests/test_object_reflection.rs
index a7dbf137..18990179 100644
--- a/rust/tvm-ffi/tests/test_object_reflection.rs
+++ b/rust/tvm-ffi/tests/test_object_reflection.rs
@@ -177,7 +177,10 @@ fn register_types() {
fn test_reflected_fields_are_registered_with_the_type() {
register_types();
let info = type_info::<PointObj>();
- let names: Vec<&str> = fields(info).iter().map(|f|
f.name.as_str()).collect();
+ let names: Vec<&str> = fields(info)
+ .iter()
+ .map(|f| unsafe { f.name.as_str() })
+ .collect();
assert_eq!(names, ["x", "y", "label"]);
let [x, y, label] = fields(info) else {
@@ -187,7 +190,7 @@ fn test_reflected_fields_are_registered_with_the_type() {
assert_eq!(x.offset, std::mem::offset_of!(PointObj, x) as i64);
assert_eq!(x.field_static_type_index, TypeIndex::kTVMFFIInt as i32);
assert_eq!(
- x.metadata.as_str(),
+ unsafe { x.metadata.as_str() },
r#"{"type_schema":"{\"type\":\"int\"}"}"#
);
assert_eq!(
@@ -212,10 +215,13 @@ fn test_reflected_fields_are_registered_with_the_type() {
// The subtype registers only its own fields, and inherits the kind.
let info = type_info::<Point3Obj>();
- let names: Vec<&str> = fields(info).iter().map(|f|
f.name.as_str()).collect();
+ let names: Vec<&str> = fields(info)
+ .iter()
+ .map(|f| unsafe { f.name.as_str() })
+ .collect();
assert_eq!(names, ["z", "labels", "type"]);
assert_eq!(
- fields(info)[1].metadata.as_str(),
+ unsafe { fields(info)[1].metadata.as_str() },
r#"{"type_schema":"{\"type\":\"ffi.Array\",\"args\":[{\"type\":\"ffi.String\"}]}"}"#
);
assert_eq!(
@@ -230,12 +236,12 @@ fn test_field_docs_come_from_doc_comments_or_reflect() {
let [x, y, label] = fields(type_info::<PointObj>()) else {
unreachable!()
};
- assert_eq!(x.doc.as_str(), "The x coordinate.\n\nSigned.");
- assert_eq!(y.doc.as_str(), "The y coordinate.");
- assert_eq!(label.doc.as_str(), "");
+ assert_eq!(unsafe { x.doc.as_str() }, "The x coordinate.\n\nSigned.");
+ assert_eq!(unsafe { y.doc.as_str() }, "The y coordinate.");
+ assert_eq!(unsafe { label.doc.as_str() }, "");
// A renamed field keeps its Rust type and offset.
let labels = &fields(type_info::<Point3Obj>())[1];
- assert_eq!(labels.name.as_str(), "labels");
+ assert_eq!(unsafe { labels.name.as_str() }, "labels");
assert_eq!(labels.offset, std::mem::offset_of!(Point3Obj, tags) as i64);
let p3 = point3(1, 2);
let getter = FieldGetter::new(Point3Obj::type_index(), "labels").unwrap();
@@ -247,29 +253,32 @@ fn test_field_docs_come_from_doc_comments_or_reflect() {
fn test_methods_are_registered_with_the_type() {
register_types();
let info = type_info::<PointObj>();
- let names: Vec<&str> = methods(info).iter().map(|m|
m.name.as_str()).collect();
+ let names: Vec<&str> = methods(info)
+ .iter()
+ .map(|m| unsafe { m.name.as_str() })
+ .collect();
assert_eq!(names, ["norm1", "shifted", "origin"]);
let [norm1, shifted, origin] = methods(info) else {
unreachable!()
};
assert_eq!(norm1.flags, 0);
- assert_eq!(norm1.doc.as_str(), "The L1 norm.");
- assert_eq!(shifted.doc.as_str(), "");
+ assert_eq!(unsafe { norm1.doc.as_str() }, "The L1 norm.");
+ assert_eq!(unsafe { shifted.doc.as_str() }, "");
assert_eq!(
origin.flags,
TVMFFIFieldFlagBitMask::kTVMFFIFieldFlagBitMaskIsStaticMethod as i64
);
- assert_eq!(origin.doc.as_str(), "The origin.");
+ assert_eq!(unsafe { origin.doc.as_str() }, "The origin.");
// Each method records its type schema, as C++ `ObjectDef::def` does.
let point = r#"{"type":"testing.rust.ReflectedPoint"}"#;
assert_eq!(
- shifted.metadata.as_str(),
+ unsafe { shifted.metadata.as_str() },
metadata_of(&format!(
r#"{{"type":"ffi.Function","named_args":{{"return":[{point}],"params":[{point},{{"type":"int"}}]}}}}"#
))
);
assert_eq!(
- origin.metadata.as_str(),
+ unsafe { origin.metadata.as_str() },
metadata_of(&format!(
r#"{{"type":"ffi.Function","named_args":{{"return":[{point}],"params":[]}}}}"#
))
diff --git a/rust/tvm-ffi/tests/test_rs_object.rs
b/rust/tvm-ffi/tests/test_rs_object.rs
index 848fef67..6d4ad954 100644
--- a/rust/tvm-ffi/tests/test_rs_object.rs
+++ b/rust/tvm-ffi/tests/test_rs_object.rs
@@ -273,7 +273,7 @@ fn type_info<N: ObjectCore>() -> &'static TVMFFITypeInfo {
fn field_names<N: ObjectCore>() -> Vec<&'static str> {
fields(type_info::<N>())
.iter()
- .map(|f| f.name.as_str())
+ .map(|f| unsafe { f.name.as_str() })
.collect()
}
@@ -333,12 +333,12 @@ fn test_fields_are_registered_with_their_flags() {
};
assert_eq!(x.flags, FieldFlags::NONE.bits());
assert_eq!(x.offset, std::mem::offset_of!(RsPointObj, x) as i64);
- assert_eq!(x.doc.as_str(), "The x coordinate.");
+ assert_eq!(unsafe { x.doc.as_str() }, "The x coordinate.");
assert_eq!(
label.flags,
(FieldFlags::WRITABLE | FieldFlags::SEQ_HASH_IGNORE).bits()
);
- assert_eq!(label.doc.as_str(), "A label.");
+ assert_eq!(unsafe { label.doc.as_str() }, "A label.");
let metadata = unsafe { &*type_info::<RsPointObj>().metadata };
assert_eq!(
metadata.total_size as usize,
@@ -370,8 +370,11 @@ fn test_reflection_registers_methods() {
let info = type_info::<RsPointObj>();
assert_eq!(info.num_methods, 1);
let method = unsafe { &*info.methods };
- assert_eq!(method.name.as_str(), "with_label");
- assert_eq!(method.doc.as_str(), "The point with another label.");
+ assert_eq!(unsafe { method.name.as_str() }, "with_label");
+ assert_eq!(
+ unsafe { method.doc.as_str() },
+ "The point with another label."
+ );
let with_label = tvm_ffi::reflection::get_method("testing.rust.RsPoint",
"with_label").unwrap();
let relabeled = with_label
.call_tuple((point(4, "a"), tvm_ffi::String::from("b")))
diff --git a/rust/tvm-ffi/tests/test_tensor.rs
b/rust/tvm-ffi/tests/test_tensor.rs
index 7cb6cd08..85704d5b 100644
--- a/rust/tvm-ffi/tests/test_tensor.rs
+++ b/rust/tvm-ffi/tests/test_tensor.rs
@@ -35,7 +35,7 @@ fn test_tensor_basic() {
assert_eq!(tensor.ndim(), 3);
assert_eq!(tensor.dtype().code, DLDataTypeCode::kDLFloat as u8);
assert_eq!(tensor.dtype().bits, 32_u8);
- assert_eq!(tensor.device().device_type, DLDeviceType::kDLCPU);
+ assert_eq!(tensor.device().device_type, DLDeviceType::kDLCPU as i32);
// Test strides (should be calculated correctly for row-major layout)
let strides = tensor.strides();
