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 e8eb3032 [EXTRA] Compose mutator storage and expose borrowed Rust
attributes (#918)
e8eb3032 is described below
commit e8eb3032c6ea5ee8d839d04d1a4955002fd29af0
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Nov 4 10:59:08 2026 -0500
[EXTRA] Compose mutator storage and expose borrowed Rust attributes (#918)
Allow structural map and mutate engines to compose with state-only
parents derived from `StructuralMutatorObj`. A shared engine layer owns
variable remapping and its reference lifetimes, while existing
`StructuralMapEngineBase` parents retain a single remap environment.
Expose the safe borrowed `TypeAttrColumn::get_raw_ffi_any` lookup and
use consistent `raw_ffi_any` names for Rust carrier helpers. The copied
carrier does not retain or transfer ownership; its documentation
describes registry validity and the owning `get` alternative.
---
include/tvm/ffi/extra/structural_mutate.h | 127 ++++++++++++++++------------
rust/tvm-ffi/src/extra/structural_common.rs | 6 +-
rust/tvm-ffi/src/extra/structural_map.rs | 4 +-
rust/tvm-ffi/src/extra/structural_mutate.rs | 40 ++++-----
rust/tvm-ffi/src/extra/structural_visit.rs | 2 +-
rust/tvm-ffi/src/extra/unchanged.rs | 4 +-
rust/tvm-ffi/src/reflection.rs | 16 +++-
src/ffi/extra/structural_mutate.cc | 12 +--
tests/cpp/extra/test_structural_mutate.cc | 61 +++++++++++--
9 files changed, 174 insertions(+), 98 deletions(-)
diff --git a/include/tvm/ffi/extra/structural_mutate.h
b/include/tvm/ffi/extra/structural_mutate.h
index f9fceef1..afc1a628 100644
--- a/include/tvm/ffi/extra/structural_mutate.h
+++ b/include/tvm/ffi/extra/structural_mutate.h
@@ -474,6 +474,8 @@ class StructuralMutatorObj : public Object {
public:
/*! \brief Callback-facing mutator type used by composed callback-driven
engines. */
using MutatorObjType = StructuralMutatorObj;
+ /*! \brief Empty callback-state protocol used when no custom Parent layer is
present. */
+ using StateTupleType = std::tuple<>;
/*!
* \brief Throwing form of \ref MutateExpected.
@@ -778,6 +780,24 @@ class StructuralMutatorObj : public Object {
*/
explicit StructuralMutatorObj(const StructuralMutatorVTable* vtable) :
vtable_(vtable) {}
+ /*! \brief Return the empty state tuple exposed to typed map callbacks. */
+ TVM_FFI_INLINE StateTupleType StateTuple() const noexcept { return {}; }
+
+ /*!
+ * \brief Append \p node to a failed result's mutate error context.
+ * \param result The failed result whose Error is annotated.
+ * \param node The borrowed node to name in the context.
+ */
+ TVM_FFI_COLD_CODE static void UpdateVisitErrorContext(const Expected<Any>&
result,
+ AnyView node) noexcept
{
+ // The Error is refcounted, so annotating the local handle annotates the
object the result
+ // holds. A non-object node has no context to add.
+ if (node.type_index() >= TypeIndex::kTVMFFIStaticObjectBegin) {
+ Error err = result.error();
+ ::tvm::ffi::detail::UpdateVisitErrorContext(err, node.cast<ObjectRef>());
+ }
+ }
+
/*!
* \brief Non-owning pointer to the required ABI dispatch table.
*/
@@ -1034,23 +1054,28 @@ namespace detail {
} // namespace detail
+namespace detail {
+
/*!
- * \brief Shared state and variable-remap dispatch for both structural-map
engines.
+ * \brief Engine-owned variable-remap storage above a semantic mutator Parent.
+ *
+ * Typed and dynamic map/mutate engines share this identity-substitution
environment for
+ * FreeVar and DAG identities. The Parent holds only its own descent and
callback state;
+ * remap storage belongs to this complete engine layer. Its vtable thunks
downcast to this
+ * specialization so a custom Parent need not inherit a C++ container or its
destructor.
*
- * This base owns the identity-substitution environment used for FreeVar and
DAG identities.
- * Its variable-remap vtable thunks downcast to this common subobject,
allowing the typed and
- * dynamic engines to share the environment even when a Parent layer sits
above this base.
+ * \tparam Parent Semantic layer deriving from StructuralMutatorObj and
forwarding the vtable.
*/
-class StructuralMapEngineBase : public StructuralMutatorObj {
+template <typename Parent>
+class StructuralMutatorEngineBase : public Parent {
public:
- /*! \brief Empty callback-state protocol used when no custom Parent layer is
present. */
- using StateTupleType = std::tuple<>;
+ static_assert(std::is_base_of_v<StructuralMutatorObj, Parent>,
+ "Structural map/mutate Parent must derive from
StructuralMutatorObj");
/*! \brief Construct the shared engine base with the concrete engine's
vtable. */
- explicit StructuralMapEngineBase(const StructuralMutatorVTable* vtable)
- : StructuralMutatorObj(vtable) {}
+ explicit StructuralMutatorEngineBase(const StructuralMutatorVTable* vtable)
: Parent(vtable) {}
- ~StructuralMapEngineBase() {
+ ~StructuralMutatorEngineBase() {
for (const auto& kv : var_remap_) {
detail::ObjectUnsafe::DecRefObjectHandle(
reinterpret_cast<TVMFFIObjectHandle>(const_cast<Object*>(kv.first)));
@@ -1058,9 +1083,6 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
}
protected:
- /*! \brief Return the empty state tuple exposed to typed map callbacks. */
- TVM_FFI_INLINE StateTupleType StateTuple() const noexcept { return {}; }
-
/// \cond Doxygen_Suppress
// Out of line so its strings stay out of the per-node dispatch function,
which TryLink inlines
// into. Shared by the typed and dynamic engines below.
@@ -1080,7 +1102,7 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
* \return Raw ``TVMFFIAny`` containing the owning replacement, FFI None, or
Error.
*/
static TVMFFIAny DispatchVarRemapGet(StructuralMutatorObj* mutator, AnyView
var) noexcept {
- auto* self = static_cast<StructuralMapEngineBase*>(mutator);
+ auto* self = static_cast<StructuralMutatorEngineBase*>(mutator);
return detail::ExpectedUnsafe::MoveToTVMFFIAny(self->VarRemapGetImpl(var));
}
@@ -1093,25 +1115,10 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
*/
static TVMFFIAny DispatchVarRemapSet(StructuralMutatorObj* mutator, AnyView
var,
AnyView mapped_value) noexcept {
- auto* self = static_cast<StructuralMapEngineBase*>(mutator);
+ auto* self = static_cast<StructuralMutatorEngineBase*>(mutator);
return detail::ExpectedUnsafe::MoveToTVMFFIAny(self->VarRemapSetImpl(var,
mapped_value));
}
- /*!
- * \brief Append \p node to a failed result's mutate error context.
- * \param result The failed result whose Error is annotated.
- * \param node The borrowed node to name in the context.
- */
- TVM_FFI_COLD_CODE static void UpdateVisitErrorContext(const Expected<Any>&
result,
- AnyView node) noexcept
{
- // The Error is refcounted, so annotating the local handle annotates the
object the result
- // holds. A non-object node has no context to add.
- if (node.type_index() >= TypeIndex::kTVMFFIStaticObjectBegin) {
- Error err = result.error();
- ::tvm::ffi::detail::UpdateVisitErrorContext(err, node.cast<ObjectRef>());
- }
- }
-
/*!
* \brief Look up a replacement in the identity-substitution environment.
* \param var The borrowed variable identity to look up.
@@ -1149,20 +1156,29 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
return Expected<void>();
}
- template <typename Parent, WalkOrder order, typename... Callbacks>
- friend class StructuralMapEngine;
- template <typename Parent, WalkOrder order>
- friend class StructuralMapDynEngine;
- template <typename Parent, typename... Callbacks>
- friend class StructuralMutateEngine;
- template <typename Parent>
- friend class detail::StructuralMutateDynEngine;
-
/*! \brief Identity-substitution environment keyed by object identity. */
// Raw-pointer key: IncRef once on first insert, DecRef all keys in the
destructor.
std::unordered_map<const Object*, Any> var_remap_;
};
+// Existing semantic layers may already inherit the original remap-owning
base. Reuse that
+// complete layer rather than adding a second map; new parents receive storage
above their state.
+template <typename Parent>
+using StructuralMutatorEngineParent =
+
std::conditional_t<std::is_base_of_v<StructuralMutatorEngineBase<StructuralMutatorObj>,
Parent>,
+ Parent, StructuralMutatorEngineBase<Parent>>;
+
+} // namespace detail
+
+/*!
+ * \brief Default remap-owning base retained for existing composed C++
mutators.
+ *
+ * New semantic parents can derive directly from StructuralMutatorObj; the
map/mutate engine
+ * supplies remap storage above them. Existing parents that derive from this
base keep one
+ * remap environment and their original construction/destruction behavior.
+ */
+using StructuralMapEngineBase =
detail::StructuralMutatorEngineBase<StructuralMutatorObj>;
+
/*!
* \brief Callback-dispatched structural mutator with a state-carrying Parent
layer.
*
@@ -1170,7 +1186,7 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
* type, using a compile-time-known ``as<TSub>()`` on the input node in
pre-order and on the
* descended node in post-order.
*
- * ``Parent`` derives from ``StructuralMapEngineBase``, publishes
``StateTupleType``, accepts
+ * ``Parent`` derives from ``StructuralMutatorObj``, publishes
``StateTupleType``, accepts
* and forwards the mutator vtable in its constructor, and provides a protected
* ``StateTuple() const noexcept``. Engine-internal descent uses the public
Expected entry points,
* so a Parent may override those entries directly. Each callback receives
every tuple entry
@@ -1180,7 +1196,7 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
* not virtual dispatch. This deliberately leaves composed deeper-layer
declarations eligible;
* spelling the calls as ``Parent::member`` would instead pin lookup at that
qualified layer.
* A Parent must not hide other engine-internal ``this->`` members because
matching ABI-vtable
- * paths deliberately terminate at ``StructuralMapEngineBase``.
+ * remap paths deliberately terminate at the selected engine storage layer.
*
* \tparam Parent Mutator layer that supplies descent and callback state
through the complete
* protocol above.
@@ -1188,10 +1204,10 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
* \tparam Callbacks The callbacks, tested in declaration order.
*/
template <typename Parent, WalkOrder order, typename... Callbacks>
-class StructuralMapEngine : public Parent {
+class StructuralMapEngine : public
detail::StructuralMutatorEngineParent<Parent> {
+ using EngineBase = detail::StructuralMutatorEngineParent<Parent>;
+
public:
- static_assert(std::is_base_of_v<StructuralMapEngineBase, Parent>,
- "StructuralMap Parent must derive from
StructuralMapEngineBase");
/*! \brief Tuple of state references supplied by the Parent layer. */
using StateTupleType = typename Parent::StateTupleType;
@@ -1200,7 +1216,7 @@ class StructuralMapEngine : public Parent {
* \param callbacks The typed callback links, tested in declaration order.
*/
explicit StructuralMapEngine(Callbacks... callbacks)
- : Parent(VTable()), callbacks_(std::move(callbacks)...) {}
+ : EngineBase(VTable()), callbacks_(std::move(callbacks)...) {}
private:
using ExpectedUnsafe = detail::ExpectedUnsafe;
@@ -1443,10 +1459,10 @@ class StructuralMapEngine : public Parent {
* \tparam order Callback placement relative to child mapping.
*/
template <typename Parent, WalkOrder order>
-class StructuralMapDynEngine : public Parent {
+class StructuralMapDynEngine : public
detail::StructuralMutatorEngineParent<Parent> {
+ using EngineBase = detail::StructuralMutatorEngineParent<Parent>;
+
public:
- static_assert(std::is_base_of_v<StructuralMapEngineBase, Parent>,
- "StructuralMap Parent must derive from
StructuralMapEngineBase");
/*!
* \brief Construct a dynamic map engine with the default Parent constructor.
* \param callbacks Runtime links invoked as ``callback(value)``.
@@ -1454,7 +1470,7 @@ class StructuralMapDynEngine : public Parent {
*/
StructuralMapDynEngine(Array<Tuple<int32_t, Function>> callbacks,
Array<Tuple<int32_t, Function>>
callbacks_with_def_region_kind)
- : Parent(VTable()),
+ : EngineBase(VTable()),
callbacks_(std::move(callbacks)),
callbacks_with_def_region_kind_(std::move(callbacks_with_def_region_kind)) {}
@@ -1643,14 +1659,13 @@ class StructuralMapDynEngine : public Parent {
* \tparam Callbacks Callable types whose first parameter selects the value
type.
*/
template <typename Parent, typename... Callbacks>
-class StructuralMutateEngine : public Parent {
- public:
- static_assert(std::is_base_of_v<StructuralMapEngineBase, Parent>,
- "StructuralMutate Parent must derive from
StructuralMapEngineBase");
+class StructuralMutateEngine : public
detail::StructuralMutatorEngineParent<Parent> {
+ using EngineBase = detail::StructuralMutatorEngineParent<Parent>;
+ public:
/*! \brief Construct a mutate engine over callbacks tested in declaration
order. */
explicit StructuralMutateEngine(Callbacks... callbacks)
- : Parent(VTable()), callbacks_(std::move(callbacks)...) {}
+ : EngineBase(VTable()), callbacks_(std::move(callbacks)...) {}
private:
/*! \brief Return this engine's immutable callback-aware mutator vtable. */
@@ -1861,7 +1876,7 @@ template <WalkOrder order, typename... Callbacks>
Expected<Any> StructuralMapExpected(
Any root, Callbacks&&... callbacks) noexcept { //
NOLINT(performance-unnecessary-value-param)
static_assert(sizeof...(Callbacks) != 0, "StructuralMap requires at least
one callback");
- using Mutator = StructuralMapEngine<StructuralMapEngineBase, order,
std::decay_t<Callbacks>...>;
+ using Mutator = StructuralMapEngine<StructuralMutatorObj, order,
std::decay_t<Callbacks>...>;
StructuralMutator
mutator(make_object<Mutator>(std::forward<Callbacks>(callbacks)...));
auto result = mutator->MutateExpected(root, InplaceMode::kAllow);
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
@@ -1934,7 +1949,7 @@ template <typename... Callbacks>
Expected<Any> StructuralMutateExpected(
Any root, Callbacks&&... callbacks) noexcept { //
NOLINT(performance-unnecessary-value-param)
static_assert(sizeof...(Callbacks) != 0, "StructuralMutate requires at least
one callback");
- using Mutator = StructuralMutateEngine<StructuralMapEngineBase,
std::decay_t<Callbacks>...>;
+ using Mutator = StructuralMutateEngine<StructuralMutatorObj,
std::decay_t<Callbacks>...>;
StructuralMutator
mutator(make_object<Mutator>(std::forward<Callbacks>(callbacks)...));
auto result = mutator->MutateExpected(root, InplaceMode::kAllow);
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
diff --git a/rust/tvm-ffi/src/extra/structural_common.rs
b/rust/tvm-ffi/src/extra/structural_common.rs
index 1ebc228b..19fdaca0 100644
--- a/rust/tvm-ffi/src/extra/structural_common.rs
+++ b/rust/tvm-ffi/src/extra/structural_common.rs
@@ -235,7 +235,7 @@ pub(crate) use impl_callback_chain_tuple_arities;
/// Object-backed values compare pointer identity. Both carriers must contain
/// initialized payload bits; no pointee is accessed or ownership transferred.
#[inline(always)]
-pub(crate) fn raw_tvm_ffi_any_same_as(lhs: TVMFFIAny, rhs: TVMFFIAny) -> bool {
+pub(crate) fn raw_ffi_any_same_as(lhs: TVMFFIAny, rhs: TVMFFIAny) -> bool {
lhs.type_index == rhs.type_index
&& lhs.small_str_len == rhs.small_str_len
&& unsafe { lhs.data_union.v_uint64 == rhs.data_union.v_uint64 }
@@ -248,7 +248,7 @@ pub(crate) fn raw_tvm_ffi_any_same_as(lhs: TVMFFIAny, rhs:
TVMFFIAny) -> bool {
/// Error objects as ordinary values must keep a separate result channel.
/// No payload is cloned, normalized, or released by this transfer.
#[inline(always)]
-pub(crate) fn result_into_raw_tvm_ffi_any(result: Result<Any>) -> TVMFFIAny {
+pub(crate) fn result_into_raw_ffi_any(result: Result<Any>) -> TVMFFIAny {
unsafe {
match result {
Ok(value) => Any::into_raw_ffi_any(value),
@@ -266,7 +266,7 @@ pub(crate) fn result_into_raw_tvm_ffi_any(result:
Result<Any>) -> TVMFFIAny {
/// if it was encoded from `Ok`; all other tags become an owning success.
/// Borrowed views and rvalue slots must be normalized before this transfer.
#[inline(always)]
-pub(crate) unsafe fn raw_tvm_ffi_any_into_result(raw: TVMFFIAny) ->
Result<Any> {
+pub(crate) unsafe fn raw_ffi_any_into_result(raw: TVMFFIAny) -> Result<Any> {
if raw.type_index != TVMFFITypeIndex::kTVMFFIError as i32 {
return Ok(Any::from_raw_ffi_any(raw));
}
diff --git a/rust/tvm-ffi/src/extra/structural_map.rs
b/rust/tvm-ffi/src/extra/structural_map.rs
index 02c541ab..6ba5c8ca 100644
--- a/rust/tvm-ffi/src/extra/structural_map.rs
+++ b/rust/tvm-ffi/src/extra/structural_map.rs
@@ -19,7 +19,7 @@
//! Ordered replacement callbacks and adapters for structural mapping.
-use super::structural_common::raw_tvm_ffi_any_same_as;
+use super::structural_common::raw_ffi_any_same_as;
use super::structural_mutate::{
IntoMutateResult, MutationView, StructuralMutateEngine,
StructuralMutatorView,
with_mutation_error_context,
@@ -115,7 +115,7 @@ where
};
let mapped =
callback_result.map_err(|error|
with_mutation_error_context(error, raw))?;
- if is_unchanged(&mapped) || raw_tvm_ffi_any_same_as(raw,
*mapped.as_raw_ffi_any()) {
+ if is_unchanged(&mapped) || raw_ffi_any_same_as(raw,
*mapped.as_raw_ffi_any()) {
// Drop temporary ownership before default descent checks
uniqueness.
drop(mapped);
mutator.default_mutate(self, value).map(Into::into)
diff --git a/rust/tvm-ffi/src/extra/structural_mutate.rs
b/rust/tvm-ffi/src/extra/structural_mutate.rs
index 4d77a081..83440788 100644
--- a/rust/tvm-ffi/src/extra/structural_mutate.rs
+++ b/rust/tvm-ffi/src/extra/structural_mutate.rs
@@ -74,8 +74,8 @@ use crate::tvm_ffi_sys::{
use crate::tvm_ffi_sys::{TVMFFIObjectHandle, TVMFFISEqHashKind};
use super::structural_common::{
- missing_field_getter_error, raw_tvm_ffi_any_into_result,
raw_tvm_ffi_any_same_as,
- result_into_raw_tvm_ffi_any, runtime_error, type_key_of,
with_structural_error_context,
+ missing_field_getter_error, raw_ffi_any_into_result, raw_ffi_any_same_as,
+ result_into_raw_ffi_any, runtime_error, type_key_of,
with_structural_error_context,
with_structural_field_error_context, with_visit_error_context,
};
use super::structural_visit::{
@@ -685,7 +685,7 @@ impl<D, Base: StructuralMutatorBase>
StructuralMutatorView<'_, D, Base> {
unsafe {
let context =
StructuralMutateEngine::<Base>::checked_dispatcher(self.runtime, dispatch);
-
raw_tvm_ffi_any_into_result(StructuralMutateEngine::<Base>::default_mutate_raw(
+
raw_ffi_any_into_result(StructuralMutateEngine::<Base>::default_mutate_raw(
self.runtime,
context,
value.value,
@@ -855,7 +855,7 @@ pub(super) fn with_mutation_error_context(error: Error,
raw: TVMFFIAny) -> Error
//
-----------------------------------------------------------------------------
fn shallow_copy(raw: TVMFFIAny) -> Result<Any> {
- let Some(attr) = SHALLOW_COPY_COLUMN.get_raw(raw.type_index) else {
+ let Some(attr) = SHALLOW_COPY_COLUMN.get_raw_ffi_any(raw.type_index) else {
return Err(missing_shallow_copy_error(
raw.type_index,
SHALLOW_COPY_ATTR,
@@ -963,7 +963,7 @@ fn object_is_unique(raw: TVMFFIAny) -> bool {
/// uniquely own each writable rvalue slot, converting it at most once. A null
/// object pointer or failed runtime conversion returns an error; successful
Error
/// objects remain ordinary values rather than entering the error channel.
-fn raw_tvm_ffi_any_to_owned_any(raw: TVMFFIAny) -> Result<Any> {
+fn raw_ffi_any_to_owned_any(raw: TVMFFIAny) -> Result<Any> {
if is_plain_inline(raw.type_index) {
return Ok(unsafe { Any::from_raw_ffi_any(raw) });
}
@@ -1115,7 +1115,7 @@ impl<Base: StructuralMutatorBase>
StructuralMutateEngine<Base> {
) {
// Map adapters attribute errors to their possibly replaced value;
// direct mutation attributes callback failures to this input.
- Some(result) => result_into_raw_tvm_ffi_any(if CALLBACK_CONTEXT {
+ Some(result) => result_into_raw_ffi_any(if CALLBACK_CONTEXT {
result.map_err(|error| with_mutation_error_context(error, raw))
} else {
result
@@ -1194,7 +1194,7 @@ impl<Base: StructuralMutatorBase>
StructuralMutateEngine<Base> {
match Self::runtime_var_remap_get(runtime, *var.as_raw_ffi_any()) {
Ok(Some(value)) => Any::into_raw_ffi_any(value),
Ok(None) => TVMFFIAny::new(),
- Err(error) => result_into_raw_tvm_ffi_any(Err(error)),
+ Err(error) => result_into_raw_ffi_any(Err(error)),
}
}
@@ -1209,9 +1209,9 @@ impl<Base: StructuralMutatorBase>
StructuralMutateEngine<Base> {
replacement: AnyView<'static>,
) -> TVMFFIAny {
// Finish replacement cleanup before constructing the ABI result.
- let result =
raw_tvm_ffi_any_to_owned_any(*replacement.as_raw_ffi_any())
+ let result = raw_ffi_any_to_owned_any(*replacement.as_raw_ffi_any())
.and_then(|value| Self::runtime_var_remap_set(runtime,
*var.as_raw_ffi_any(), &value));
- result_into_raw_tvm_ffi_any(result.map(|()| Any::new()))
+ result_into_raw_ffi_any(result.map(|()| Any::new()))
}
/// Run the public mutate/map adapter with owned `root` and borrowed
`dispatch`.
@@ -1311,8 +1311,8 @@ impl StructuralMutatorObj {
small_str_len,
data_union,
};
- result_into_raw_tvm_ffi_any(
- raw_tvm_ffi_any_into_result(result)
+ result_into_raw_ffi_any(
+ raw_ffi_any_into_result(result)
.map_err(|error| with_mutation_error_context(error, raw)),
)
}
@@ -1426,7 +1426,7 @@ impl StructuralMutatorObj {
Self::call_mutator(runtime, child_raw, InplaceMode::Disallow)
})
.map_err(|error| with_structural_field_error_context(error, "map",
field.name.as_str()))?;
- if is_unchanged(&mapped) || raw_tvm_ffi_any_same_as(child_raw,
*mapped.as_raw_ffi_any()) {
+ if is_unchanged(&mapped) || raw_ffi_any_same_as(child_raw,
*mapped.as_raw_ffi_any()) {
return Ok(());
}
@@ -1456,7 +1456,7 @@ impl StructuralMutatorObj {
};
unsafe {
let view = AnyView::from_raw_ffi_any(raw);
- raw_tvm_ffi_any_into_result(callback(mutator, view))
+ raw_ffi_any_into_result(callback(mutator, view))
}
}
@@ -1467,7 +1467,7 @@ impl StructuralMutatorObj {
fn structural_mutate_hook(raw: TVMFFIAny, mode: InplaceMode) ->
Option<TVMFFIAny> {
let use_inplace = mode == InplaceMode::Allow && object_is_unique(raw);
if use_inplace {
- if let Some(attr) =
structural_mutate_column::<true>().get_raw(raw.type_index) {
+ if let Some(attr) =
structural_mutate_column::<true>().get_raw_ffi_any(raw.type_index) {
if attr.type_index == TVMFFITypeIndex::kTVMFFIOpaquePtr as i32
|| attr.type_index == TVMFFITypeIndex::kTVMFFIFunction as
i32
{
@@ -1475,7 +1475,7 @@ impl StructuralMutatorObj {
}
}
}
- let attr =
structural_mutate_column::<false>().get_raw(raw.type_index)?;
+ let attr =
structural_mutate_column::<false>().get_raw_ffi_any(raw.type_index)?;
(attr.type_index != TVMFFITypeIndex::kTVMFFINone as
i32).then_some(attr)
}
@@ -1558,7 +1558,7 @@ impl StructuralMutatorObj {
Err(invalid_mutation_hook_error())
};
// Preserve successful Error-valued returns until the public
conversion.
- result_into_raw_tvm_ffi_any(result.map_err(|error|
with_mutation_error_context(error, raw)))
+ result_into_raw_ffi_any(result.map_err(|error|
with_mutation_error_context(error, raw)))
}
/// Borrow a live runtime mutator as an object-valued ABI argument.
@@ -1709,7 +1709,7 @@ impl StructuralMutatorObj {
small_str_len,
data_union,
};
- result_into_raw_tvm_ffi_any(
+ result_into_raw_ffi_any(
Self::mutate_without_hook::<ORDINARY_ENGINE>(runtime, raw,
def_region_kind)
.map_err(|error| with_mutation_error_context(error, raw)),
)
@@ -1723,7 +1723,7 @@ impl StructuralMutatorObj {
raw: TVMFFIAny,
) -> Result<Option<Any>> {
let callback = (*(*runtime).vtable).var_remap_get;
- let value = raw_tvm_ffi_any_into_result(callback(runtime,
AnyView::from_raw_ffi_any(raw)))?;
+ let value = raw_ffi_any_into_result(callback(runtime,
AnyView::from_raw_ffi_any(raw)))?;
Ok((value.type_index() != TVMFFITypeIndex::kTVMFFINone as
i32).then_some(value))
}
@@ -1736,7 +1736,7 @@ impl StructuralMutatorObj {
value: &Any,
) -> Result<()> {
let callback = (*(*runtime).vtable).var_remap_set;
- raw_tvm_ffi_any_into_result(callback(
+ raw_ffi_any_into_result(callback(
runtime,
AnyView::from_raw_ffi_any(raw),
AnyView::from_raw_ffi_any(*value.as_raw_ffi_any()),
@@ -1755,7 +1755,7 @@ impl StructuralMutatorObj {
def_region_kind: DefRegionKind,
) -> Result<Any> {
if raw.type_index < TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32 {
- return raw_tvm_ffi_any_to_owned_any(raw);
+ return raw_ffi_any_to_owned_any(raw);
}
let type_info = checked_type_info(raw.type_index)?;
diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs
b/rust/tvm-ffi/src/extra/structural_visit.rs
index afde0d76..82d05f4e 100644
--- a/rust/tvm-ffi/src/extra/structural_visit.rs
+++ b/rust/tvm-ffi/src/extra/structural_visit.rs
@@ -695,7 +695,7 @@ impl StructuralVisitorObj {
/// owner survives descent. Hooks may not retain the runtime or input
borrow.
#[inline(always)]
unsafe fn visit_children_raw(runtime: StructuralVisitorHandle, value:
TVMFFIAny) -> TVMFFIAny {
- if let Some(attr) =
structural_visit_column().get_raw(value.type_index) {
+ if let Some(attr) =
structural_visit_column().get_raw_ffi_any(value.type_index) {
if attr.type_index != TVMFFITypeIndex::kTVMFFINone as i32 {
return Self::call_structural_visit_hook_raw(runtime, value,
attr);
}
diff --git a/rust/tvm-ffi/src/extra/unchanged.rs
b/rust/tvm-ffi/src/extra/unchanged.rs
index 631f7237..862fa6f2 100644
--- a/rust/tvm-ffi/src/extra/unchanged.rs
+++ b/rust/tvm-ffi/src/extra/unchanged.rs
@@ -21,7 +21,7 @@
use std::marker::PhantomData;
-use super::structural_common::{raw_tvm_ffi_any_same_as,
structural_error_from_raw};
+use super::structural_common::{raw_ffi_any_same_as, structural_error_from_raw};
use crate::any::{Any, AnyView, TryFromTemp};
use crate::error::{Error, Result, TYPE_ERROR};
@@ -90,7 +90,7 @@ impl<T: ContainerElement> UnchangedOr<T> {
unsafe {
T::container_copy_to_any_view(original, &mut raw);
}
- raw_tvm_ffi_any_same_as(raw, *self.data.as_raw_ffi_any())
+ raw_ffi_any_same_as(raw, *self.data.as_raw_ffi_any())
}
/// Move out the replacement, returning None when unchanged.
diff --git a/rust/tvm-ffi/src/reflection.rs b/rust/tvm-ffi/src/reflection.rs
index 019573d5..0de590f7 100644
--- a/rust/tvm-ffi/src/reflection.rs
+++ b/rust/tvm-ffi/src/reflection.rs
@@ -61,16 +61,26 @@ impl TypeAttrColumn {
/// Return an owning copy of this attribute for `type_index`.
pub fn get(self, type_index: i32) -> Option<Any> {
- let raw = self.get_raw(type_index)?;
+ let raw = self.get_raw_ffi_any(type_index)?;
if raw.type_index == TVMFFITypeIndex::kTVMFFINone as i32 {
return None;
}
Some(Any::from(unsafe { AnyView::from_raw_ffi_any(raw) }))
}
- /// Copy one borrowed cell without taking ownership.
+ /// Copy the borrowed raw carrier for this attribute and `type_index`.
+ ///
+ /// This does not retain the value or transfer ownership. The
registry-owned
+ /// value must remain valid while the returned carrier is used, and
registration
+ /// must not race with reads. Passing the carrier to unsafe APIs still
requires
+ /// satisfying their lifetime and ownership contracts. Use [`Self::get`]
for an
+ /// owning [`Any`].
+ ///
+ /// Returns `None` when the index is outside the column or the column has
no data.
+ /// An in-range cell with the `kTVMFFINone` tag is returned as `Some`,
unlike
+ /// [`Self::get`], which filters it out.
#[inline(always)]
- pub(crate) fn get_raw(self, type_index: i32) -> Option<TVMFFIAny> {
+ pub fn get_raw_ffi_any(self, type_index: i32) -> Option<TVMFFIAny> {
unsafe {
let column = self.0.as_ref();
let index = type_index - column.begin_index;
diff --git a/src/ffi/extra/structural_mutate.cc
b/src/ffi/extra/structural_mutate.cc
index f133d28f..72ffc6ad 100644
--- a/src/ffi/extra/structural_mutate.cc
+++ b/src/ffi/extra/structural_mutate.cc
@@ -55,11 +55,11 @@ Expected<Any> StructuralMapExpected(
const Array<Tuple<int32_t, Function>>& callbacks_with_def_region_kind, int
order) noexcept {
Expected<UnchangedOr<Any>> result = [&]() -> Expected<UnchangedOr<Any>> {
if (order == static_cast<int>(WalkOrder::kPreOrder)) {
- using Mutator = StructuralMapDynEngine<StructuralMapEngineBase,
WalkOrder::kPreOrder>;
+ using Mutator = StructuralMapDynEngine<StructuralMutatorObj,
WalkOrder::kPreOrder>;
StructuralMutator mutator(make_object<Mutator>(callbacks,
callbacks_with_def_region_kind));
return mutator->MutateExpected(root, InplaceMode::kAllow);
}
- using Mutator = StructuralMapDynEngine<StructuralMapEngineBase,
WalkOrder::kPostOrder>;
+ using Mutator = StructuralMapDynEngine<StructuralMutatorObj,
WalkOrder::kPostOrder>;
StructuralMutator mutator(make_object<Mutator>(callbacks,
callbacks_with_def_region_kind));
return mutator->MutateExpected(root, InplaceMode::kAllow);
}();
@@ -77,10 +77,12 @@ Expected<Any> StructuralMapExpected(
/*! \brief Runtime counterpart of the typed callback-owned mutate engine. */
template <typename Parent>
-class StructuralMutateDynEngine : public Parent {
+class StructuralMutateDynEngine : public StructuralMutatorEngineParent<Parent>
{
+ using EngineBase = StructuralMutatorEngineParent<Parent>;
+
public:
explicit StructuralMutateDynEngine(Array<Tuple<int32_t, Function, bool>>
callbacks)
- : Parent(VTable()), callbacks_(std::move(callbacks)) {}
+ : EngineBase(VTable()), callbacks_(std::move(callbacks)) {}
private:
static const StructuralMutatorVTable* VTable() {
@@ -153,7 +155,7 @@ class StructuralMutateDynEngine : public Parent {
Expected<Any> StructuralMutateExpected(
Any root, // NOLINT(performance-unnecessary-value-param)
const Array<Tuple<int32_t, Function, bool>>& callbacks) noexcept {
- using Mutator = StructuralMutateDynEngine<StructuralMapEngineBase>;
+ using Mutator = StructuralMutateDynEngine<StructuralMutatorObj>;
StructuralMutator mutator(make_object<Mutator>(callbacks));
auto result = mutator->MutateExpected(root, InplaceMode::kAllow);
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
diff --git a/tests/cpp/extra/test_structural_mutate.cc
b/tests/cpp/extra/test_structural_mutate.cc
index 1a25b87d..0acf0257 100644
--- a/tests/cpp/extra/test_structural_mutate.cc
+++ b/tests/cpp/extra/test_structural_mutate.cc
@@ -333,12 +333,12 @@ struct MutateCount {
int maybe_inplace_expected = 0;
};
-class StructuralMapWithMutateCount : public StructuralMapEngineBase {
+class StructuralMapWithMutateCount : public StructuralMutatorObj {
public:
using StateTupleType = std::tuple<const MutateCount&, const int&>;
explicit StructuralMapWithMutateCount(const StructuralMutatorVTable* vtable)
- : StructuralMapEngineBase(vtable) {}
+ : StructuralMutatorObj(vtable) {}
const MutateCount& count() const { return count_; }
@@ -350,7 +350,7 @@ class StructuralMapWithMutateCount : public
StructuralMapEngineBase {
} else {
++count_.mutate_expected;
}
- return StructuralMapEngineBase::DefaultMutateExpected(value, inplace_mode);
+ return StructuralMutatorObj::DefaultMutateExpected(value, inplace_mode);
}
protected:
@@ -361,16 +361,65 @@ class StructuralMapWithMutateCount : public
StructuralMapEngineBase {
int marker_ = 17;
};
-class StructuralMutateLayer : public StructuralMapEngineBase {
+class StructuralMutateLayer : public StructuralMutatorObj {
public:
using MutatorObjType = StructuralMutateLayer;
explicit StructuralMutateLayer(const StructuralMutatorVTable* vtable)
- : StructuralMapEngineBase(vtable) {}
+ : StructuralMutatorObj(vtable) {}
- int callback_tag() const { return 23; }
+ int callback_tag() const { return callback_tag_; }
+
+ private:
+ int callback_tag_ = 23;
};
+TEST(StructuralMutate, EngineOwnsRemapAboveParentState) {
+ auto check = [](auto make_engine) {
+ TVar key("key");
+ TVar first("first");
+ TVar second("second");
+ {
+ auto engine = make_engine();
+ auto* mutator = static_cast<StructuralMutatorObj*>(engine.get());
+ EXPECT_EQ(mutator->VarRemapGetExpected(key).value(), nullptr);
+ mutator->VarRemapSetExpected(key, first).value();
+ EXPECT_EQ(key.use_count(), 2);
+ EXPECT_EQ(first.use_count(), 2);
+ EXPECT_TRUE(mutator->VarRemapGetExpected(key).value().same_as(first));
+ mutator->VarRemapSetExpected(key, second).value();
+ EXPECT_EQ(key.use_count(), 2);
+ EXPECT_EQ(first.use_count(), 1);
+ EXPECT_EQ(second.use_count(), 2);
+ EXPECT_TRUE(mutator->VarRemapGetExpected(key).value().same_as(second));
+ EXPECT_EQ(mutator->VarRemapSetExpected(1, second).error().kind(),
"TypeError");
+ }
+ EXPECT_EQ(key.use_count(), 1);
+ EXPECT_EQ(first.use_count(), 1);
+ EXPECT_EQ(second.use_count(), 1);
+ };
+ check([] {
+ auto callback = [](int64_t, StructuralMutatorObj*) -> Expected<Any> {
return Unchanged(); };
+ return make_object<StructuralMutateEngine<StructuralMapEngineBase,
decltype(callback)>>(
+ callback);
+ });
+ check([] {
+ auto callback = [](int64_t, StructuralMutateLayer*) -> Expected<Any> {
return Unchanged(); };
+ return make_object<StructuralMutateEngine<StructuralMutateLayer,
decltype(callback)>>(callback);
+ });
+ check([] {
+ auto callback = [](int64_t, const MutateCount&, const int&) ->
Expected<Any> {
+ return Unchanged();
+ };
+ return make_object<StructuralMapEngine<StructuralMapWithMutateCount,
WalkOrder::kPostOrder,
+ decltype(callback)>>(callback);
+ });
+ check([] {
+ return make_object<StructuralMapDynEngine<StructuralMapWithMutateCount,
WalkOrder::kPreOrder>>(
+ Array<Tuple<int32_t, Function>>(), Array<Tuple<int32_t, Function>>());
+ });
+}
+
TEST(StructuralMap, ParentLayerOwnsBothDescentsAndProvidesState) {
std::vector<int> callback_counts;
auto identity = [&](const AnyArray& value, const MutateCount& live_count,
const int& live_marker,
