tvm_ffi/collections/
map.rs

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//! Immutable [`Map`] container backed by the C++ `ffi.Map` object.
20//!
21//! The `MapObj` header mirrors C++ `MapBaseObj`, so [`Map::len`] reads `size`
22//! directly (no FFI). The hash-table storage itself is an implementation detail,
23//! so lookups and iteration are delegated to the global functions the C++ runtime
24//! registers (`ffi.MapGetItem`, `ffi.MapForwardIterFunctor`, ...): there is no
25//! Map-specific C ABI, and a Rust re-implementation would have to replicate the
26//! hashing (`AnyHash`/`AnyEqual`) and dense/small probing, just as the Python
27//! bindings also delegate.
28//!
29//! `Any` conversions follow C++ `MapTypeTraitsBase` (map_base.h): the strict
30//! check walks every entry, and a failed strict `try_from` falls back to
31//! converting the entries one by one into a new map.
32use std::fmt::Debug;
33use std::marker::PhantomData;
34use std::ops::Deref;
35
36use crate::any::TryFromTemp;
37use crate::derive::Object;
38use crate::function::Function;
39use crate::object::{Object, ObjectArc};
40use crate::type_traits::ContainerElement;
41use crate::{Any, AnyCompatible, AnyView, Error, ObjectRefCore, Result};
42use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
43use tvm_ffi_sys::{TVMFFIAny, TVMFFIObject};
44
45#[inline]
46fn element_view<T: ContainerElement>(value: &T) -> AnyView<'_> {
47    unsafe {
48        let mut data = TVMFFIAny::new();
49        T::container_copy_to_any_view(value, &mut data);
50        AnyView::from_raw_ffi_any(data)
51    }
52}
53
54fn element_from_any<T: ContainerElement>(value: Any) -> Result<T> {
55    unsafe {
56        if T::container_check_any_strict(value.as_raw_ffi_any()) {
57            let mut value = std::mem::ManuallyDrop::new(value);
58            return Ok(T::container_move_from_any_after_check(
59                &mut *value.as_data_ptr(),
60            ));
61        }
62        T::container_try_cast_from_any_view(value.as_raw_ffi_any()).map_err(|()| {
63            let message = format!(
64                "Cannot convert from type `{}` to `{}`",
65                T::container_get_mismatch_type_info(value.as_raw_ffi_any()),
66                T::container_type_str()
67            );
68            Error::new(crate::error::TYPE_ERROR, &message, "")
69        })
70    }
71}
72
73/// Container object for [`Map`]. The header fields mirror C++ `MapBaseObj`
74/// (`include/tvm/ffi/container/map_base.h`) so [`Map::len`] can read `size`
75/// without an FFI call; the storage `data` points to stays opaque.
76///
77/// This is a partial mirror by design: `MapBaseObj` has one more trailing field
78/// after `slots_` (a `data_deleter_` function pointer), omitted here because this
79/// binding is read-only — it only reads fields up to `size`, which precede it,
80/// and never allocates a `MapObj` itself (the C++ runtime allocates the larger
81/// `Small`/`DenseMapBaseObj` subclass), so the trailing field is never touched.
82/// The layout also assumes `TVM_FFI_DEBUG_WITH_ABI_CHANGE` is off (the default):
83/// with that debug flag set, `MapBaseObj` gains a *leading* `state_marker` field
84/// that would shift every offset below.
85#[repr(C)]
86#[derive(Object)]
87#[type_key = "ffi.Map"]
88#[type_index(TypeIndex::kTVMFFIMap)]
89pub struct MapObj {
90    pub object: Object,
91    /// Pointer to the (opaque) key/value storage region (`MapBaseObj::data_`).
92    pub data: *mut core::ffi::c_void,
93    /// Number of entries (`MapBaseObj::size_`).
94    pub size: u64,
95    /// Number of hash slots; the MSB is a small-map tag (`MapBaseObj::slots_`).
96    pub slots: u64,
97}
98
99/// Immutable, reference-counted map from `K` to `V`, sharing its underlying
100/// `MapObj` with C++. Cloning is cheap (it bumps the refcount).
101#[repr(C)]
102pub struct Map<K, V> {
103    data: ObjectArc<MapObj>,
104    _marker: PhantomData<(K, V)>,
105}
106
107// Manual `Clone`: the data is the shared `ObjectArc`, so `Map<K, V>` is `Clone`
108// regardless of whether `K`/`V` are (they are only phantom markers).
109impl<K, V> Clone for Map<K, V> {
110    fn clone(&self) -> Self {
111        Self {
112            data: self.data.clone(),
113            _marker: PhantomData,
114        }
115    }
116}
117
118unsafe impl<K, V> ObjectRefCore for Map<K, V> {
119    type ContainerType = MapObj;
120
121    fn data(this: &Self) -> &ObjectArc<MapObj> {
122        &this.data
123    }
124
125    fn into_data(this: Self) -> ObjectArc<MapObj> {
126        this.data
127    }
128
129    unsafe fn from_data(data: ObjectArc<MapObj>) -> Self {
130        Self {
131            data,
132            _marker: PhantomData,
133        }
134    }
135}
136
137// A `Map<K, V>` is a counted handle to its `MapObj`, so it derefs to it (like
138// `ObjectArc` does): methods can read header fields as `self.size` instead of
139// `self.data.size`.
140impl<K, V> Deref for Map<K, V> {
141    type Target = MapObj;
142    #[inline]
143    fn deref(&self) -> &MapObj {
144        &self.data
145    }
146}
147
148impl<K, V> Map<K, V>
149where
150    K: ContainerElement,
151    V: ContainerElement,
152{
153    /// Creates a new, empty map (via an `ffi.Map()` call to the C++ runtime).
154    pub fn new() -> Self {
155        Self::from_pairs(&[]).expect("ffi.Map() failed to construct an empty map")
156    }
157
158    /// Builds a map by calling the C++ `ffi.Map` constructor with a flattened
159    /// `[k0, v0, k1, v1, ...]` argument list.
160    fn from_pairs(pairs: &[(K, V)]) -> Result<Self> {
161        let mut args: Vec<AnyView<'_>> = Vec::with_capacity(pairs.len() * 2);
162        for (k, v) in pairs {
163            args.push(element_view(k));
164            args.push(element_view(v));
165        }
166        let result = crate::cached_global_func!("ffi.Map").call_packed(&args)?;
167        Self::try_from(result)
168    }
169
170    /// Returns the number of entries in the map by reading the `MapObj` header
171    /// directly (no FFI call), like [`Array::len`](crate::Array).
172    pub fn len(&self) -> usize {
173        self.size as usize
174    }
175
176    /// Returns `true` if the map contains no entries.
177    pub fn is_empty(&self) -> bool {
178        self.len() == 0
179    }
180
181    /// Returns whether `key` is present, propagating any FFI error (e.g. a key
182    /// the C++ runtime cannot hash) rather than panicking. Backs both
183    /// [`Map::contains_key`] and [`Map::get`].
184    fn try_contains_key(&self, key: &K) -> Result<bool> {
185        let result = crate::cached_global_func!("ffi.MapCount")
186            .call_packed(&[AnyView::from(self), element_view(key)])?;
187        Ok(i64::try_from(result)? != 0)
188    }
189
190    /// In debug builds, panics if `K` does not match the map's actual key type,
191    /// checked against one stored key. Compiles to nothing in release, so
192    /// [`Map::get`] / [`Map::contains_key`] keep their single-FFI-call fast path.
193    ///
194    /// A pure key lookup *hashes* the key rather than retrieving one, so a wrong
195    /// `K` simply fails to match and reads as "absent" — unlike [`Map::iter`],
196    /// which retrieves keys and surfaces the mismatch. This assertion catches
197    /// that misuse in dev/tests without taxing release lookups. Only meaningful
198    /// on a miss (a hit already proves `K` matched a stored key).
199    #[inline]
200    fn debug_assert_key_type(&self) {
201        #[cfg(debug_assertions)]
202        {
203            if !self.is_empty() {
204                let functor = self.iter_functor();
205                let first_key = functor
206                    .call_packed(&[AnyView::from(&0i64)])
207                    .expect("map iterator: reading current key failed");
208                assert!(
209                    unsafe { K::container_check_any_strict(first_key.as_raw_ffi_any()) },
210                    "Map lookup: key type `{}` does not match the map's stored key type",
211                    std::any::type_name::<K>(),
212                );
213            }
214        }
215    }
216
217    /// Returns `true` if `key` is present in the map.
218    ///
219    /// Panics if the lookup fails in the C++ runtime (e.g. `key` is not
220    /// hashable); use [`Map::get`] when such failures must be recovered from.
221    /// Passing a `key` whose *type* does not match the map's key type is
222    /// **undefined** (currently reads as `false`; see [`Map::get`]); debug builds
223    /// assert against it.
224    pub fn contains_key(&self, key: &K) -> bool {
225        let present = self
226            .try_contains_key(key)
227            .expect("ffi.MapCount call failed");
228        if !present {
229            self.debug_assert_key_type();
230        }
231        present
232    }
233
234    /// Looks up `key`, returning `Ok(None)` if absent, or `Err` if the lookup
235    /// fails in the C++ runtime (e.g. `key` is not hashable) or the stored value
236    /// cannot be converted to `V`. The map is immutable, so the lookup
237    /// (existence check then `ffi.MapGetItem`) cannot race with itself.
238    ///
239    /// Passing a `key` whose *type* does not match the map's key type is
240    /// **undefined**: the result is unspecified — it currently reads as absent
241    /// (`Ok(None)`) because a pure lookup hashes the key rather than retrieving
242    /// one, so the mismatch cannot be surfaced as an `Err` the way a value
243    /// mismatch (or [`Map::iter`], which retrieves keys) is. Debug builds assert
244    /// against this misuse.
245    pub fn get(&self, key: &K) -> Result<Option<V>> {
246        if !self.try_contains_key(key)? {
247            self.debug_assert_key_type();
248            return Ok(None);
249        }
250        let result = crate::cached_global_func!("ffi.MapGetItem")
251            .call_packed(&[AnyView::from(self), element_view(key)])?;
252        let value = element_from_any(result)?;
253        Ok(Some(value))
254    }
255
256    /// Returns an iterator over the `(key, value)` pairs of the map.
257    pub fn iter(&self) -> MapItems<K, V> {
258        self.make_iter(|f| (iter_read::<K>(f, 0, "key"), iter_read::<V>(f, 1, "value")))
259    }
260
261    /// Returns an iterator over the keys of the map.
262    pub fn keys(&self) -> MapKeys<K> {
263        self.make_iter(|f| iter_read::<K>(f, 0, "key"))
264    }
265
266    /// Returns an iterator over the values of the map.
267    pub fn values(&self) -> MapValues<V> {
268        self.make_iter(|f| iter_read::<V>(f, 1, "value"))
269    }
270
271    /// Builds a [`MapIter`] whose `read` extracts each entry as `T`. The
272    /// forward-iteration functor is requested only for a non-empty map, so
273    /// iterating an empty map makes no FFI call and allocates no functor.
274    fn make_iter<T>(&self, read: fn(&Function) -> T) -> MapIter<T> {
275        let remaining = self.len();
276        MapIter {
277            functor: (remaining != 0).then(|| self.iter_functor()),
278            remaining,
279            _keepalive: self.data.clone(),
280            read,
281        }
282    }
283
284    /// Obtains a fresh stateful forward-iteration functor from the C++ runtime.
285    fn iter_functor(&self) -> Function {
286        let result = crate::cached_global_func!("ffi.MapForwardIterFunctor")
287            .call_packed(&[AnyView::from(self)])
288            .expect("ffi.MapForwardIterFunctor call failed");
289        Function::try_from(result).expect("ffi.MapForwardIterFunctor returned a non-function")
290    }
291
292    /// Reads all entries as raw `(Any, Any)` pairs without converting to
293    /// `K`/`V`, propagating FFI failures instead of panicking (unlike the
294    /// iterator API): this backs `check_any_strict` and
295    /// `try_cast_from_any_view`, which run during argument decoding where a
296    /// panic could unwind across the C ABI.
297    fn try_raw_entries(&self) -> Result<Vec<(Any, Any)>> {
298        let mut entries = Vec::with_capacity(self.len());
299        let mut remaining = self.len();
300        if remaining == 0 {
301            return Ok(entries);
302        }
303        let functor = crate::cached_global_func!("ffi.MapForwardIterFunctor")
304            .call_packed(&[AnyView::from(self)])
305            .and_then(Function::try_from)?;
306        loop {
307            let k = functor.call_packed(&[AnyView::from(&0i64)])?;
308            let v = functor.call_packed(&[AnyView::from(&1i64)])?;
309            entries.push((k, v));
310            remaining -= 1;
311            if remaining == 0 {
312                return Ok(entries);
313            }
314            functor.call_packed(&[AnyView::from(&2i64)])?;
315        }
316    }
317}
318
319impl<K, V> Default for Map<K, V>
320where
321    K: ContainerElement,
322    V: ContainerElement,
323{
324    fn default() -> Self {
325        Self::new()
326    }
327}
328
329impl<K, V> Debug for Map<K, V>
330where
331    K: ContainerElement,
332    V: ContainerElement,
333{
334    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
335        fn short(name: &str) -> &str {
336            name.split("::").last().unwrap_or(name)
337        }
338        write!(
339            f,
340            "Map<{}, {}>[{}]",
341            short(std::any::type_name::<K>()),
342            short(std::any::type_name::<V>()),
343            self.len()
344        )
345    }
346}
347
348impl<K, V> FromIterator<(K, V)> for Map<K, V>
349where
350    K: ContainerElement,
351    V: ContainerElement,
352{
353    /// Duplicate keys follow C++ `ffi.Map` semantics: a later pair overwrites an
354    /// earlier one, so the resulting map may be smaller than the iterator.
355    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
356        let pairs: Vec<(K, V)> = iter.into_iter().collect();
357        Self::from_pairs(&pairs).expect("ffi.Map() failed to construct a map")
358    }
359}
360
361// --- Iterators ---
362//
363// The functor from `ffi.MapForwardIterFunctor` does not keep the map alive, so
364// each iterator carries an `ObjectArc<MapObj>` keepalive. Functor command codes:
365// `0` = read key, `1` = read value, `2` = advance.
366//
367// These are `ExactSizeIterator`s, so a key/value not matching the declared
368// `K`/`V` panics rather than ending iteration early (which would drop entries).
369// Use [`Map::get`], which returns `Err` on a type mismatch, when the element
370// types are uncertain.
371
372/// Reads the functor's current key (`command` 0) or value (`command` 1) as `T`,
373/// panicking on a type mismatch (see the note above on `ExactSizeIterator`).
374fn iter_read<T: ContainerElement>(functor: &Function, command: i64, kind: &str) -> T {
375    let any = functor
376        .call_packed(&[AnyView::from(&command)])
377        .expect("map iterator: reading current element failed");
378    element_from_any(any)
379        .unwrap_or_else(|_| panic!("map iterator: {kind} does not match the map's {kind} type"))
380}
381
382/// Consumes one entry: decrements `remaining` and advances the functor unless
383/// the map is now exhausted. Callers must guard `remaining > 0` (every `next`
384/// returns early when it hits 0), so the decrement can never underflow.
385fn iter_advance(functor: &Function, remaining: &mut usize) {
386    debug_assert!(
387        *remaining > 0,
388        "iter_advance called with no remaining entries"
389    );
390    *remaining -= 1;
391    if *remaining > 0 {
392        functor
393            .call_packed(&[AnyView::from(&2i64)])
394            .expect("map iterator: advancing failed");
395    }
396}
397
398/// Forward iterator over a [`Map`], yielding `T` per entry. The `read` fn pulls
399/// the current entry (key, value, or both) from the C++ functor, so each variant
400/// touches only the command(s) it needs — `keys()` never reads values, and vice
401/// versa. Created via the [`MapItems`] / [`MapKeys`] / [`MapValues`] aliases.
402pub struct MapIter<T> {
403    /// `None` for an empty map (no functor is requested); `Some` while at least
404    /// one entry remains to be yielded.
405    functor: Option<Function>,
406    remaining: usize,
407    _keepalive: ObjectArc<MapObj>,
408    read: fn(&Function) -> T,
409}
410
411impl<T> Iterator for MapIter<T> {
412    type Item = T;
413
414    fn next(&mut self) -> Option<T> {
415        if self.remaining == 0 {
416            return None;
417        }
418        // `remaining > 0` implies `functor` is `Some` (see `Map::make_iter`).
419        let functor = self
420            .functor
421            .as_ref()
422            .expect("non-empty map iterator has a functor");
423        let item = (self.read)(functor);
424        iter_advance(functor, &mut self.remaining);
425        Some(item)
426    }
427
428    fn size_hint(&self) -> (usize, Option<usize>) {
429        (self.remaining, Some(self.remaining))
430    }
431}
432
433impl<T> ExactSizeIterator for MapIter<T> {}
434
435/// Iterator over `(key, value)` pairs, created by [`Map::iter`].
436pub type MapItems<K, V> = MapIter<(K, V)>;
437/// Iterator over keys, created by [`Map::keys`].
438pub type MapKeys<K> = MapIter<K>;
439/// Iterator over values, created by [`Map::values`].
440pub type MapValues<V> = MapIter<V>;
441
442impl<K, V> IntoIterator for &Map<K, V>
443where
444    K: ContainerElement,
445    V: ContainerElement,
446{
447    type Item = (K, V);
448    type IntoIter = MapItems<K, V>;
449
450    fn into_iter(self) -> Self::IntoIter {
451        self.iter()
452    }
453}
454
455// --- Any Type System Conversions ---
456
457unsafe impl<K, V> AnyCompatible for Map<K, V>
458where
459    K: ContainerElement,
460    V: ContainerElement,
461{
462    fn type_str() -> String {
463        format!(
464            "Map<{}, {}>",
465            K::container_type_str(),
466            V::container_type_str()
467        )
468    }
469
470    unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
471        // Mirrors C++ `CheckAnyStrict` (map_base.h): every entry must strictly
472        // match `K`/`V`. An FFI failure reads as "no match" — this fn has no
473        // error channel and must not panic (see `try_raw_entries`).
474        if data.type_index != TypeIndex::kTVMFFIMap as i32 {
475            return false;
476        }
477        let map = <Self as AnyCompatible>::copy_from_any_view_after_check(data);
478        match map.try_raw_entries() {
479            Ok(entries) => entries.iter().all(|(k, v)| unsafe {
480                K::container_check_any_strict(k.as_raw_ffi_any())
481                    && V::container_check_any_strict(v.as_raw_ffi_any())
482            }),
483            Err(_) => false,
484        }
485    }
486
487    unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
488        data.type_index = TypeIndex::kTVMFFIMap as i32;
489        data.data_union.v_obj = ObjectArc::as_raw(Self::data(src)) as *mut TVMFFIObject;
490        data.small_str_len = 0;
491    }
492
493    unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
494        data.type_index = TypeIndex::kTVMFFIMap as i32;
495        data.data_union.v_obj = ObjectArc::into_raw(Self::into_data(src)) as *mut TVMFFIObject;
496        data.small_str_len = 0;
497    }
498
499    unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
500        let ptr = data.data_union.v_obj as *const MapObj;
501        crate::object::unsafe_::inc_ref(ptr as *mut TVMFFIObject);
502        Self::from_data(ObjectArc::from_raw(ptr))
503    }
504
505    unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
506        let ptr = data.data_union.v_obj as *const MapObj;
507        let obj = Self::from_data(ObjectArc::from_raw(ptr));
508        data.type_index = TypeIndex::kTVMFFINone as i32;
509        data.data_union.v_int64 = 0;
510        obj
511    }
512
513    unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
514        if data.type_index != TypeIndex::kTVMFFIMap as i32 {
515            return Err(());
516        }
517
518        // Fast path: if all entries match strictly, we can just copy the reference.
519        if <Self as AnyCompatible>::check_any_strict(data) {
520            return Ok(<Self as AnyCompatible>::copy_from_any_view_after_check(
521                data,
522            ));
523        }
524
525        // Slow path: try to convert entry by entry into a new map, as C++
526        // `TryCastFromAnyView` does.
527        let src = <Self as AnyCompatible>::copy_from_any_view_after_check(data);
528        let mut pairs = Vec::with_capacity(src.len());
529        for (k, v) in src.try_raw_entries().map_err(|_| ())? {
530            let k = element_from_any::<K>(k).map_err(|_| ())?;
531            let v = element_from_any::<V>(v).map_err(|_| ())?;
532            pairs.push((k, v));
533        }
534        Self::from_pairs(&pairs).map_err(|_| ())
535    }
536}
537
538impl<K, V> TryFrom<Any> for Map<K, V>
539where
540    K: ContainerElement,
541    V: ContainerElement,
542{
543    type Error = Error;
544
545    fn try_from(value: Any) -> Result<Self> {
546        let temp: TryFromTemp<Self> = TryFromTemp::try_from(value)?;
547        Ok(TryFromTemp::into_value(temp))
548    }
549}
550
551impl<'a, K, V> TryFrom<AnyView<'a>> for Map<K, V>
552where
553    K: ContainerElement,
554    V: ContainerElement,
555{
556    type Error = Error;
557
558    fn try_from(value: AnyView<'a>) -> Result<Self> {
559        let temp: TryFromTemp<Self> = TryFromTemp::try_from(value)?;
560        Ok(TryFromTemp::into_value(temp))
561    }
562}
tvm_ffi/collections/map.rs

tvm_ffi/collections/
map.rs
