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 28d09fe7 [FEAT][RUST]Add tvm_ffi::Map in Rust (#623)
28d09fe7 is described below
commit 28d09fe733a24fdf0b84249131dced8eee3e3574
Author: Linzhang Li <[email protected]>
AuthorDate: Fri Jul 3 18:36:01 2026 -0400
[FEAT][RUST]Add tvm_ffi::Map in Rust (#623)
This PR adds `tvm_ffi::Map` in Rust.
`tvm_ffi::MapObj` in Rust and the `ffi::MapObj` in C++ share the same
memory layer at run time(when `TVM_FFI_DEBUG_WITH_ABI_CHANGE` is
`False`). and `tvm_ffi::MapObj` will call the corresponding method of
`ffi::Obj` in C++ through global function call directly in most cases.
---------
Signed-off-by: yuchuan <[email protected]>
---
rust/tvm-ffi/src/collections/map.rs | 527 ++++++++++++++++++++++++++++++++++
rust/tvm-ffi/src/collections/mod.rs | 1 +
rust/tvm-ffi/src/extra/module.rs | 8 +-
rust/tvm-ffi/src/function_internal.rs | 27 ++
rust/tvm-ffi/src/lib.rs | 1 +
rust/tvm-ffi/src/macros.rs | 19 ++
rust/tvm-ffi/tests/test_map.rs | 288 +++++++++++++++++++
7 files changed, 865 insertions(+), 6 deletions(-)
diff --git a/rust/tvm-ffi/src/collections/map.rs
b/rust/tvm-ffi/src/collections/map.rs
new file mode 100644
index 00000000..cdb87adc
--- /dev/null
+++ b/rust/tvm-ffi/src/collections/map.rs
@@ -0,0 +1,527 @@
+/*
+ * 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.
+ */
+//! Immutable [`Map`] container backed by the C++ `ffi.Map` object.
+//!
+//! The `MapObj` header mirrors C++ `MapBaseObj`, so [`Map::len`] reads `size`
+//! directly (no FFI). The hash-table storage itself is an implementation
detail,
+//! so lookups and iteration are delegated to the global functions the C++
runtime
+//! registers (`ffi.MapGetItem`, `ffi.MapForwardIterFunctor`, ...): there is no
+//! Map-specific C ABI, and a Rust re-implementation would have to replicate
the
+//! hashing (`AnyHash`/`AnyEqual`) and dense/small probing, just as the Python
+//! bindings also delegate.
+//!
+//! `Any` conversions follow C++ `MapTypeTraitsBase` (map_base.h): the strict
+//! check walks every entry, and a failed strict `try_from` falls back to
+//! converting the entries one by one into a new map.
+use std::fmt::Debug;
+use std::marker::PhantomData;
+use std::ops::Deref;
+
+use crate::any::TryFromTemp;
+use crate::derive::Object;
+use crate::function::Function;
+use crate::object::{Object, ObjectArc};
+use crate::{Any, AnyCompatible, AnyView, Error, ObjectRefCore, Result};
+use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+use tvm_ffi_sys::{TVMFFIAny, TVMFFIObject};
+
+/// Container object for [`Map`]. The header fields mirror C++ `MapBaseObj`
+/// (`include/tvm/ffi/container/map_base.h`) so [`Map::len`] can read `size`
+/// without an FFI call; the storage `data` points to stays opaque.
+///
+/// This is a partial mirror by design: `MapBaseObj` has one more trailing
field
+/// after `slots_` (a `data_deleter_` function pointer), omitted here because
this
+/// binding is read-only — it only reads fields up to `size`, which precede it,
+/// and never allocates a `MapObj` itself (the C++ runtime allocates the larger
+/// `Small`/`DenseMapBaseObj` subclass), so the trailing field is never
touched.
+/// The layout also assumes `TVM_FFI_DEBUG_WITH_ABI_CHANGE` is off (the
default):
+/// with that debug flag set, `MapBaseObj` gains a *leading* `state_marker`
field
+/// that would shift every offset below.
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "ffi.Map"]
+#[type_index(TypeIndex::kTVMFFIMap)]
+pub struct MapObj {
+ pub object: Object,
+ /// Pointer to the (opaque) key/value storage region (`MapBaseObj::data_`).
+ pub data: *mut core::ffi::c_void,
+ /// Number of entries (`MapBaseObj::size_`).
+ pub size: u64,
+ /// Number of hash slots; the MSB is a small-map tag
(`MapBaseObj::slots_`).
+ pub slots: u64,
+}
+
+/// Immutable, reference-counted map from `K` to `V`, sharing its underlying
+/// `MapObj` with C++. Cloning is cheap (it bumps the refcount).
+#[repr(C)]
+pub struct Map<K, V> {
+ data: ObjectArc<MapObj>,
+ _marker: PhantomData<(K, V)>,
+}
+
+// Manual `Clone`: the data is the shared `ObjectArc`, so `Map<K, V>` is
`Clone`
+// regardless of whether `K`/`V` are (they are only phantom markers).
+impl<K, V> Clone for Map<K, V> {
+ fn clone(&self) -> Self {
+ Self {
+ data: self.data.clone(),
+ _marker: PhantomData,
+ }
+ }
+}
+
+unsafe impl<K, V> ObjectRefCore for Map<K, V> {
+ type ContainerType = MapObj;
+
+ fn data(this: &Self) -> &ObjectArc<MapObj> {
+ &this.data
+ }
+
+ fn into_data(this: Self) -> ObjectArc<MapObj> {
+ this.data
+ }
+
+ fn from_data(data: ObjectArc<MapObj>) -> Self {
+ Self {
+ data,
+ _marker: PhantomData,
+ }
+ }
+}
+
+// A `Map<K, V>` is a counted handle to its `MapObj`, so it derefs to it (like
+// `ObjectArc` does): methods can read header fields as `self.size` instead of
+// `self.data.size`.
+impl<K, V> Deref for Map<K, V> {
+ type Target = MapObj;
+ #[inline]
+ fn deref(&self) -> &MapObj {
+ &self.data
+ }
+}
+
+impl<K, V> Map<K, V>
+where
+ K: AnyCompatible,
+ V: AnyCompatible,
+{
+ /// Creates a new, empty map (via an `ffi.Map()` call to the C++ runtime).
+ pub fn new() -> Self {
+ Self::from_pairs(&[]).expect("ffi.Map() failed to construct an empty
map")
+ }
+
+ /// Builds a map by calling the C++ `ffi.Map` constructor with a flattened
+ /// `[k0, v0, k1, v1, ...]` argument list.
+ fn from_pairs(pairs: &[(K, V)]) -> Result<Self> {
+ let mut args: Vec<AnyView<'_>> = Vec::with_capacity(pairs.len() * 2);
+ for (k, v) in pairs {
+ args.push(AnyView::from(k));
+ args.push(AnyView::from(v));
+ }
+ let result = crate::cached_global_func!("ffi.Map").call_packed(&args)?;
+ Self::try_from(result)
+ }
+
+ /// Returns the number of entries in the map by reading the `MapObj` header
+ /// directly (no FFI call), like [`Array::len`](crate::Array).
+ pub fn len(&self) -> usize {
+ self.size as usize
+ }
+
+ /// Returns `true` if the map contains no entries.
+ pub fn is_empty(&self) -> bool {
+ self.len() == 0
+ }
+
+ /// Returns whether `key` is present, propagating any FFI error (e.g. a key
+ /// the C++ runtime cannot hash) rather than panicking. Backs both
+ /// [`Map::contains_key`] and [`Map::get`].
+ fn try_contains_key(&self, key: &K) -> Result<bool> {
+ let result = crate::cached_global_func!("ffi.MapCount")
+ .call_packed(&[AnyView::from(self), AnyView::from(key)])?;
+ Ok(i64::try_from(result)? != 0)
+ }
+
+ /// In debug builds, panics if `K` does not match the map's actual key
type,
+ /// checked against one stored key. Compiles to nothing in release, so
+ /// [`Map::get`] / [`Map::contains_key`] keep their single-FFI-call fast
path.
+ ///
+ /// A pure key lookup *hashes* the key rather than retrieving one, so a
wrong
+ /// `K` simply fails to match and reads as "absent" — unlike [`Map::iter`],
+ /// which retrieves keys and surfaces the mismatch. This assertion catches
+ /// that misuse in dev/tests without taxing release lookups. Only
meaningful
+ /// on a miss (a hit already proves `K` matched a stored key).
+ #[inline]
+ fn debug_assert_key_type(&self) {
+ #[cfg(debug_assertions)]
+ {
+ if !self.is_empty() {
+ let functor = self.iter_functor();
+ let first_key = functor
+ .call_packed(&[AnyView::from(&0i64)])
+ .expect("map iterator: reading current key failed");
+ assert!(
+ first_key.try_as::<K>().is_some(),
+ "Map lookup: key type `{}` does not match the map's stored
key type",
+ std::any::type_name::<K>(),
+ );
+ }
+ }
+ }
+
+ /// Returns `true` if `key` is present in the map.
+ ///
+ /// Panics if the lookup fails in the C++ runtime (e.g. `key` is not
+ /// hashable); use [`Map::get`] when such failures must be recovered from.
+ /// Passing a `key` whose *type* does not match the map's key type is
+ /// **undefined** (currently reads as `false`; see [`Map::get`]); debug
builds
+ /// assert against it.
+ pub fn contains_key(&self, key: &K) -> bool {
+ let present = self
+ .try_contains_key(key)
+ .expect("ffi.MapCount call failed");
+ if !present {
+ self.debug_assert_key_type();
+ }
+ present
+ }
+
+ /// Looks up `key`, returning `Ok(None)` if absent, or `Err` if the lookup
+ /// fails in the C++ runtime (e.g. `key` is not hashable) or the stored
value
+ /// cannot be converted to `V`. The map is immutable, so the lookup
+ /// (existence check then `ffi.MapGetItem`) cannot race with itself.
+ ///
+ /// Passing a `key` whose *type* does not match the map's key type is
+ /// **undefined**: the result is unspecified — it currently reads as absent
+ /// (`Ok(None)`) because a pure lookup hashes the key rather than
retrieving
+ /// one, so the mismatch cannot be surfaced as an `Err` the way a value
+ /// mismatch (or [`Map::iter`], which retrieves keys) is. Debug builds
assert
+ /// against this misuse.
+ pub fn get(&self, key: &K) -> Result<Option<V>> {
+ if !self.try_contains_key(key)? {
+ self.debug_assert_key_type();
+ return Ok(None);
+ }
+ let result = crate::cached_global_func!("ffi.MapGetItem")
+ .call_packed(&[AnyView::from(self), AnyView::from(key)])?;
+ let value =
TryFromTemp::<V>::try_from(result).map(TryFromTemp::into_value)?;
+ Ok(Some(value))
+ }
+
+ /// Returns an iterator over the `(key, value)` pairs of the map.
+ pub fn iter(&self) -> MapItems<K, V> {
+ self.make_iter(|f| (iter_read::<K>(f, 0, "key"), iter_read::<V>(f, 1,
"value")))
+ }
+
+ /// Returns an iterator over the keys of the map.
+ pub fn keys(&self) -> MapKeys<K> {
+ self.make_iter(|f| iter_read::<K>(f, 0, "key"))
+ }
+
+ /// Returns an iterator over the values of the map.
+ pub fn values(&self) -> MapValues<V> {
+ self.make_iter(|f| iter_read::<V>(f, 1, "value"))
+ }
+
+ /// Builds a [`MapIter`] whose `read` extracts each entry as `T`. The
+ /// forward-iteration functor is requested only for a non-empty map, so
+ /// iterating an empty map makes no FFI call and allocates no functor.
+ fn make_iter<T>(&self, read: fn(&Function) -> T) -> MapIter<T> {
+ let remaining = self.len();
+ MapIter {
+ functor: (remaining != 0).then(|| self.iter_functor()),
+ remaining,
+ _keepalive: self.data.clone(),
+ read,
+ }
+ }
+
+ /// Obtains a fresh stateful forward-iteration functor from the C++
runtime.
+ fn iter_functor(&self) -> Function {
+ let result = crate::cached_global_func!("ffi.MapForwardIterFunctor")
+ .call_packed(&[AnyView::from(self)])
+ .expect("ffi.MapForwardIterFunctor call failed");
+ Function::try_from(result).expect("ffi.MapForwardIterFunctor returned
a non-function")
+ }
+
+ /// Reads all entries as raw `(Any, Any)` pairs without converting to
+ /// `K`/`V`, propagating FFI failures instead of panicking (unlike the
+ /// iterator API): this backs `check_any_strict` and
+ /// `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();
+ 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)])?;
+ }
+ }
+}
+
+impl<K, V> Default for Map<K, V>
+where
+ K: AnyCompatible,
+ V: AnyCompatible,
+{
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl<K, V> Debug for Map<K, V>
+where
+ K: AnyCompatible,
+ V: AnyCompatible,
+{
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ fn short(name: &str) -> &str {
+ name.split("::").last().unwrap_or(name)
+ }
+ write!(
+ f,
+ "Map<{}, {}>[{}]",
+ short(std::any::type_name::<K>()),
+ short(std::any::type_name::<V>()),
+ self.len()
+ )
+ }
+}
+
+impl<K, V> FromIterator<(K, V)> for Map<K, V>
+where
+ K: AnyCompatible,
+ V: AnyCompatible,
+{
+ /// Duplicate keys follow C++ `ffi.Map` semantics: a later pair overwrites
an
+ /// earlier one, so the resulting map may be smaller than the iterator.
+ fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
+ let pairs: Vec<(K, V)> = iter.into_iter().collect();
+ Self::from_pairs(&pairs).expect("ffi.Map() failed to construct a map")
+ }
+}
+
+// --- Iterators ---
+//
+// The functor from `ffi.MapForwardIterFunctor` does not keep the map alive, so
+// each iterator carries an `ObjectArc<MapObj>` keepalive. Functor command
codes:
+// `0` = read key, `1` = read value, `2` = advance.
+//
+// These are `ExactSizeIterator`s, so a key/value not matching the declared
+// `K`/`V` panics rather than ending iteration early (which would drop
entries).
+// Use [`Map::get`], which returns `Err` on a type mismatch, when the element
+// types are uncertain.
+
+/// Reads the functor's current key (`command` 0) or value (`command` 1) as
`T`,
+/// panicking on a type mismatch (see the note above on `ExactSizeIterator`).
+fn iter_read<T: AnyCompatible>(functor: &Function, command: i64, kind: &str)
-> T {
+ let any = functor
+ .call_packed(&[AnyView::from(&command)])
+ .expect("map iterator: reading current element failed");
+ TryFromTemp::<T>::try_from(any)
+ .map(TryFromTemp::into_value)
+ .unwrap_or_else(|_| panic!("map iterator: {kind} does not match the
map's {kind} type"))
+}
+
+/// Consumes one entry: decrements `remaining` and advances the functor unless
+/// the map is now exhausted. Callers must guard `remaining > 0` (every `next`
+/// returns early when it hits 0), so the decrement can never underflow.
+fn iter_advance(functor: &Function, remaining: &mut usize) {
+ debug_assert!(
+ *remaining > 0,
+ "iter_advance called with no remaining entries"
+ );
+ *remaining -= 1;
+ if *remaining > 0 {
+ functor
+ .call_packed(&[AnyView::from(&2i64)])
+ .expect("map iterator: advancing failed");
+ }
+}
+
+/// Forward iterator over a [`Map`], yielding `T` per entry. The `read` fn
pulls
+/// the current entry (key, value, or both) from the C++ functor, so each
variant
+/// touches only the command(s) it needs — `keys()` never reads values, and
vice
+/// versa. Created via the [`MapItems`] / [`MapKeys`] / [`MapValues`] aliases.
+pub struct MapIter<T> {
+ /// `None` for an empty map (no functor is requested); `Some` while at
least
+ /// one entry remains to be yielded.
+ functor: Option<Function>,
+ remaining: usize,
+ _keepalive: ObjectArc<MapObj>,
+ read: fn(&Function) -> T,
+}
+
+impl<T> Iterator for MapIter<T> {
+ type Item = T;
+
+ fn next(&mut self) -> Option<T> {
+ if self.remaining == 0 {
+ return None;
+ }
+ // `remaining > 0` implies `functor` is `Some` (see `Map::make_iter`).
+ let functor = self
+ .functor
+ .as_ref()
+ .expect("non-empty map iterator has a functor");
+ let item = (self.read)(functor);
+ iter_advance(functor, &mut self.remaining);
+ Some(item)
+ }
+
+ fn size_hint(&self) -> (usize, Option<usize>) {
+ (self.remaining, Some(self.remaining))
+ }
+}
+
+impl<T> ExactSizeIterator for MapIter<T> {}
+
+/// Iterator over `(key, value)` pairs, created by [`Map::iter`].
+pub type MapItems<K, V> = MapIter<(K, V)>;
+/// Iterator over keys, created by [`Map::keys`].
+pub type MapKeys<K> = MapIter<K>;
+/// Iterator over values, created by [`Map::values`].
+pub type MapValues<V> = MapIter<V>;
+
+impl<K, V> IntoIterator for &Map<K, V>
+where
+ K: AnyCompatible,
+ V: AnyCompatible,
+{
+ type Item = (K, V);
+ type IntoIter = MapItems<K, V>;
+
+ fn into_iter(self) -> Self::IntoIter {
+ self.iter()
+ }
+}
+
+// --- Any Type System Conversions ---
+
+unsafe impl<K, V> AnyCompatible for Map<K, V>
+where
+ K: AnyCompatible,
+ V: AnyCompatible,
+{
+ fn type_str() -> String {
+ format!("Map<{}, {}>", K::type_str(), V::type_str())
+ }
+
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ // Mirrors C++ `CheckAnyStrict` (map_base.h): every entry must strictly
+ // match `K`/`V`. An FFI failure reads as "no match" — this fn has no
+ // error channel and must not panic (see `try_raw_entries`).
+ if data.type_index != TypeIndex::kTVMFFIMap as i32 {
+ return false;
+ }
+ let map = Self::copy_from_any_view_after_check(data);
+ match map.try_raw_entries() {
+ Ok(entries) => entries
+ .iter()
+ .all(|(k, v)| k.try_as::<K>().is_some() &&
v.try_as::<V>().is_some()),
+ Err(_) => false,
+ }
+ }
+
+ unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIMap as i32;
+ data.data_union.v_obj = ObjectArc::as_raw(Self::data(src)) as *mut
TVMFFIObject;
+ data.small_str_len = 0;
+ }
+
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIMap as i32;
+ data.data_union.v_obj = ObjectArc::into_raw(Self::into_data(src)) as
*mut TVMFFIObject;
+ data.small_str_len = 0;
+ }
+
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
+ let ptr = data.data_union.v_obj as *const MapObj;
+ crate::object::unsafe_::inc_ref(ptr as *mut TVMFFIObject);
+ Self::from_data(ObjectArc::from_raw(ptr))
+ }
+
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
+ let ptr = data.data_union.v_obj as *const MapObj;
+ let obj = Self::from_data(ObjectArc::from_raw(ptr));
+ data.type_index = TypeIndex::kTVMFFINone as i32;
+ data.data_union.v_int64 = 0;
+ obj
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index != TypeIndex::kTVMFFIMap as i32 {
+ return Err(());
+ }
+
+ // Fast path: if all entries match strictly, we can just copy the
reference.
+ if Self::check_any_strict(data) {
+ return Ok(Self::copy_from_any_view_after_check(data));
+ }
+
+ // Slow path: try to convert entry by entry into a new map, as C++
+ // `TryCastFromAnyView` does.
+ let src = Self::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 = TryFromTemp::<K>::try_from(k).map_err(|_| ())?;
+ let v = TryFromTemp::<V>::try_from(v).map_err(|_| ())?;
+ pairs.push((TryFromTemp::into_value(k),
TryFromTemp::into_value(v)));
+ }
+ Self::from_pairs(&pairs).map_err(|_| ())
+ }
+}
+
+impl<K, V> TryFrom<Any> for Map<K, V>
+where
+ K: AnyCompatible,
+ V: AnyCompatible,
+{
+ type Error = Error;
+
+ fn try_from(value: Any) -> Result<Self> {
+ let temp: TryFromTemp<Self> = TryFromTemp::try_from(value)?;
+ Ok(TryFromTemp::into_value(temp))
+ }
+}
+
+impl<'a, K, V> TryFrom<AnyView<'a>> for Map<K, V>
+where
+ K: AnyCompatible,
+ V: AnyCompatible,
+{
+ type Error = Error;
+
+ fn try_from(value: AnyView<'a>) -> Result<Self> {
+ let temp: TryFromTemp<Self> = TryFromTemp::try_from(value)?;
+ Ok(TryFromTemp::into_value(temp))
+ }
+}
diff --git a/rust/tvm-ffi/src/collections/mod.rs
b/rust/tvm-ffi/src/collections/mod.rs
index ad17dcca..791ff755 100644
--- a/rust/tvm-ffi/src/collections/mod.rs
+++ b/rust/tvm-ffi/src/collections/mod.rs
@@ -18,5 +18,6 @@
*/
/// Collection types
pub mod array;
+pub mod map;
pub mod shape;
pub mod tensor;
diff --git a/rust/tvm-ffi/src/extra/module.rs b/rust/tvm-ffi/src/extra/module.rs
index e3009e26..82630dfb 100644
--- a/rust/tvm-ffi/src/extra/module.rs
+++ b/rust/tvm-ffi/src/extra/module.rs
@@ -51,10 +51,8 @@ impl Module {
/// # Returns
/// * `Result<Module>` - A `Module` instance on success
pub fn load_from_file<Str: AsRef<str>>(file_name: Str) -> Result<Module> {
- static API_FUNC: std::sync::LazyLock<Function> =
- std::sync::LazyLock::new(||
Function::get_global("ffi.ModuleLoadFromFile").unwrap());
let file_name = crate::string::String::from(file_name);
- (*API_FUNC)
+ crate::cached_global_func!("ffi.ModuleLoadFromFile")
.call_tuple_with_len::<1, _>((file_name,))?
.try_into()
}
@@ -67,10 +65,8 @@ impl Module {
/// # Returns
/// * `Result<Function>` - A `Function` instance on success
pub fn get_function<Str: AsRef<str>>(&self, name: Str) -> Result<Function>
{
- static API_FUNC: std::sync::LazyLock<Function> =
- std::sync::LazyLock::new(||
Function::get_global("ffi.ModuleGetFunction").unwrap());
let name = crate::string::String::from(name);
- (*API_FUNC)
+ crate::cached_global_func!("ffi.ModuleGetFunction")
.call_tuple_with_len::<3, _>((self, name, true))?
.try_into()
}
diff --git a/rust/tvm-ffi/src/function_internal.rs
b/rust/tvm-ffi/src/function_internal.rs
index e059051c..ffd0b634 100644
--- a/rust/tvm-ffi/src/function_internal.rs
+++ b/rust/tvm-ffi/src/function_internal.rs
@@ -152,6 +152,33 @@ crate::impl_arg_into_ref!(
bool, i8, i16, i32, i64, isize, u8, u16, u32, u64, usize, f32, f64,
String, Bytes
);
+// `Map<K, V>` passes by value/reference like the scalars above, but its type
+// parameters keep it out of the `impl_*!` macros, so the impls are spelled
out.
+impl<K: AnyCompatible, V: AnyCompatible> IntoArgHolder for crate::Map<K, V> {
+ type Target = crate::Map<K, V>;
+ fn into_arg_holder(self) -> Self::Target {
+ self
+ }
+}
+impl<'a, K: AnyCompatible, V: AnyCompatible> IntoArgHolder for &'a
crate::Map<K, V> {
+ type Target = &'a crate::Map<K, V>;
+ fn into_arg_holder(self) -> Self::Target {
+ self
+ }
+}
+impl<K: AnyCompatible, V: AnyCompatible> ArgIntoRef for crate::Map<K, V> {
+ type Target = crate::Map<K, V>;
+ fn to_ref(&self) -> &Self::Target {
+ self
+ }
+}
+impl<K: AnyCompatible, V: AnyCompatible> ArgIntoRef for &crate::Map<K, V> {
+ type Target = crate::Map<K, V>;
+ fn to_ref(&self) -> &Self::Target {
+ self
+ }
+}
+
//-----------------------------------------------------------
// TupleAsPackedArgs
//
diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs
index fad82601..891f734e 100644
--- a/rust/tvm-ffi/src/lib.rs
+++ b/rust/tvm-ffi/src/lib.rs
@@ -33,6 +33,7 @@ pub use tvm_ffi_sys;
pub use crate::any::{Any, AnyView};
pub use crate::collections::array::Array;
+pub use crate::collections::map::Map;
pub use crate::collections::shape::Shape;
pub use crate::collections::tensor::{CPUNDAlloc, NDAllocator, Tensor};
pub use crate::device::{current_stream, with_stream};
diff --git a/rust/tvm-ffi/src/macros.rs b/rust/tvm-ffi/src/macros.rs
index 7c2cad74..b9665623 100644
--- a/rust/tvm-ffi/src/macros.rs
+++ b/rust/tvm-ffi/src/macros.rs
@@ -40,6 +40,25 @@ macro_rules! function_name {
}};
}
+/// Resolves a registered global [`Function`](crate::function::Function) by
name
+/// and caches it for the lifetime of the process (lock-free after the first
+/// lookup). Each call site gets its own cache and the macro evaluates to a
+/// `&'static Function`. Panics — naming the function — if it is not
registered.
+///
+/// Usage: `cached_global_func!("ffi.Map").call_packed(args)`
+#[macro_export]
+macro_rules! cached_global_func {
+ ($name:literal) => {{
+ static FUNC: std::sync::LazyLock<$crate::function::Function> =
+ std::sync::LazyLock::new(|| {
+
$crate::function::Function::get_global($name).unwrap_or_else(|_| {
+ panic!(concat!("global function `", $name, "` is not
registered"))
+ })
+ });
+ &*FUNC
+ }};
+}
+
/// Check the return code of the safe call
///
/// # Arguments
diff --git a/rust/tvm-ffi/tests/test_map.rs b/rust/tvm-ffi/tests/test_map.rs
new file mode 100644
index 00000000..ff50c576
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_map.rs
@@ -0,0 +1,288 @@
+/*
+ * 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.
+ */
+use tvm_ffi::*;
+
+/// Helper to create a CPU `f32` tensor whose first element is `val`.
+fn create_tensor(val: f32, shape: &[i64]) -> Tensor {
+ let dtype = DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1);
+ let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
+ let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, shape, dtype, device);
+ if let Ok(slice) = tensor.data_as_slice_mut::<f32>() {
+ slice[0] = val;
+ }
+ tensor
+}
+
+/// Helper to read the first `f32` of a tensor.
+fn get_val(tensor: &Tensor) -> f32 {
+ tensor
+ .data_as_slice::<f32>()
+ .expect("Type mismatch or null")[0]
+}
+
+#[test]
+fn test_map_basic_lookup() {
+ let map: Map<i64, i64> = [(1i64, 10i64), (2, 20), (3,
30)].into_iter().collect();
+
+ assert_eq!(map.len(), 3);
+ assert!(!map.is_empty());
+
+ assert_eq!(map.get(&1).unwrap(), Some(10));
+ assert_eq!(map.get(&2).unwrap(), Some(20));
+ assert_eq!(map.get(&3).unwrap(), Some(30));
+ assert_eq!(map.get(&99).unwrap(), None);
+
+ assert!(map.contains_key(&1));
+ assert!(!map.contains_key(&99));
+}
+
+#[test]
+fn test_map_empty() {
+ let map: Map<i64, i64> = Map::new();
+ assert_eq!(map.len(), 0);
+ assert!(map.is_empty());
+ assert_eq!(map.get(&1).unwrap(), None);
+ // All three iterators must handle the empty (remaining == 0) case.
+ assert_eq!(map.iter().count(), 0);
+ assert_eq!(map.keys().count(), 0);
+ assert_eq!(map.values().count(), 0);
+}
+
+/// `FromIterator` follows C++ `ffi.Map` last-wins semantics for duplicate
keys,
+/// so the resulting map is smaller than the input iterator.
+#[test]
+fn test_map_from_iter_duplicate_keys_last_wins() {
+ let map: Map<i64, i64> = [(1i64, 10i64), (1, 20), (2,
30)].into_iter().collect();
+ assert_eq!(map.len(), 2);
+ assert_eq!(map.get(&1).unwrap(), Some(20));
+ assert_eq!(map.get(&2).unwrap(), Some(30));
+}
+
+#[test]
+fn test_map_iteration() {
+ let map: Map<i64, i64> = [(1i64, 10i64), (2, 20), (3,
30)].into_iter().collect();
+
+ let mut items: Vec<(i64, i64)> = map.iter().collect();
+ items.sort();
+ assert_eq!(items, vec![(1, 10), (2, 20), (3, 30)]);
+
+ let mut keys: Vec<i64> = map.keys().collect();
+ keys.sort();
+ assert_eq!(keys, vec![1, 2, 3]);
+
+ let mut values: Vec<i64> = map.values().collect();
+ values.sort();
+ assert_eq!(values, vec![10, 20, 30]);
+
+ // `&map` IntoIterator yields the same pairs.
+ let mut via_ref: Vec<(i64, i64)> = (&map).into_iter().collect();
+ via_ref.sort();
+ assert_eq!(via_ref, vec![(1, 10), (2, 20), (3, 30)]);
+}
+
+#[test]
+fn test_map_string_keys() {
+ let map: Map<String, i64> = [
+ (String::from("a"), 1i64),
+ (String::from("b"), 2),
+ (String::from("c"), 3),
+ ]
+ .into_iter()
+ .collect();
+
+ assert_eq!(map.len(), 3);
+ assert_eq!(map.get(&String::from("a")).unwrap(), Some(1));
+ assert_eq!(map.get(&String::from("c")).unwrap(), Some(3));
+ assert_eq!(map.get(&String::from("z")).unwrap(), None);
+ assert!(map.contains_key(&String::from("b")));
+}
+
+#[test]
+fn test_map_any_roundtrip() {
+ let map: Map<i64, i64> = [(1i64, 10i64), (2, 20)].into_iter().collect();
+
+ let any = Any::from(map.clone());
+ assert_eq!(any.type_index(), TypeIndex::kTVMFFIMap as i32);
+
+ let back: Map<i64, i64> = Map::try_from(any).expect("Any -> Map failed");
+ assert_eq!(back.len(), 2);
+ assert_eq!(back.get(&1).unwrap(), Some(10));
+ assert_eq!(back.get(&2).unwrap(), Some(20));
+}
+
+#[test]
+fn test_map_shares_underlying_object() {
+ let map: Map<i64, i64> = [(1i64, 10i64)].into_iter().collect();
+ // Cloning shares the same underlying MapObj rather than copying entries.
+ let clone = map.clone();
+ assert_eq!(clone.len(), 1);
+ assert_eq!(clone.get(&1).unwrap(), Some(10));
+}
+
+#[test]
+fn test_map_object_values_and_refcount() {
+ let t = create_tensor(7.0, &[2, 3]);
+ let base = AnyView::from(&t)
+ .debug_strong_count()
+ .expect("tensor is reference counted");
+ assert_eq!(base, 1);
+
+ // Building the map stores one internal reference to the tensor object.
+ let map: Map<String, Tensor> = [(String::from("x"),
t.clone())].into_iter().collect();
+ assert_eq!(AnyView::from(&t).debug_strong_count().unwrap(), base + 1);
+
+ // The value round-trips correctly and `get` hands back a fresh handle.
+ let got = map
+ .get(&String::from("x"))
+ .unwrap()
+ .expect("key should be present");
+ assert_eq!(get_val(&got), 7.0);
+ assert_eq!(got.ndim(), 2);
+ assert_eq!(AnyView::from(&t).debug_strong_count().unwrap(), base + 2);
+
+ // Iteration yields object handles too.
+ let collected: Vec<(String, Tensor)> = map.iter().collect();
+ assert_eq!(collected.len(), 1);
+ assert_eq!(get_val(&collected[0].1), 7.0);
+ // Peak: `t` + the map's internal ref + `got` + the iterated handle.
+ assert_eq!(AnyView::from(&t).debug_strong_count().unwrap(), base + 3);
+
+ // Dropping the borrowed handles restores the count to just `t` + the map.
+ drop(collected);
+ drop(got);
+ assert_eq!(AnyView::from(&t).debug_strong_count().unwrap(), base + 1);
+
+ // Dropping the map releases its internal reference; only `t` remains.
+ drop(map);
+ assert_eq!(AnyView::from(&t).debug_strong_count().unwrap(), base);
+}
+
+#[test]
+fn test_map_as_function_argument() {
+ let map: Map<i64, i64> = [(1i64, 10i64), (2, 20), (3,
30)].into_iter().collect();
+
+ // Exercises the `ArgIntoRef`/`call_tuple` path, by value and by reference.
+ let map_size = Function::get_global("ffi.MapSize").unwrap();
+ let n_by_value: i64 = map_size
+ .call_tuple((map.clone(),))
+ .unwrap()
+ .try_into()
+ .unwrap();
+ assert_eq!(n_by_value, 3);
+ let n_by_ref: i64 =
map_size.call_tuple((&map,)).unwrap().try_into().unwrap();
+ assert_eq!(n_by_ref, 3);
+
+ // Two-argument call mixing `&Map` with a key value.
+ let get_item = Function::get_global("ffi.MapGetItem").unwrap();
+ let v: i64 = get_item
+ .call_tuple((&map, 2i64))
+ .unwrap()
+ .try_into()
+ .unwrap();
+ assert_eq!(v, 20);
+}
+
+/// `get` distinguishes a present object value from an absent key even when `V`
+/// is object-typed (an absent lookup must not be miscast into a `V`).
+#[test]
+fn test_map_get_missing_with_object_values() {
+ let map: Map<String, Tensor> = [(String::from("x"), create_tensor(7.0,
&[2]))]
+ .into_iter()
+ .collect();
+
+ let present = map.get(&String::from("x")).unwrap();
+ assert!(present.is_some());
+ assert_eq!(get_val(&present.unwrap()), 7.0);
+
+ // Absent key returns `None`, never an object miscast as a tensor.
+ assert!(map.get(&String::from("missing")).unwrap().is_none());
+}
+
+/// Views a map behind an `Any` as `Map<K, V>` while bypassing the strict
+/// element check (`try_from` now walks all entries, C++-style, so a mistyped
+/// view can only be produced this way), to exercise the access-time guards.
+fn view_map_unchecked<K, V>(mut any: Any) -> Map<K, V>
+where
+ K: AnyCompatible,
+ V: AnyCompatible,
+{
+ assert_eq!(any.type_index(), TypeIndex::kTVMFFIMap as i32);
+ unsafe { Map::copy_from_any_view_after_check(&*any.as_data_ptr()) }
+}
+
+/// Builds an `i64`-valued map but views it as `Map<_, String>`, bypassing the
+/// strict check, so the mismatch only surfaces on access.
+fn mistyped_value_map() -> Map<i64, String> {
+ let real: Map<i64, i64> = [(1i64, 10i64)].into_iter().collect();
+ view_map_unchecked(Any::from(real))
+}
+
+/// `get` reports a value type mismatch as `Err`, not a panic.
+#[test]
+fn test_map_get_type_mismatch_is_err() {
+ let map = mistyped_value_map();
+ assert!(map.get(&1).is_err());
+}
+
+/// Iterators panic on a value type mismatch rather than silently truncating
+/// (which would violate their `ExactSizeIterator` contract).
+#[test]
+#[should_panic(expected = "value does not match")]
+fn test_map_iter_type_mismatch_panics() {
+ let map = mistyped_value_map();
+ let _ = map.values().next();
+}
+
+/// A key whose type does not match the map's key type reads as absent (keys
are
+/// hashed, not retrieved, so a pure lookup can't distinguish it from a real
+/// miss). In debug builds the misuse is caught by an assertion on the miss
path.
+#[cfg(debug_assertions)]
+#[test]
+#[should_panic(expected = "does not match the map's stored key type")]
+fn test_map_get_mistyped_key_debug_asserts() {
+ // Real `Map<i64, i64>` viewed (unchecked) as `Map<String, i64>`: the i64
+ // keys do not match `K = String`.
+ let real: Map<i64, i64> = [(1i64, 10i64)].into_iter().collect();
+ let map: Map<String, i64> = view_map_unchecked(Any::from(real));
+ let _ = map.get(&String::from("anything"));
+}
+
+/// `try_from` rejects a map whose values can neither strictly match nor cast
+/// to `V` (aligned with C++ `CheckAnyStrict`/`TryCastFromAnyView`).
+#[test]
+fn test_map_any_cast_value_mismatch_is_err() {
+ let real: Map<i64, i64> = [(1i64, 10i64)].into_iter().collect();
+ assert!(Map::<i64, String>::try_from(Any::from(real)).is_err());
+}
+
+/// When entries do not strictly match but are castable (i64 -> f64),
+/// `try_from` builds a NEW map with converted entries instead of sharing the
+/// object (C++ slow path).
+#[test]
+fn test_map_any_cast_slow_path_converts() {
+ let real: Map<i64, i64> = [(1i64, 10i64), (2, 20)].into_iter().collect();
+ let converted: Map<f64, f64> =
+ Map::try_from(Any::from(real.clone())).expect("i64 entries cast to
f64");
+ assert_eq!(converted.len(), 2);
+ assert_eq!(converted.get(&1.0).unwrap(), Some(10.0));
+ assert_eq!(converted.get(&2.0).unwrap(), Some(20.0));
+ // Each map owns its own object: nothing is shared between them.
+ assert_eq!(AnyView::from(&real).debug_strong_count(), Some(1));
+ assert_eq!(AnyView::from(&converted).debug_strong_count(), Some(1));
+}
