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 ca8803c9 [FEAT] Add structural map, maybe_inplace_mutate, and
var_remap APIs (#649)
ca8803c9 is described below
commit ca8803c9f34d26cdc665f5dc5adfc0bc816771eb
Author: Kathryn (Jinqi) Chen <[email protected]>
AuthorDate: Thu Aug 13 20:03:29 2026 -0700
[FEAT] Add structural map, maybe_inplace_mutate, and var_remap APIs (#649)
## Summary
This PR adds structural transformation support to TVM FFI, complementing
the existing structural equality, hashing, and walking APIs.
The new `StructuralMutator` recursively transforms reflected object
graphs, while `structural_map` provides a convenient callback-based
interface for compiler passes.
## Structural mapping
`structural_map` follows the same typed callback model as
`structural_walk`. Callbacks are matched by runtime type and return
either the unchanged value or a replacement.
Mapping defaults to post-order, so callbacks observe values whose
children have already been transformed. This makes bottom-up rewrites
such as constant folding straightforward:
```py
def fold_add(expr):
if isinstance(expr.lhs, IntImm) and isinstance(expr.rhs, IntImm):
return IntImm(expr.lhs.value + expr.rhs.value)
return expr
optimized = tvm_ffi.structural_map(function, (Add, fold_add))
```
The PR exposes:
- structural_map in Python.
- StructuralMap and StructuralMapExpected in C++.
- Ordered typed callbacks and grouped callback types.
- Pre-order and post-order transformation.
- Callbacks receiving either value or (value, def_region_kind).
- Structural error context when transformation fails.
## StructuralMutator
`StructuralMutator` provides the low-level transformation engine through
two operations:
- `Mutate` transforms a value without intentionally modifying the input.
- `MaybeInplaceMutate` permits implementations to reuse an object when
doing so is safe.
## Custom mutation behavior
Object types can customize mutation using:
- `__s_mutate__` for canonical non-in-place transformation.
- `__s_maybe_inplace_mutate__` for an optional type-specific in-place
optimization.
The maybe-in-place hook owns its safety policy and may reuse the input,
delegate to normal mutation, or return another value. A type providing
it must also provide `__s_mutate__`.
Built-in hooks are registered for Array, List, Map, and Dict.
## Variable identity remapping
The mutator maintains an identity-substitution environment for FreeVar
objects. Once an identity is mapped, later occurrences reuse the same
result.
For example:
```py
def specialize_shape_var(var):
if var.name == "n":
return IntImm(10)
return var
specialized = tvm_ffi.structural_map(
function,
(Var, specialize_shape_var),
)
```
If the same `Var` occurs in a function parameter and its body, both
occurrences resolve to the same mapped result.
The mutator also exposes `get_var_remap` and `set_var_remap`, allowing
custom DAG-style variable wrappers to use their underlying identity
object as the remapping key.
---
CMakeLists.txt | 1 +
docs/concepts/structural_eq_hash.rst | 362 +++++++----
docs/guides/rust_lang_guide.md | 4 +
docs/reference/python/index.rst | 3 +
include/tvm/ffi/container/dict.h | 9 +
include/tvm/ffi/container/map.h | 9 +
include/tvm/ffi/extra/structural_mutate.h | 919 ++++++++++++++++++++++++++++
include/tvm/ffi/extra/structural_visit.h | 89 +--
include/tvm/ffi/reflection/accessor.h | 44 ++
python/tvm_ffi/__init__.py | 4 +
python/tvm_ffi/_ffi_api.py | 16 +-
python/tvm_ffi/dataclasses/py_class.py | 2 +
python/tvm_ffi/structural.py | 305 +++++++--
rust/tvm-ffi/src/extra/structural_visit.rs | 50 +-
rust/tvm-ffi/tests/test_structural_visit.rs | 10 +-
src/ffi/extra/structural_mutate.cc | 378 ++++++++++++
src/ffi/extra/structural_visit.cc | 7 +-
tests/cpp/extra/test_structural_mutate.cc | 259 ++++++++
tests/cpp/extra/test_structural_visit.cc | 7 +-
tests/python/test_structural.py | 223 ++++++-
20 files changed, 2433 insertions(+), 268 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 96f214c2..2e76a758 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -77,6 +77,7 @@ set(_tvm_ffi_extra_objs_sources
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_equal.cc"
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_hash.cc"
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_visit.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_mutate.cc"
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/visit_error_context.cc"
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_parser.cc"
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_writer.cc"
diff --git a/docs/concepts/structural_eq_hash.rst
b/docs/concepts/structural_eq_hash.rst
index f3acc3a7..6b998918 100644
--- a/docs/concepts/structural_eq_hash.rst
+++ b/docs/concepts/structural_eq_hash.rst
@@ -15,13 +15,18 @@
specific language governing permissions and limitations
under the License.
-Structural Equality, Hashing, and Walk
-======================================
+Structural Equality, Hashing, Walking, and Mapping
+==================================================
TVM FFI provides ``structural_equal`` and ``structural_hash`` for the
object graph. These compare objects by **content** — recursively walking
fields — rather than by pointer identity.
+The same reflection metadata also drives ``structural_walk`` for analyses and
+``structural_map`` for rewrites. Their low-level engines,
+``StructuralVisitor`` and ``StructuralMutator``, let custom object types
+participate in the same traversal protocol.
+
The behavior is controlled by two layers of annotation on
:func:`~tvm_ffi.dataclasses.py_class`:
@@ -971,185 +976,294 @@ And in Python:
assert structural_hash(f1) == structural_hash(f2) # same hash
-Structural Walk and Visit
--------------------------
+Structural Walk and Map
+------------------------------
-``structural_equal`` and ``structural_hash`` are built on a structural
traversal
-of the value graph. ``structural_walk`` exposes that traversal directly: it
-visits containers, object fields, and POD leaves, and invokes user callbacks
for
-values whose runtime type matches a callback entry.
+Structural walk and map use the same type metadata, field flags, and
+container registrations as structural equality and hashing. The default
+reflected traversal visits only structural fields, skips fields marked
+``structural_eq="ignore"``, and preserves definition-region information from
+fields marked as definitions.
-It is useful when you want to collect information, validate a tree, find a
node, or
-stop traversal early without writing a custom equality/hash hook.
+``Map`` and ``Dict`` keys are structural anchors. Both APIs recurse through
+container values only: walk callbacks do not observe keys, and map callbacks do
+not replace them. The map or dict object itself still participates in callback
+dispatch normally.
-Basic Walk
-~~~~~~~~~~
+.. code-block:: python
+
+ import tvm_ffi
+
+ table = tvm_ffi.Map({1: 2})
+ visited = []
+
+ tvm_ffi.structural_walk(table, (int, visited.append))
+ assert visited == [2]
+
+ mapped = tvm_ffi.structural_map(table, (int, lambda value: value + 10))
+ assert mapped[1] == 12
+ assert 11 not in mapped
+
+There are two layers of API:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 24 38 38
+
+ * - API
+ - Purpose
+ - Typical use
+ * - :func:`~tvm_ffi.structural_walk`
+ - Inspect a value graph without replacing values
+ - Collect information, validate IR, or stop at a match
+ * - :func:`~tvm_ffi.structural_map`
+ - Recursively replace values and rebuild changed paths
+ - Rewriting and compiler optimization passes
+ * - :class:`~tvm_ffi.StructuralVisitor`
+ - Low-level recursive visit engine
+ - Implementing ``__s_visit__`` or a language binding
+ * - :class:`~tvm_ffi.StructuralMutator`
+ - Low-level recursive mutation engine
+ - Implementing custom mutation hooks and identity substitution
+
+``structural_walk`` and ``structural_map`` construct the corresponding
low-level
+object, install callback-aware dispatch, run it on the root, and return the
final
+result. Applications normally use these two functions directly. Custom object
+hooks receive the low-level visitor or mutator so that recursive calls remain
in
+the same traversal.
+
+StructuralVisitor and StructuralMutator
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+A :class:`~tvm_ffi.StructuralVisitor` carries recursive dispatch, the current
+definition-region kind, and any early-interruption state. Its main operations
+are:
+
+- ``visitor.visit(value)`` visits a child with the same visitor.
+- ``visitor.def_region_kind()`` reports the active definition-region kind.
+- ``visitor.with_def_region_kind(kind, callback)`` temporarily changes that
kind
+ while ``callback`` performs recursive visits.
+
+The default visitor dispatches to a type's ``__s_visit__`` hook when present.
+Otherwise POD values are leaves and object-backed values are visited through
+their reflected structural fields. Array and List have built-in hooks that
+visit their elements; Map and Dict hooks visit values while skipping keys.
+
+A :class:`~tvm_ffi.StructuralMutator` adds ownership and replacement semantics.
+Its main operations are:
+
+- ``mutator.mutate(value)`` maps without intentionally modifying ``value``.
+- ``mutator.maybe_inplace_mutate(value)`` permits a type-specific
implementation
+ to reuse a safely mutable value and otherwise falls back to ``mutate``.
+- ``mutator.var_remap_get(var)`` and ``mutator.var_remap_set(var, mapped)``
+ access the current identity-substitution environment.
+- ``def_region_kind`` and ``with_def_region_kind`` have the same role as on the
+ visitor.
+
+String and Bytes are returned unchanged by default and are never mutated in
+place. For reflected objects, ``mutate`` starts from a shallow copy,
+recursively maps each structural field, and installs mapped fields in that
copy.
+If no field changes, it returns the original object instead. A nested change
+therefore copies only the objects along the changed path; unchanged children
+remain shared.
+``maybe_inplace_mutate`` is an explicit optimization path. A type-specific
+``__s_maybe_inplace_mutate__`` hook owns the safety policy and may reuse its
+input. Without that hook, the default implementation calls ``mutate``.
+
+.. note::
+
+ Visitor and mutator instances are supplied by an active traversal. Python
+ code normally receives them as arguments to ``__s_visit__``,
``__s_mutate__``,
+ or ``__s_maybe_inplace_mutate__`` rather than constructing them directly.
+
+Structural Walk
+~~~~~~~~~~~~~~~
+
+:func:`~tvm_ffi.structural_walk` invokes an analysis callback at each matching
+value. Callback entries are ordered, and only the first matching entry runs.
+A walk is post-order by default.
+A callback may return:
+
+- :attr:`~tvm_ffi.WalkResult.ADVANCE` or ``None`` to continue.
+- :attr:`~tvm_ffi.WalkResult.SKIP` to skip the current value's children. This
is
+ primarily useful with pre-order traversal.
+- :class:`~tvm_ffi.VisitInterrupt` to stop the entire traversal and return a
+ payload.
-Pass callbacks as ordered ``(type, callback)`` entries. The first matching
-entry runs for each visited value. Normal Python callbacks receive one
-argument, ``value``.
+For example, the following analysis records integer leaves and stops at the
+first negative value:
.. code-block:: python
import tvm_ffi
- visited = []
+ integers = []
- def on_int(value):
- visited.append(value)
- if value == 0:
+ def visit_int(value):
+ integers.append(value)
+ if value < 0:
return tvm_ffi.VisitInterrupt(value)
return tvm_ffi.WalkResult.ADVANCE
- result = tvm_ffi.structural_walk(root, (int, on_int), order="pre")
+ interrupted = tvm_ffi.structural_walk(
+ function,
+ (int, visit_int),
+ )
- if result is not None:
- print("stopped at", result.value)
+ if interrupted is not None:
+ print("first negative value:", interrupted.value)
-Callbacks may return:
+Walking never replaces a value. Side effects should be limited to the
+analysis state owned by the callback.
-- ``WalkResult.ADVANCE`` to continue into children.
-- ``WalkResult.SKIP`` to skip the current value's children.
-- ``VisitInterrupt(payload)`` to stop the entire walk and return an interrupt
- carrying ``payload``.
-- ``None`` as shorthand for ``WalkResult.ADVANCE``.
+Structural Map
+~~~~~~~~~~~~~~
-Grouped Types
-~~~~~~~~~~~~~
+:func:`~tvm_ffi.structural_map` uses the same typed callback selection but each
+callback returns the mapped value: either its input unchanged or a replacement.
+Mapping always visits all structural children; there is no ``SKIP`` or
+``VisitInterrupt`` result. A callback exception aborts the mapping and
+is propagated with structural visit context. Mapping is post-order by default,
+so a callback receives a value whose children have already been mapped.
-Several types can share one callback by passing a tuple of types:
+Post-order is natural for bottom-up compiler rewrites because children have
+already been mapped when the callback runs:
.. code-block:: python
- numbers = []
- strings = []
+ def fold_add(add):
+ if isinstance(add.lhs, IntImm) and isinstance(add.rhs, IntImm):
+ return IntImm(add.lhs.value + add.rhs.value)
+ return add
- tvm_ffi.structural_walk(
- root,
- [
- ((int, float), lambda value: numbers.append(value)),
- (str, lambda value: strings.append(value)),
- ],
+ optimized = tvm_ffi.structural_map(
+ function,
+ (Add, fold_add),
)
-This is normalized as if the same callback had been registered separately for
-``int`` and ``float``. Callback entries are still tried in order, so broad
-callbacks should usually come after more specific ones.
-
-Catch-All and Object Callbacks
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Map callbacks must follow map semantics: they must not mutate their input in
+place. The surrounding traversal may still reuse storage through an explicit
+``__s_maybe_inplace_mutate__`` hook. In pre-order, an unchanged or uniquely
+owned callback result may continue through that hook; in post-order, optional
+in-place mutation happens before the callback runs.
-``object`` and ``typing.Any`` are catch-all callbacks. They match POD leaves
-and object-backed values.
+Callback Selection and Order
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-.. code-block:: python
+Both Python functions accept the same callback forms:
- from typing import Any
+- ``(Type, callback)`` for one type.
+- ``((TypeA, TypeB), callback)`` to share one callback across types.
+- A sequence of callback entries.
+- A bare callable as a ``typing.Any`` catch-all.
- seen = []
- tvm_ffi.structural_walk(root, (Any, lambda value: seen.append(value)))
+``typing.Any`` and ``object`` match both POD and object-backed values.
+``tvm_ffi.Object`` matches only object-backed FFI values. Entries are tested
in
+the order supplied, so place specific types before broad catch-all callbacks.
-``tvm_ffi.Object`` is different: it matches only object-backed FFI values, such
-as ``Array``, ``Map``, ``Function``, ``String`` objects, or registered object
-classes. It does not match POD leaves such as ``int`` or ``float``.
+Both APIs default to post-order. The ``order`` argument controls the
+relationship between callbacks and children:
-.. code-block:: python
+- In pre-order, a walk callback runs before the children. For mapping, the
+ callback result becomes the value whose children are subsequently mapped.
+- In post-order, children are processed first. A map callback therefore
receives
+ the value with its mapped children already installed.
- objects = []
- leaves = []
+Definition Regions
+~~~~~~~~~~~~~~~~~~
- tvm_ffi.structural_walk(
- root,
- [
- (tvm_ffi.Object, lambda value: objects.append(value)),
- (object, lambda value: leaves.append(value)),
- ],
- )
+Callbacks passed through ``with_def_region_kind`` receive
+``(value, def_region_kind)``. The kind is one of:
-Def-Region Aware Walk
-~~~~~~~~~~~~~~~~~~~~~
+- ``DefRegionKind.NONE`` for an ordinary use.
+- ``DefRegionKind.DEF_RECURSIVE`` for a recursive definition region.
+- ``DefRegionKind.DEF_NON_RECURSIVE`` for a non-recursive definition.
-Callbacks passed to ``with_def_region_kind`` receive a second argument that
-reports whether the current value is visited as a definition or a use. This is
-useful for analyses such as collecting variable uses while skipping binders:
+The field annotations described earlier in this document establish these
+regions. Recursive definitions propagate the definition mode into the defined
+value's children. A non-recursive definition applies to the FreeVar identity
+itself, while its unannotated children are treated as ordinary uses. Explicit
+nested definition annotations establish their own region.
.. code-block:: python
uses = []
tvm_ffi.structural_walk(
- func,
+ function,
with_def_region_kind=(
Var,
lambda var, kind: (
- uses.append(var) if kind == tvm_ffi.DefRegionKind.NONE else None
+ uses.append(var)
+ if kind == tvm_ffi.DefRegionKind.NONE
+ else None
),
),
)
-For a function node, parameters are visited in a definition region, while
-occurrences in the body are visited with ``DefRegionKind.NONE``.
+``structural_map`` accepts the same def-region-aware callback form, but the
+callback must return the mapped value.
-Traversal Order
-~~~~~~~~~~~~~~~
+Custom Visit and Mutation Hooks
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-The default order is pre-order: callbacks run before visiting children.
-Post-order callbacks run after children.
+A type with non-standard child storage can define ``__s_visit__``. The hook
+receives the active visitor and the current value, recursively visits every
+structural child, and returns an interrupt if one occurs:
.. code-block:: python
- trace = []
-
- tvm_ffi.structural_walk(
- tvm_ffi.Array([tvm_ffi.Array([1]), 2]),
- [
- (tvm_ffi.Array, lambda value: trace.append(f"array:{len(value)}")),
- (int, lambda value: trace.append(f"int:{value}")),
- ],
- order="post",
- )
-
- assert trace == ["int:1", "array:1", "int:2", "array:2"]
-
-C++ Walk
+ @staticmethod
+ def __s_visit__(visitor, value):
+ return visitor.visit(value.children)
+
+A custom ``__s_mutate__`` hook similarly receives the active mutator. It
should
+recursively call ``mutator.mutate`` and return a new value only when needed.
+An optional ``__s_maybe_inplace_mutate__`` hook may implement an in-place
+optimization. Callers use ``mutate`` for shared objects and call
+``maybe_inplace_mutate`` only when the input is safe to mutate, so the optional
+hook may rely on that ownership guarantee. A type defining it must also define
+``__s_mutate__``. If the optional hook is absent, ``maybe_inplace_mutate``
uses
+the default non-in-place mutation; generic reflected fields are never mutated
+in place automatically.
+
+When an object marked ``structural_eq="var"`` or ``structural_eq="dag"``
registers
+either ``__s_mutate__`` or ``__s_maybe_inplace_mutate__`` hooks, it should:
+
+1. call ``var_remap_get`` before recursively mutating the value;
+2. immediately return the mapped value on a hit;
+3. compute the final result on a miss; and
+4. call ``var_remap_set`` with that final result before returning it.
+
+Structural-map callbacks apply the same identity rule automatically: the
complete
+callback/default result is recorded on the first occurrence and reused for
later
+occurrences of the same FreeVar or DAG node.
+
+C++ APIs
~~~~~~~~
-C++ code can use ``StructuralWalk`` with typed callbacks. Callbacks are tried
-in order and dispatch on the first argument type:
+C++ provides typed counterparts. Callback dispatch uses the first argument
+type and accepts an optional second ``TVMFFIDefRegionKind`` argument. The
+``Expected`` forms report failures without throwing:
.. code-block:: cpp
- Optional<VisitInterrupt> result = StructuralWalk<WalkOrder::kPreOrder>(
+ Expected<Optional<VisitInterrupt>> walked =
+ StructuralWalkExpected<WalkOrder::kPreOrder>(
+ root,
+ [&](const Add& add) -> Expected<WalkResult> {
+ ++num_adds;
+ return WalkResult::Advance();
+ });
+
+ Expected<Any> mapped = StructuralMapExpected<WalkOrder::kPostOrder>(
root,
- [&](const Add& add) -> Expected<WalkResult> {
- ++num_adds;
- return WalkResult::Advance();
- },
- [&](const Mul& mul) -> Expected<WalkResult> {
- return WalkResult::Skip();
+ [&](const Add& add) -> Expected<Any> {
+ return FoldAdd(add);
});
-C++ callbacks dispatch on their first argument, which may be ``AnyView``,
-``Any``, an ``ObjectRef`` subclass, an ``Object`` pointer type, or another
-FFI-convertible POD type. They may also take an optional second
-``TVMFFIDefRegionKind`` argument to distinguish definition sites from uses.
-Errors should be returned as ``Expected<WalkResult>``.
-
-Low-Level ``StructuralVisitor``
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-``StructuralVisitor`` is the lower-level traversal object. It is mainly useful
-inside structural visit hooks or C++ integrations that need to participate in
-the same recursive traversal protocol.
-
-Python users normally call ``structural_walk`` instead. The low-level visitor
-API exposes:
-
-- ``visitor.visit(value)`` to recursively visit a child value.
-- ``visitor.def_region_kind()`` to inspect the current definition-region mode.
-- ``visitor.with_def_region_kind(kind, callback)`` to run a recursive visit
- under a temporary definition-region mode.
-
-Custom visit hooks are registered as the ``__s_visit__`` type attribute. They
-receive the active visitor and the current object, and are responsible for
-calling ``visitor.visit(child)`` on structural children.
+Walk callbacks return ``Expected<WalkResult>``. Map callbacks return
+``Expected<Any>`` and must obey the same non-in-place callback contract as the
+Python API. For ``Map`` and ``Dict``, both APIs process values and skip keys.
+``StructuralWalk`` and ``StructuralMap`` are the corresponding throwing
+convenience forms.
diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md
index 226b5e31..717f1205 100644
--- a/docs/guides/rust_lang_guide.md
+++ b/docs/guides/rust_lang_guide.md
@@ -241,6 +241,10 @@ structural_walk(&values, &mut probe, WalkOrder::PreOrder)?;
assert_eq!(probe.total, 6);
```
+For `Map` and `Dict`, structural walk treats keys as structural anchors: it
+visits container values but does not pass keys to handlers. The map or dict
+object itself is still visited normally.
+
Lambdas also work — pass a single typed lambda, or a tuple of them (up to 8)
tried in order with the first matching argument type winning, like the
variadic C++ `StructuralWalk(root, callbacks...)` chain. Unmatched values
diff --git a/docs/reference/python/index.rst b/docs/reference/python/index.rst
index b5227f94..80a66725 100644
--- a/docs/reference/python/index.rst
+++ b/docs/reference/python/index.rst
@@ -83,13 +83,16 @@ Structural
StructuralKey
StructuralVisitor
+ StructuralMutator
VisitInterrupt
+ DefRegionKind
WalkOrder
WalkResult
get_first_structural_mismatch
structural_equal
structural_hash
structural_walk
+ structural_map
Global Registry
diff --git a/include/tvm/ffi/container/dict.h b/include/tvm/ffi/container/dict.h
index e75ea2cd..32751a09 100644
--- a/include/tvm/ffi/container/dict.h
+++ b/include/tvm/ffi/container/dict.h
@@ -41,6 +41,15 @@ namespace ffi {
/*! \brief Dict object — mutable map with shared reference semantics. */
class DictObj : public MapBaseObj {
public:
+ /*!
+ * \brief Create a shallow copy with the same iteration order.
+ * \param src The source dictionary object.
+ * \return The copied underlying dictionary storage.
+ */
+ static ObjectPtr<Object> ShallowCopy(const DictObj* src) {
+ return MapBaseObj::CopyFrom<DictObj>(const_cast<DictObj*>(src));
+ }
+
/// \cond Doxygen_Suppress
static constexpr const int32_t _type_index = TypeIndex::kTVMFFIDict;
static const constexpr bool _type_final = true;
diff --git a/include/tvm/ffi/container/map.h b/include/tvm/ffi/container/map.h
index 10013be9..5c0807ca 100644
--- a/include/tvm/ffi/container/map.h
+++ b/include/tvm/ffi/container/map.h
@@ -39,6 +39,15 @@ namespace ffi {
/*! \brief Map object */
class MapObj : public MapBaseObj {
public:
+ /*!
+ * \brief Create a shallow copy with the same iteration order.
+ * \param src The source map object.
+ * \return The copied underlying map storage.
+ */
+ static ObjectPtr<Object> ShallowCopy(const MapObj* src) {
+ return MapBaseObj::CopyFrom<MapObj>(const_cast<MapObj*>(src));
+ }
+
/// \cond Doxygen_Suppress
static constexpr const int32_t _type_index = TypeIndex::kTVMFFIMap;
static const constexpr bool _type_final = true;
diff --git a/include/tvm/ffi/extra/structural_mutate.h
b/include/tvm/ffi/extra/structural_mutate.h
new file mode 100644
index 00000000..908bb5e8
--- /dev/null
+++ b/include/tvm/ffi/extra/structural_mutate.h
@@ -0,0 +1,919 @@
+/*
+ * 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.
+ */
+/*!
+ * \file tvm/ffi/extra/structural_mutate.h
+ * \brief Structural mutation API with optional in-place optimization.
+ */
+#ifndef TVM_FFI_EXTRA_STRUCTURAL_MUTATE_H_
+#define TVM_FFI_EXTRA_STRUCTURAL_MUTATE_H_
+
+#include <tvm/ffi/any.h>
+#include <tvm/ffi/c_api.h>
+#include <tvm/ffi/cast.h>
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/container/tuple.h>
+#include <tvm/ffi/container/variant.h>
+#include <tvm/ffi/expected.h>
+#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ffi/extra/visit_error_context.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/function_details.h>
+#include <tvm/ffi/optional.h>
+#include <tvm/ffi/reflection/accessor.h>
+
+#include <cstddef>
+#include <exception>
+#include <optional>
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <utility>
+
+namespace tvm {
+namespace ffi {
+
+class StructuralMutatorObj;
+
+/*!
+ * \brief ABI callback type for structural mutation.
+ *
+ * \param mutator The active structural mutator.
+ * \param value The borrowed value to mutate.
+ * \return Raw ``TVMFFIAny`` containing the mutated value or an Error.
+ */
+using FStructuralMutate = TVMFFIAny (*)(StructuralMutatorObj* mutator, AnyView
value) noexcept;
+
+/*!
+ * \brief ABI callback type for looking up an identity substitution.
+ *
+ * \param mutator The active structural mutator.
+ * \param var The borrowed variable identity to look up.
+ * \return Raw ``TVMFFIAny`` containing the owning mapped value, FFI None when
no mapping exists,
+ * or an Error.
+ */
+using FStructuralVarRemapGet = TVMFFIAny (*)(StructuralMutatorObj* mutator,
AnyView var) noexcept;
+
+/*!
+ * \brief ABI callback type for recording an identity substitution.
+ *
+ * \param mutator The active structural mutator.
+ * \param var The borrowed variable identity to bind.
+ * \param mapped_value The borrowed replacement value.
+ * \return Raw ``TVMFFIAny`` containing FFI None on success or an Error.
+ */
+using FStructuralVarRemapSet = TVMFFIAny (*)(StructuralMutatorObj* mutator,
AnyView var,
+ AnyView mapped_value) noexcept;
+
+namespace details {
+
+// Copy and structurally mutate the reflected fields of an object-backed value.
+TVM_FFI_INLINE static Expected<Any>
MutateReflectedFieldsExpected(StructuralMutatorObj* mutator,
+ AnyView
value) noexcept;
+
+} // namespace details
+
+/*!
+ * \brief VTable ABI for \ref StructuralMutator dispatch.
+ */
+struct StructuralMutatorVTable {
+ /*!
+ * \brief Mutate a value without modifying the source in place.
+ *
+ * \param mutator The active structural mutator.
+ * \param value The borrowed value to mutate.
+ * \return Raw ``TVMFFIAny`` carrying the mutated value or Error.
+ */
+ FStructuralMutate mutate = nullptr;
+ /*!
+ * \brief Mutate a value, permitting an in-place implementation when it is
safe.
+ *
+ * \param mutator The active structural mutator.
+ * \param value The borrowed value to mutate.
+ * \return Raw ``TVMFFIAny`` carrying the mutated value or Error.
+ *
+ * The returned value may refer to the same object as \p value when the
implementation mutates
+ * that object in place.
+ */
+ FStructuralMutate maybe_inplace_mutate = nullptr;
+ /*!
+ * \brief Look up the replacement for a variable identity.
+ *
+ * \param mutator The active structural mutator.
+ * \param var The borrowed variable identity to look up.
+ * \return Raw ``TVMFFIAny`` carrying the owning replacement, FFI None on a
miss, or Error.
+ */
+ FStructuralVarRemapGet var_remap_get = nullptr;
+ /*!
+ * \brief Record the replacement for a variable identity.
+ *
+ * \param mutator The active structural mutator.
+ * \param var The borrowed variable identity to bind.
+ * \param mapped_value The borrowed replacement value.
+ * \return Raw ``TVMFFIAny`` carrying None or Error.
+ */
+ FStructuralVarRemapSet var_remap_set = nullptr;
+};
+
+/*!
+ * \brief Object node of a structural mutator.
+ */
+class StructuralMutatorObj : public Object {
+ public:
+ /*!
+ * \brief Mutate a value through the mutator vtable.
+ *
+ * \param value The value to mutate.
+ * \return The mutated owning value.
+ * \throws Error if mutation fails.
+ *
+ * This entry point never intentionally mutates \p value in place. Recursive
mutations
+ * also use \ref Mutate.
+ */
+ TVM_FFI_INLINE Any Mutate(AnyView value) { return
MutateExpected(value).value(); }
+
+ /*!
+ * \brief Exception-free form of \ref Mutate.
+ *
+ * \param value The value to mutate.
+ * \return The mutated owning value, or an Error if mutation failed.
+ */
+ TVM_FFI_INLINE Expected<Any> MutateExpected(AnyView value) noexcept {
+ return
details::ExpectedUnsafe::MoveFromTVMFFIAny<Any>((*vtable_->mutate)(this,
value));
+ }
+
+ /*!
+ * \brief Mutate a value, permitting an in-place implementation when it is
safe.
+ *
+ * \param value The borrowed value to mutate.
+ * \return The mutated owning value.
+ * \throws Error if mutation fails.
+ *
+ * The returned value may refer to the same object as \p value. Callers must
use the return value
+ * as the result of the mutation rather than assuming that the input object
was reused.
+ */
+ TVM_FFI_INLINE Any MaybeInplaceMutate(AnyView value) {
+ return MaybeInplaceMutateExpected(value).value();
+ }
+
+ /*!
+ * \brief Exception-free form of \ref MaybeInplaceMutate.
+ *
+ * \param value The borrowed value to mutate.
+ * \return The mutated owning value, or an Error if mutation failed.
+ */
+ TVM_FFI_INLINE Expected<Any> MaybeInplaceMutateExpected(AnyView value)
noexcept {
+ return details::ExpectedUnsafe::MoveFromTVMFFIAny<Any>(
+ (*vtable_->maybe_inplace_mutate)(this, value));
+ }
+
+ /*!
+ * \brief Mutate a value, using in-place mutation only for a uniquely owned
object.
+ *
+ * \param value The borrowed value to mutate.
+ * \return The mutated owning value, or an Error if mutation failed.
+ */
+ TVM_FFI_INLINE Expected<Any> MaybeInplaceMutateIfUniqueExpected(AnyView
value) noexcept {
+ const Object* obj = value.as<Object>();
+ if (obj != nullptr && obj->unique()) {
+ return MaybeInplaceMutateExpected(value);
+ }
+ return MutateExpected(value);
+ }
+
+ /*!
+ * \brief Apply the default structural mutation with copy-on-write behavior.
+ *
+ * \param value The value to mutate.
+ * \return The mutated value, or an Error if hook dispatch, copying, or
field mutation failed.
+ *
+ * \note A registered ``__s_mutate__`` hook is dispatched before the
reflected fallback and is
+ * responsible for variable-remap lookup and insertion when it
represents a FreeVar or DAG
+ * identity. Automatic remapping applies only to the reflected
fallback.
+ */
+ TVM_FFI_INLINE Expected<Any> DefaultMutateExpected(AnyView value) noexcept {
+ int32_t type_index = value.type_index();
+ static reflection::TypeAttrColumn
column(reflection::type_attr::kStructuralMutate);
+ AnyView attr = column[type_index];
+ if (attr.type_index() != TypeIndex::kTVMFFINone) {
+ if (attr.type_index() == TypeIndex::kTVMFFIOpaquePtr) {
+ auto* hook = reinterpret_cast<FStructuralMutate>(attr.cast<void*>());
+ return details::ExpectedUnsafe::MoveFromTVMFFIAny<Any>((*hook)(this,
value));
+ }
+ if (attr.type_index() == TypeIndex::kTVMFFIFunction) {
+ return attr.cast<Function>().CallExpected<Any>(this, value);
+ }
+ return Unexpected(Error(
+ "TypeError", "__s_mutate__ must be an opaque function pointer or
ffi.Function", ""));
+ }
+ if (type_index < TypeIndex::kTVMFFIStaticObjectBegin) {
+ return Any(value);
+ }
+ return MutateWithIdentityRemapExpected(value, [&]() -> Expected<Any> {
+ return details::MutateReflectedFieldsExpected(this, value);
+ });
+ }
+
+ /*!
+ * \brief Apply custom maybe-in-place mutation, or fall back to non-in-place
mutation.
+ *
+ * \param value The borrowed value to mutate.
+ * \return The mutated owning value, or an Error if mutation failed.
In-place changes
+ * completed before an Error are not rolled back.
+ *
+ * \note In-place mutation is explicitly opt-in. A registered
+ * ``__s_maybe_inplace_mutate__`` hook may rely on its input being
safe to mutate and owns
+ * any variable-remap handling. When the hook is absent, this method
calls
+ * \ref DefaultMutateExpected.
+ */
+ TVM_FFI_INLINE Expected<Any> DefaultMaybeInplaceMutateExpected(AnyView
value) noexcept {
+ int32_t type_index = value.type_index();
+ static reflection::TypeAttrColumn maybe_inplace_mutate_column(
+ reflection::type_attr::kStructuralMaybeInplaceMutate);
+ AnyView maybe_inplace_mutate_attr =
maybe_inplace_mutate_column[type_index];
+ if (maybe_inplace_mutate_attr.type_index() == TypeIndex::kTVMFFIOpaquePtr)
{
+ auto* hook =
reinterpret_cast<FStructuralMutate>(maybe_inplace_mutate_attr.cast<void*>());
+ return details::ExpectedUnsafe::MoveFromTVMFFIAny<Any>((*hook)(this,
value));
+ }
+ if (maybe_inplace_mutate_attr.type_index() == TypeIndex::kTVMFFIFunction) {
+ return
maybe_inplace_mutate_attr.cast<Function>().CallExpected<Any>(this, value);
+ }
+ return DefaultMutateExpected(value);
+ }
+
+ /*!
+ * \brief Look up the replacement recorded for a variable identity.
+ *
+ * \param var The borrowed variable identity to look up.
+ * \return The owning replacement, FFI None if no replacement exists, or an
Error if lookup
+ * fails.
+ *
+ * \note The identity must have ``kTVMFFISEqHashKindFreeVar`` or
+ * ``kTVMFFISEqHashKindDAGNode`` structural-equality metadata.
+ */
+ TVM_FFI_INLINE Expected<Any> VarRemapGetExpected(AnyView var) noexcept {
+ return
details::ExpectedUnsafe::MoveFromTVMFFIAny<Any>((*vtable_->var_remap_get)(this,
var));
+ }
+
+ /*!
+ * \brief Record the replacement for a variable identity.
+ *
+ * \param var The borrowed variable identity to bind.
+ * \param mapped_value The borrowed replacement value.
+ * \return Successful completion, or an Error if the binding is invalid or
cannot be stored.
+ *
+ * \note The identity must have ``kTVMFFISEqHashKindFreeVar`` or
+ * ``kTVMFFISEqHashKindDAGNode`` structural-equality metadata.
+ */
+ TVM_FFI_INLINE Expected<void> VarRemapSetExpected(AnyView var, AnyView
mapped_value) noexcept {
+ return details::ExpectedUnsafe::MoveFromTVMFFIAny<void>(
+ (*vtable_->var_remap_set)(this, var, mapped_value));
+ }
+
+ /*!
+ * \brief Return the current def-region context.
+ * \return The active def-region kind.
+ */
+ TVM_FFI_INLINE TVMFFIDefRegionKind def_region_kind() const { return
def_region_mode_; }
+
+ /*!
+ * \brief Temporarily switch the def-region context while invoking \p
callback.
+ *
+ * \param kind The def-region kind to set during the callback.
+ * \param callback A nullary callable that performs recursive mutation.
+ * \return The value returned by \p callback.
+ */
+ template <typename Callback>
+ TVM_FFI_INLINE auto WithDefRegionKind(TVMFFIDefRegionKind kind, Callback&&
callback)
+ -> decltype(std::forward<Callback>(callback)()) {
+ class Scope {
+ public:
+ Scope(StructuralMutatorObj* mutator, TVMFFIDefRegionKind kind)
+ : mutator_(mutator), old_kind_(mutator->def_region_mode_) {
+ mutator_->def_region_mode_ = kind;
+ }
+ ~Scope() { mutator_->def_region_mode_ = old_kind_; }
+ Scope(const Scope&) = delete;
+ Scope& operator=(const Scope&) = delete;
+
+ private:
+ StructuralMutatorObj* mutator_;
+ TVMFFIDefRegionKind old_kind_;
+ };
+ Scope scope(this, kind);
+ return std::forward<Callback>(callback)();
+ }
+
+ /// \cond Doxygen_Suppress
+ static constexpr const bool _type_mutable = true;
+ TVM_FFI_DECLARE_OBJECT_INFO("ffi.StructuralMutator", StructuralMutatorObj,
Object);
+ /// \endcond
+
+ protected:
+ /*!
+ * \brief Mutate a FreeVar or DAG identity once and reuse its final result.
+ *
+ * \tparam Mutation Nullary callable returning ``Expected<Any>``.
+ * \param value The borrowed value to mutate.
+ * \param mutation The complete mutation to apply on a cache miss.
+ * \return The cached or newly computed owning result, or an Error.
+ */
+ template <typename Mutation>
+ TVM_FFI_INLINE Expected<Any> MutateWithIdentityRemapExpected(AnyView value,
+ Mutation&&
mutation) noexcept {
+ int32_t type_index = value.type_index();
+ if (type_index < TypeIndex::kTVMFFIStaticObjectBegin) {
+ return mutation();
+ }
+
+ const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index);
+ bool is_remappable_identity =
+ type_info->metadata != nullptr &&
+ (type_info->metadata->structural_eq_hash_kind ==
kTVMFFISEqHashKindFreeVar ||
+ type_info->metadata->structural_eq_hash_kind ==
kTVMFFISEqHashKindDAGNode);
+ if (!is_remappable_identity) {
+ return mutation();
+ }
+
+ Expected<Any> mapped_value = VarRemapGetExpected(value);
+ if (TVM_FFI_PREDICT_FALSE(mapped_value.is_err())) {
+ return Unexpected(std::move(mapped_value).error());
+ }
+ if (details::ExpectedUnsafe::GetData(mapped_value).type_index() !=
TypeIndex::kTVMFFINone) {
+ return mapped_value;
+ }
+
+ Expected<Any> result = mutation();
+ if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
+ return result;
+ }
+ Expected<void> set_result =
+ VarRemapSetExpected(value, details::ExpectedUnsafe::GetData(result));
+ if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
+ return Unexpected(std::move(set_result).error());
+ }
+ return result;
+ }
+
+ /*!
+ * \brief Construct a structural mutator from an immutable dispatch vtable.
+ * \param vtable The non-null dispatch table for this mutator. It must
outlive this object.
+ */
+ explicit StructuralMutatorObj(const StructuralMutatorVTable* vtable) :
vtable_(vtable) {}
+
+ /*!
+ * \brief Non-owning pointer to the required ABI dispatch table.
+ */
+ const StructuralMutatorVTable* vtable_ = nullptr;
+
+ /*!
+ * \brief Current def-region context for def-region-aware structural
mutation.
+ */
+ TVMFFIDefRegionKind def_region_mode_ = kTVMFFIDefRegionKindNone;
+};
+
+/*!
+ * \brief ObjectRef wrapper for \ref StructuralMutatorObj.
+ *
+ * \sa StructuralMutatorObj
+ */
+class StructuralMutator : public ObjectRef {
+ public:
+ /*!
+ * \brief Construct from an existing mutator object pointer.
+ * \param n The object pointer to wrap.
+ */
+ explicit StructuralMutator(ObjectPtr<StructuralMutatorObj> n) :
ObjectRef(std::move(n)) {}
+
+ /// \cond Doxygen_Suppress
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StructuralMutator, ObjectRef,
StructuralMutatorObj);
+ /// \endcond
+};
+
+namespace details {
+
+/*!
+ * \brief Mutate the reflected structural fields of an object-backed value.
+ *
+ * \param mutator The active structural mutator.
+ * \param value The object-backed value to mutate.
+ * \return The original value when no field changes, a mutated shallow copy
otherwise, or an
+ * Error if copying or mutation failed.
+ */
+TVM_FFI_INLINE static Expected<Any>
MutateReflectedFieldsExpected(StructuralMutatorObj* mutator,
+ AnyView
value) noexcept {
+ const Object* obj = value.as<Object>();
+ int32_t type_index = obj->type_index();
+
+ static reflection::TypeAttrColumn
column(reflection::type_attr::kShallowCopy);
+ AnyView attr = column[type_index];
+ if (TVM_FFI_PREDICT_FALSE(attr.type_index() != TypeIndex::kTVMFFIFunction)) {
+ return Unexpected(Error("TypeError", "__ffi_shallow_copy__ must be an
ffi.Function", ""));
+ }
+
+ Expected<Any> result = attr.cast<Function>().CallExpected<Any>(value);
+ if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
+ return result;
+ }
+
+ const Any& result_value = details::ExpectedUnsafe::GetData(result);
+ Object* new_obj = const_cast<Object*>(result_value.as<Object>());
+ // Copy-on-write mutation requires a distinct target so partial updates
cannot modify the source.
+ if (TVM_FFI_PREDICT_FALSE(new_obj == nullptr || result.type_index() !=
value.type_index() ||
+ new_obj == obj)) {
+ return Unexpected(Error(
+ "TypeError",
+ "Shallow copy callback must return a distinct object with the same
type as its input", ""));
+ }
+
+ const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(new_obj->type_index());
+ bool field_changed = false;
+ auto mutate_fields = [&]() {
+ reflection::ForEachFieldInfoWithEarlyStop(
+ type_info, [&](const TVMFFIFieldInfo* field_info) -> bool {
+ if (field_info->flags & kTVMFFIFieldFlagBitMaskSEqHashIgnore) {
+ return false;
+ }
+
+ Any field_value;
+ void* field_addr = reinterpret_cast<char*>(new_obj) +
field_info->offset;
+ int ret_code = field_info->getter(field_addr,
reinterpret_cast<TVMFFIAny*>(&field_value));
+ if (TVM_FFI_PREDICT_FALSE(ret_code != 0)) {
+ result = Unexpected(details::MoveFromSafeCallRaised());
+ return true;
+ }
+
+ Expected<Any> mutated_field = [&]() -> Expected<Any> {
+ if (field_info->flags &
kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive) {
+ return
mutator->WithDefRegionKind(kTVMFFIDefRegionKindNonRecursive, [&]() {
+ return mutator->MutateExpected(field_value);
+ });
+ } else if (field_info->flags &
kTVMFFIFieldFlagBitMaskSEqHashDefRecursive) {
+ return mutator->WithDefRegionKind(kTVMFFIDefRegionKindRecursive,
[&]() {
+ return mutator->MutateExpected(field_value);
+ });
+ } else {
+ return mutator->MutateExpected(field_value);
+ }
+ }();
+ if (TVM_FFI_PREDICT_FALSE(mutated_field.is_err())) {
+ result = Unexpected(std::move(mutated_field).error());
+ return true;
+ }
+ const Any& new_field =
details::ExpectedUnsafe::GetData(mutated_field);
+ if (field_value.same_as(new_field)) {
+ return false;
+ }
+
+ if (TVM_FFI_PREDICT_FALSE(field_info->setter == nullptr)) {
+ result = Unexpected(Error(
+ "TypeError",
+ "Cannot structurally mutate field `" +
+ std::string(field_info->name.data, field_info->name.size)
+ "` of type `" +
+ std::string(type_info->type_key.data,
type_info->type_key.size) +
+ "` because it does not define a setter",
+ ""));
+ return true;
+ }
+
+ ret_code = reflection::CallFieldSetter(field_info, field_addr,
+ reinterpret_cast<const
TVMFFIAny*>(&new_field));
+ if (TVM_FFI_PREDICT_FALSE(ret_code != 0)) {
+ result = Unexpected(details::MoveFromSafeCallRaised());
+ return true;
+ }
+ field_changed = true;
+ return false;
+ });
+ };
+
+ // A non-recursive definition applies to the FreeVar itself, but its fields
are uses. The
+ // complete field traversal are clamped to None, then the definition region
is restored.
+ if (mutator->def_region_kind() == kTVMFFIDefRegionKindNonRecursive &&
+ type_info->metadata != nullptr &&
+ type_info->metadata->structural_eq_hash_kind ==
kTVMFFISEqHashKindFreeVar) {
+ mutator->WithDefRegionKind(kTVMFFIDefRegionKindNone, mutate_fields);
+ } else {
+ mutate_fields();
+ }
+
+ if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
+ return result;
+ }
+ if (!field_changed) {
+ return Any(value);
+ }
+ return result;
+}
+
+} // namespace details
+
+// ---------------------------------------------------------------------------
+// Structural Map API.
+// ---------------------------------------------------------------------------
+
+namespace details {
+
+/// \cond Doxygen_Suppress
+// Return from the current mutation function if Result is an Error.
+// Append Node to the mutate error context before returning.
+#define TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(Result, Node)
\
+ do {
\
+ auto&& tvm_ffi_res_ = (Result);
\
+ if (TVM_FFI_PREDICT_FALSE(tvm_ffi_res_.type_index() ==
::tvm::ffi::TypeIndex::kTVMFFIError)) { \
+ if ((Node).type_index() >=
::tvm::ffi::TypeIndex::kTVMFFIStaticObjectBegin) { \
+ ::tvm::ffi::Error tvm_ffi_mutate_err_ = tvm_ffi_res_.error();
\
+ ::tvm::ffi::details::UpdateVisitErrorContext(tvm_ffi_mutate_err_,
\
+
(Node).cast<::tvm::ffi::ObjectRef>()); \
+ }
\
+ return ::std::move(tvm_ffi_res_);
\
+ }
\
+ } while (0)
+/// \endcond
+
+/*!
+ * \brief Structural mutator that invokes typed callbacks during recursive
mapping.
+ *
+ * \tparam order Callback placement relative to child mapping.
+ * \tparam Dispatch Callback dispatcher with signature
+ * ``Expected<Any>(AnyView, TVMFFIDefRegionKind)``.
+ * \sa StructuralMapCallbackChain
+ */
+template <WalkOrder order, typename Dispatch>
+class StructuralMapMutatorObj : public StructuralMutatorObj {
+ public:
+ /*!
+ * \brief Construct a callback-aware mutator.
+ * \param dispatch The composed callback dispatcher.
+ */
+ explicit StructuralMapMutatorObj(Dispatch dispatch)
+ : StructuralMutatorObj(VTable()), dispatch_(std::move(dispatch)) {}
+
+ private:
+ /*!
+ * \brief Return the shared callback-aware mutator vtable.
+ * \return Pointer to the immutable mutator vtable for this specialization.
+ */
+ static const StructuralMutatorVTable* VTable() {
+ static const StructuralMutatorVTable vtable{
+ &StructuralMapMutatorObj::DispatchMutate,
+ &StructuralMapMutatorObj::DispatchMaybeInplaceMutate,
+ &StructuralMapMutatorObj::DispatchVarRemapGet,
+ &StructuralMapMutatorObj::DispatchVarRemapSet,
+ };
+ return &vtable;
+ }
+
+ /*!
+ * \brief Dispatch variable-remap lookup through this mutator's vtable.
+ * \param mutator The erased callback-aware mutator.
+ * \param var The borrowed variable identity to look up.
+ * \return Raw ``TVMFFIAny`` containing the owning replacement, FFI None, or
Error.
+ */
+ static TVMFFIAny DispatchVarRemapGet(StructuralMutatorObj* mutator, AnyView
var) noexcept {
+ auto* self = static_cast<StructuralMapMutatorObj*>(mutator);
+ Expected<Any> result = self->VarRemapGetImpl(var);
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ }
+
+ /*!
+ * \brief Dispatch variable-remap insertion through this mutator's vtable.
+ * \param mutator The erased callback-aware mutator.
+ * \param var The borrowed variable identity to bind.
+ * \param mapped_value The borrowed replacement value.
+ * \return Raw ``TVMFFIAny`` containing FFI None or Error.
+ */
+ static TVMFFIAny DispatchVarRemapSet(StructuralMutatorObj* mutator, AnyView
var,
+ AnyView mapped_value) noexcept {
+ auto* self = static_cast<StructuralMapMutatorObj*>(mutator);
+ Expected<void> result = self->VarRemapSetImpl(var, mapped_value);
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ }
+
+ /*!
+ * \brief Look up a replacement in this mutator's identity-substitution
environment.
+ * \param var The borrowed variable identity to look up.
+ * \return The owning replacement, FFI None on a miss, or an Error.
+ */
+ Expected<Any> VarRemapGetImpl(AnyView var) noexcept {
+ if (var.type_index() < TypeIndex::kTVMFFIStaticObjectBegin) {
+ return Unexpected(
+ Error("TypeError", "Variable-remap key must be an object-backed
value", ""));
+ }
+ try {
+ ObjectRef var_ref = var.cast<ObjectRef>();
+ std::optional<Any> result = var_remap_.Get(var_ref);
+ if (!result.has_value()) {
+ return Any(nullptr);
+ }
+ return *std::move(result);
+ } catch (const Error& err) {
+ return Unexpected(err);
+ }
+ }
+
+ /*!
+ * \brief Record a replacement in this mutator's identity-substitution
environment.
+ * \param var The borrowed variable identity to bind.
+ * \param mapped_value The borrowed replacement value.
+ * \return Successful completion, or an Error if the binding cannot be
stored.
+ */
+ Expected<void> VarRemapSetImpl(AnyView var, AnyView mapped_value) noexcept {
+ if (var.type_index() < TypeIndex::kTVMFFIStaticObjectBegin) {
+ return Unexpected(
+ Error("TypeError", "Variable-remap key must be an object-backed
value", ""));
+ }
+ try {
+ ObjectRef var_ref = var.cast<ObjectRef>();
+ Any owned_mapped_value(mapped_value);
+ var_remap_.Set(var_ref, owned_mapped_value);
+ return Expected<void>();
+ } catch (const Error& err) {
+ return Unexpected(err);
+ }
+ }
+
+ /*!
+ * \brief Dispatch callback-aware optional in-place mutation through the ABI
vtable.
+ * \param mutator The erased callback-aware mutator.
+ * \param value The borrowed value to mutate.
+ * \return Raw ``TVMFFIAny`` containing the mutated value or Error.
+ */
+ static TVMFFIAny DispatchMaybeInplaceMutate(StructuralMutatorObj* mutator,
+ AnyView value) noexcept {
+ auto* self = static_cast<StructuralMapMutatorObj*>(mutator);
+ return
ExpectedUnsafe::MoveToTVMFFIAny(self->MaybeInplaceMutateImpl(value));
+ }
+
+ /*!
+ * \brief Dispatch callback-aware mutation through the ABI vtable.
+ * \param mutator The erased callback-aware mutator.
+ * \param value The borrowed value to mutate.
+ * \return Raw ``TVMFFIAny`` containing the mutated value or Error.
+ */
+ static TVMFFIAny DispatchMutate(StructuralMutatorObj* mutator, AnyView
value) noexcept {
+ auto* self = static_cast<StructuralMapMutatorObj*>(mutator);
+ return ExpectedUnsafe::MoveToTVMFFIAny(self->MutateImpl(value));
+ }
+
+ /*!
+ * \brief Invoke the callback and select mutation according to its result
and walk order.
+ *
+ * \param value The borrowed value to mutate.
+ * \return The mutated value or an Error.
+ */
+ Expected<Any> MaybeInplaceMutateImpl(AnyView value) noexcept {
+ return MutateWithIdentityRemapExpected(value, [&]() -> Expected<Any> {
+ if constexpr (order == WalkOrder::kPreOrder) {
+ Expected<Any> callback_result = dispatch_(value, def_region_kind());
+
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(callback_result, value);
+ // A pre-order result can be mutated in place if unchanged or uniquely
owned.
+ const Any& mapped_value = ExpectedUnsafe::GetData(callback_result);
+ const TVMFFIAny* mapped_data =
AnyUnsafe::TVMFFIAnyPtrFromAny(mapped_value);
+ const TVMFFIAny input_data = value.CopyToTVMFFIAny();
+ if (mapped_data->type_index != input_data.type_index ||
+ mapped_data->zero_padding != input_data.zero_padding ||
+ mapped_data->v_int64 != input_data.v_int64) {
+ const Object* mapped_obj = mapped_value.as<Object>();
+ bool can_mutate_mapped_value_inplace = mapped_obj != nullptr &&
mapped_obj->unique();
+ Expected<Any> result = can_mutate_mapped_value_inplace
+ ?
DefaultMaybeInplaceMutateExpected(mapped_value)
+ : DefaultMutateExpected(mapped_value);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(result,
mapped_value);
+ return result;
+ }
+ Expected<Any> result = DefaultMaybeInplaceMutateExpected(value);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(result, value);
+ return result;
+ } else {
+ Expected<Any> result = DefaultMaybeInplaceMutateExpected(value);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(result, value);
+
+ const Any& mapped_value = ExpectedUnsafe::GetData(result);
+ Expected<Any> callback_result = dispatch_(mapped_value,
def_region_kind());
+
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(callback_result,
mapped_value);
+ return callback_result;
+ }
+ });
+ }
+
+ /*!
+ * \brief Mutate one value and invoke its matching callback in the
configured order.
+ * \param value The borrowed input value.
+ * \return The mutated value or an Error.
+ */
+ Expected<Any> MutateImpl(AnyView value) noexcept {
+ return MutateWithIdentityRemapExpected(value, [&]() -> Expected<Any> {
+ if constexpr (order == WalkOrder::kPreOrder) {
+ Expected<Any> callback_result = dispatch_(value, def_region_kind());
+
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(callback_result, value);
+
+ const Any& mapped_value = ExpectedUnsafe::GetData(callback_result);
+ Expected<Any> result = DefaultMutateExpected(mapped_value);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(result,
mapped_value);
+ return result;
+ } else {
+ Expected<Any> result = DefaultMutateExpected(value);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(result, value);
+
+ const Any& mapped_value = ExpectedUnsafe::GetData(result);
+ Expected<Any> callback_result = dispatch_(mapped_value,
def_region_kind());
+
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(callback_result,
mapped_value);
+ return callback_result;
+ }
+ });
+ }
+
+ /*! \brief Composed callback dispatcher owned by this mutator. */
+ Dispatch dispatch_;
+
+ /*! \brief Identity-substitution table. */
+ Map<ObjectRef, Any> var_remap_;
+};
+
+/*!
+ * \brief Build a callback dispatcher from a typed callback chain.
+ */
+struct StructuralMapCallbackChain {
+ public:
+ /*!
+ * \brief Construct a dispatcher owning \p callbacks.
+ * \tparam Callbacks Callback types.
+ * \param callbacks Callbacks tested in declaration order.
+ * \return A callback-aware structural-map dispatcher.
+ */
+ template <typename... Callbacks>
+ static auto FromChain(Callbacks... callbacks) {
+ auto callback_tuple = std::make_tuple(std::move(callbacks)...);
+ return [callbacks = std::move(callback_tuple)](
+ AnyView value, TVMFFIDefRegionKind kind) mutable ->
Expected<Any> {
+ try {
+ std::optional<Expected<Any>> result;
+ // Fold expression: each TryCallLink returns empty std::optional on
no-match
+ // (falsy) or a result on match (truthy); || short-circuits on first
match.
+ std::apply(
+ [&](auto&... callback) { (... || (result = TryCallLink(callback,
value, kind))); },
+ callbacks);
+ if (result.has_value()) {
+ return *std::move(result);
+ }
+ return Any(value);
+ } catch (const Error& err) {
+ return Unexpected(err);
+ }
+ };
+ }
+
+ private:
+ /*!
+ * \brief Invoke \p callback when \p value matches its first argument.
+ * \tparam Callback Callback type whose first argument selects the value
type.
+ * \param callback The callback under test.
+ * \param value The value to match and pass to the callback.
+ * \param kind The active def-region kind.
+ * \return The callback result on match, or an empty ``std::optional``
otherwise.
+ */
+ template <typename Callback>
+ TVM_FFI_INLINE static std::optional<Expected<Any>> TryCallLink(Callback&
callback, AnyView value,
+
TVMFFIDefRegionKind kind) {
+ using FuncInfo = FunctionInfo<std::decay_t<Callback>>;
+ static_assert(FuncInfo::num_args == 1 || FuncInfo::num_args == 2,
+ "StructuralMap callbacks must take one argument (value) or
two arguments "
+ "(value, def-region kind)");
+ using FirstArg = std::tuple_element_t<0, typename FuncInfo::ArgType>;
+ using TSub = std::remove_cv_t<std::remove_reference_t<FirstArg>>;
+ if constexpr (std::is_same_v<TSub, AnyView>) {
+ return InvokeCallbackLink(callback, value, kind);
+ } else if constexpr (std::is_same_v<TSub, Any>) {
+ return InvokeCallbackLink(callback, Any(value), kind);
+ } else {
+ if (auto opt = value.template as<TSub>()) {
+ return InvokeCallbackLink(callback, *std::move(opt), kind);
+ }
+ }
+ return std::nullopt;
+ }
+
+ /*!
+ * \brief Invoke a matched callback with optional def-region context.
+ * \tparam Callback Callable whose result is convertible to
``Expected<Any>``.
+ * \tparam Value Type of the converted value passed to the callback.
+ * \param callback The matched callback.
+ * \param value The converted value passed to the callback.
+ * \param kind The active def-region kind.
+ * \return The callback result converted to ``Expected<Any>``.
+ */
+ template <typename Callback, typename Value>
+ TVM_FFI_INLINE static Expected<Any> InvokeCallbackLink(Callback& callback,
Value&& value,
+ TVMFFIDefRegionKind
kind) {
+ using FuncInfo = FunctionInfo<std::decay_t<Callback>>;
+ if constexpr (FuncInfo::num_args == 1) {
+ return callback(std::forward<Value>(value));
+ } else {
+ return callback(std::forward<Value>(value), kind);
+ }
+ }
+};
+
+} // namespace details
+
+/*!
+ * \brief Map a structured value graph and invoke typed replacement callbacks.
+ *
+ * Each callback is selected by the type of its first argument. The argument
may be ``AnyView``,
+ * ``Any``, an object reference type, an object pointer type, or another
FFI-convertible POD type. A
+ * callback may optionally take a second ``TVMFFIDefRegionKind`` argument.
Callbacks are tested in
+ * declaration order and only the first strict type match is invoked. An
``AnyView`` callback
+ * argument is borrowed and must not be retained after the callback returns.
+ *
+ * Each callback should follow map semantics: it must not mutate the input in
place and should
+ * return ``Expected<Any>`` containing either the unchanged input or its
replacement. An ``Error``
+ * stops the mapping. In pre-order, an unchanged input or uniquely owned
replacement may
+ * continue through ``MaybeInplaceMutate``; a shared replacement uses
``Mutate``. In post-order,
+ * the callback runs after the node's optional in-place mutation. In-place
mutation is available
+ * only through an explicit ``__s_maybe_inplace_mutate__`` hook.
+ *
+ * Objects marked ``kTVMFFISEqHashKindFreeVar`` or
``kTVMFFISEqHashKindDAGNode`` are
+ * identity-substituted. A callback is invoked only for the first occurrence
of each identity; its
+ * final result, including an unchanged result, is reused for every later
occurrence in the same
+ * structural-map invocation.
+ *
+ * \sa WalkOrder, StructuralMutator
+ *
+ * Example:
+ *
+ * \code{.cpp}
+ * Expected<Any> result = StructuralMapExpected<WalkOrder::kPostOrder>(
+ * root,
+ * [](const IntImm& value) -> Expected<Any> {
+ * if (value->value < 0) {
+ * return Unexpected(Error("ValueError", "negative constant", ""));
+ * }
+ * return Any(IntImm(value->value + 1));
+ * },
+ * [](const Add& add, TVMFFIDefRegionKind kind) -> Expected<Any> {
+ * // In post-order, add->lhs and add->rhs have already been mapped.
+ * return Any(add);
+ * });
+ * \endcode
+ *
+ * \tparam order Whether callbacks run before or after recursively mapping
children.
+ * \tparam Callbacks Callback types whose first parameters select matching
values.
+ * \param root The borrowed root value to map.
+ * \param callbacks Callbacks tested in declaration order. Each accepts
``(value)`` or
+ * ``(value, def_region_kind)`` and should return ``Expected<Any>``.
+ * \return The mapped owning value, or an Error if mapping or a callback fails.
+ *
+ * \note Return type of each callback should be ``Expected<Any>``.
+ */
+template <WalkOrder order, typename... Callbacks>
+Expected<Any> StructuralMapExpected(AnyView root, Callbacks&&... callbacks)
noexcept {
+ static_assert(sizeof...(Callbacks) != 0, "StructuralMap requires at least
one callback");
+ auto dispatch =
+
details::StructuralMapCallbackChain::FromChain(std::forward<Callbacks>(callbacks)...);
+ using Mutator = details::StructuralMapMutatorObj<order, decltype(dispatch)>;
+ StructuralMutator mutator(make_object<Mutator>(std::move(dispatch)));
+ return mutator->MaybeInplaceMutateIfUniqueExpected(root);
+}
+
+/*!
+ * \brief Throwing form of \ref tvm::ffi::StructuralMapExpected.
+ *
+ * See \ref tvm::ffi::StructuralMapExpected for callback dispatch, ordering,
and ownership
+ * semantics.
+ *
+ * \tparam order Whether callbacks run before or after recursively mapping
children.
+ * \tparam Callbacks Callback types whose first parameters select matching
values.
+ * \param root The borrowed root value to map.
+ * \param callbacks Callbacks tested in declaration order. Each accepts
``(value)`` or
+ * ``(value, def_region_kind)`` and should return ``Expected<Any>``.
+ * \return The mapped owning value.
+ * \throws Error if mapping or a callback fails.
+ *
+ * \note Return type of each callback should be ``Expected<Any>``.
+ */
+template <WalkOrder order, typename... Callbacks>
+Any StructuralMap(AnyView root, Callbacks&&... callbacks) {
+ return StructuralMapExpected<order>(root,
std::forward<Callbacks>(callbacks)...).value();
+}
+
+} // namespace ffi
+} // namespace tvm
+
+#endif // TVM_FFI_EXTRA_STRUCTURAL_MUTATE_H_
diff --git a/include/tvm/ffi/extra/structural_visit.h
b/include/tvm/ffi/extra/structural_visit.h
index e4ea5fad..463843b6 100644
--- a/include/tvm/ffi/extra/structural_visit.h
+++ b/include/tvm/ffi/extra/structural_visit.h
@@ -218,10 +218,8 @@ class StructuralVisitorObj : public Object {
}
if (TVM_FFI_PREDICT_FALSE(attr.type_index() != TypeIndex::kTVMFFINone)) {
- return Unexpected(Error("TypeError",
-
std::string(reflection::type_attr::kStructuralVisit) +
- " must be an opaque function pointer or
ffi.Function",
- ""));
+ return Unexpected(
+ Error("TypeError", "__s_visit__ must be an opaque function pointer
or ffi.Function", ""));
}
if (type_index < TypeIndex::kTVMFFIStaticObjectBegin) {
@@ -305,42 +303,47 @@ TVM_FFI_INLINE static Expected<Optional<VisitInterrupt>>
VisitReflectedFieldsExp
StructuralVisitorObj* visitor, const Object* obj) noexcept {
int32_t type_index = obj->type_index();
const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index);
- // A non-recursive definition applies to a FreeVar itself, but not to its
children. All other
- // inherited modes propagate until an explicit field annotation overrides
them.
- TVMFFIDefRegionKind inherited_kind = visitor->def_region_kind();
- if (inherited_kind == kTVMFFIDefRegionKindNonRecursive &&
type_info->metadata != nullptr &&
+ auto visit_fields = [&]() -> Expected<Optional<VisitInterrupt>> {
+ Expected<Optional<VisitInterrupt>> result =
Optional<VisitInterrupt>(std::nullopt);
+ reflection::ForEachFieldInfoWithEarlyStop(
+ type_info, [&](const TVMFFIFieldInfo* field_info) -> bool {
+ if (field_info->flags & kTVMFFIFieldFlagBitMaskSEqHashIgnore) {
+ return false;
+ }
+
+ Any field_value;
+ const void* field_addr = reinterpret_cast<const char*>(obj) +
field_info->offset;
+ int ret_code = field_info->getter(const_cast<void*>(field_addr),
+
reinterpret_cast<TVMFFIAny*>(&field_value));
+ if (TVM_FFI_PREDICT_FALSE(ret_code != 0)) {
+ result = Unexpected(details::MoveFromSafeCallRaised());
+ return true;
+ }
+
+ if (field_info->flags &
kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive) {
+ result =
visitor->WithDefRegionKind(kTVMFFIDefRegionKindNonRecursive, [&]() {
+ return visitor->VisitExpected(field_value);
+ });
+ } else if (field_info->flags &
kTVMFFIFieldFlagBitMaskSEqHashDefRecursive) {
+ result = visitor->WithDefRegionKind(kTVMFFIDefRegionKindRecursive,
[&]() {
+ return visitor->VisitExpected(field_value);
+ });
+ } else {
+ result = visitor->VisitExpected(field_value);
+ }
+ return StructuralVisitNeedEarlyReturn(result);
+ });
+ return result;
+ };
+
+ // A non-recursive definition applies to the FreeVar itself, but its fields
are uses. The
+ // complete field traversal are clamped to None, then the definition region
is restored.
+ if (visitor->def_region_kind() == kTVMFFIDefRegionKindNonRecursive &&
+ type_info->metadata != nullptr &&
type_info->metadata->structural_eq_hash_kind ==
kTVMFFISEqHashKindFreeVar) {
- inherited_kind = kTVMFFIDefRegionKindNone;
+ return visitor->WithDefRegionKind(kTVMFFIDefRegionKindNone, visit_fields);
}
-
- Expected<Optional<VisitInterrupt>> result =
Optional<VisitInterrupt>(std::nullopt);
- reflection::ForEachFieldInfoWithEarlyStop(
- type_info, [&](const TVMFFIFieldInfo* field_info) -> bool {
- if (field_info->flags & kTVMFFIFieldFlagBitMaskSEqHashIgnore) {
- return false;
- }
-
- Any field_value;
- const void* field_addr = reinterpret_cast<const char*>(obj) +
field_info->offset;
- int ret_code = field_info->getter(const_cast<void*>(field_addr),
-
reinterpret_cast<TVMFFIAny*>(&field_value));
- if (TVM_FFI_PREDICT_FALSE(ret_code != 0)) {
- result = Unexpected(details::MoveFromSafeCallRaised());
- return true;
- }
-
- TVMFFIDefRegionKind kind = inherited_kind;
- if (field_info->flags & kTVMFFIFieldFlagBitMaskSEqHashDefNonRecursive)
{
- kind = kTVMFFIDefRegionKindNonRecursive;
- } else if (field_info->flags &
kTVMFFIFieldFlagBitMaskSEqHashDefRecursive) {
- kind = kTVMFFIDefRegionKindRecursive;
- }
-
- result =
- visitor->WithDefRegionKind(kind, [&]() { return
visitor->VisitExpected(field_value); });
- return StructuralVisitNeedEarlyReturn(result);
- });
- return result;
+ return visit_fields();
}
} // namespace details
@@ -483,13 +486,13 @@ namespace details {
* accept either ``(value)`` or ``(value, def_region_kind)``.
*/
template <WalkOrder order, typename Dispatch>
-class StructuralWalkCallbackVisitorObj : public StructuralVisitorObj {
+class StructuralWalkVisitorObj : public StructuralVisitorObj {
public:
/*!
* \brief Construct a structural walk visitor.
* \param dispatch The composed dispatcher invoked on each visited node.
*/
- explicit StructuralWalkCallbackVisitorObj(Dispatch dispatch)
+ explicit StructuralWalkVisitorObj(Dispatch dispatch)
: StructuralVisitorObj(VTable()), dispatch_(std::move(dispatch)) {}
private:
@@ -499,7 +502,7 @@ class StructuralWalkCallbackVisitorObj : public
StructuralVisitorObj {
*/
static const StructuralVisitorVTable* VTable() {
static const StructuralVisitorVTable vtable{
- &StructuralWalkCallbackVisitorObj::DispatchVisit,
+ &StructuralWalkVisitorObj::DispatchVisit,
};
return &vtable;
}
@@ -511,7 +514,7 @@ class StructuralWalkCallbackVisitorObj : public
StructuralVisitorObj {
* \return Interrupt state, or an error if traversal failed.
*/
static TVMFFIAny DispatchVisit(StructuralVisitorObj* self, AnyView value)
noexcept {
- return
static_cast<StructuralWalkCallbackVisitorObj*>(self)->VisitImpl(value);
+ return static_cast<StructuralWalkVisitorObj*>(self)->VisitImpl(value);
}
/*!
@@ -698,7 +701,7 @@ Expected<Optional<VisitInterrupt>>
StructuralWalkExpected(AnyView root,
static_assert(sizeof...(Callbacks) != 0, "StructuralWalk requires at least
one callback");
auto dispatch =
details::StructuralWalkCallbackChain::FromChain(std::forward<Callbacks>(callbacks)...);
- using Visitor = details::StructuralWalkCallbackVisitorObj<order,
decltype(dispatch)>;
+ using Visitor = details::StructuralWalkVisitorObj<order, decltype(dispatch)>;
StructuralVisitor visitor(make_object<Visitor>(std::move(dispatch)));
return visitor->VisitExpected(root);
}
diff --git a/include/tvm/ffi/reflection/accessor.h
b/include/tvm/ffi/reflection/accessor.h
index daa0f26d..0f8337d9 100644
--- a/include/tvm/ffi/reflection/accessor.h
+++ b/include/tvm/ffi/reflection/accessor.h
@@ -497,6 +497,50 @@ inline constexpr const char* kSEqual = "__s_equal__";
* visit structural children.
*/
inline constexpr const char* kStructuralVisit = "__s_visit__";
+
+/*!
+ * \brief Custom structural mutation hook used by ``StructuralMutator``.
+ *
+ * The hook receives the active mutator and the input value. It should
recursively mutate structural
+ * children through the mutator, return the input unchanged if no fields are
changed, and avoid
+ * intentionally mutating the source object. The input is borrowed and remains
valid for the
+ * duration of the call.
+ *
+ * Value type: either an opaque function pointer to a C++ structural mutation
hook
+ *
+ * ``TVMFFIAny (*)(StructuralMutatorObj* mutator, AnyView value) noexcept``
+ *
+ * returning raw ``Expected<Any>`` storage, or an ``ffi::Function`` with
signature
+ *
+ * ``(StructuralMutator mutator, Any value) -> Any``.
+ *
+ * This is the canonical mutation hook. A type that defines
+ * ``kStructuralMaybeInplaceMutate`` must also define this attribute.
+ */
+inline constexpr const char* kStructuralMutate = "__s_mutate__";
+/*!
+ * \brief Optional custom mutation hook that may reuse its input object.
+ *
+ * The hook receives the active mutator and a borrowed input value that is
safe to mutate in place.
+ * Callers must route shared objects through non-in-place mutation instead.
The hook may mutate and
+ * return the source object, delegate to non-in-place mutation, or return
another replacement.
+ *
+ * Value type: either an opaque function pointer to a C++ structural mutation
hook
+ *
+ * ``TVMFFIAny (*)(StructuralMutatorObj* mutator, AnyView value) noexcept``
+ *
+ * returning raw ``Expected<Any>`` storage, or an
+ * ``ffi::Function`` with signature
+ *
+ * ``(StructuralMutator mutator, Any value) -> Any``.
+ *
+ * This hook is optional. When it is absent,
``DefaultMaybeInplaceMutateExpected`` falls back to
+ * non-in-place mutation through ``kStructuralMutate`` or reflected structural
fields. In-place
+ * mutation is therefore explicitly opt-in and is never inferred from
ownership by the reflected
+ * fallback.
+ */
+inline constexpr const char* kStructuralMaybeInplaceMutate =
"__s_maybe_inplace_mutate__";
+
/*!
* \brief Serialize object data to a JSON-compatible value.
*
diff --git a/python/tvm_ffi/__init__.py b/python/tvm_ffi/__init__.py
index 9ce96656..59d55d89 100644
--- a/python/tvm_ffi/__init__.py
+++ b/python/tvm_ffi/__init__.py
@@ -78,6 +78,7 @@ if TYPE_CHECKING or not _is_config_mode():
from .structural import (
DefRegionKind,
StructuralKey,
+ StructuralMutator,
StructuralVisitor,
VisitInterrupt,
WalkOrder,
@@ -85,6 +86,7 @@ if TYPE_CHECKING or not _is_config_mode():
get_first_structural_mismatch,
structural_equal,
structural_hash,
+ structural_map,
structural_walk,
)
from . import serialization
@@ -154,6 +156,7 @@ __all__ = [
"Shape",
"StreamContext",
"StructuralKey",
+ "StructuralMutator",
"StructuralVisitor",
"Tensor",
"VisitInterrupt",
@@ -184,6 +187,7 @@ __all__ = [
"structural",
"structural_equal",
"structural_hash",
+ "structural_map",
"structural_walk",
"system_lib",
"use_raw_stream",
diff --git a/python/tvm_ffi/_ffi_api.py b/python/tvm_ffi/_ffi_api.py
index 38a54fff..c19828eb 100644
--- a/python/tvm_ffi/_ffi_api.py
+++ b/python/tvm_ffi/_ffi_api.py
@@ -25,7 +25,7 @@ from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Mapping, MutableMapping, MutableSequence,
Sequence
from ctypes import c_void_p
- from tvm_ffi import Device, Module, Object, StructuralKey as
_StructuralKey, StructuralVisitor as _StructuralVisitor, VisitInterrupt as
_VisitInterrupt
+ from tvm_ffi import Device, Module, Object, StructuralKey as
_StructuralKey, StructuralMutator as _StructuralMutator, StructuralVisitor as
_StructuralVisitor, VisitInterrupt as _VisitInterrupt
from tvm_ffi.access_path import AccessPath
from typing import Any, Callable
# isort: on
@@ -111,6 +111,13 @@ if TYPE_CHECKING:
def StructuralHash(_0: Any, _1: bool, _2: bool, /) -> int: ...
def StructuralKey(_0: Any, /) -> _StructuralKey: ...
def StructuralKeyEqual(_0: Any, _1: Any, /) -> bool: ...
+ def StructuralMap(_0: Any, _1: Sequence[tuple[int, Callable[..., Any]]],
_2: Sequence[tuple[int, Callable[..., Any]]], _3: int, /) -> Any: ...
+ def StructuralMutatorDefRegionKind(_0: _StructuralMutator, /) -> int: ...
+ def StructuralMutatorMaybeInplaceMutate(_0: _StructuralMutator, _1: Any,
/) -> Any: ...
+ def StructuralMutatorMutate(_0: _StructuralMutator, _1: Any, /) -> Any: ...
+ def StructuralMutatorVarRemapGet(_0: _StructuralMutator, _1: Any, /) ->
Any: ...
+ def StructuralMutatorVarRemapSet(_0: _StructuralMutator, _1: Any, _2: Any,
/) -> None: ...
+ def StructuralMutatorWithDefRegionKind(_0: _StructuralMutator, _1: int,
_2: Callable[..., Any], /) -> Any: ...
def StructuralVisitorDefRegionKind(_0: _StructuralVisitor, /) -> int: ...
def StructuralVisitorVisit(_0: _StructuralVisitor, _1: Any, /) ->
_VisitInterrupt | None: ...
def StructuralVisitorWithDefRegionKind(_0: _StructuralVisitor, _1: int,
_2: Callable[..., Any], /) -> Any: ...
@@ -202,6 +209,13 @@ __all__ = [
"StructuralHash",
"StructuralKey",
"StructuralKeyEqual",
+ "StructuralMap",
+ "StructuralMutatorDefRegionKind",
+ "StructuralMutatorMaybeInplaceMutate",
+ "StructuralMutatorMutate",
+ "StructuralMutatorVarRemapGet",
+ "StructuralMutatorVarRemapSet",
+ "StructuralMutatorWithDefRegionKind",
"StructuralVisitorDefRegionKind",
"StructuralVisitorVisit",
"StructuralVisitorWithDefRegionKind",
diff --git a/python/tvm_ffi/dataclasses/py_class.py
b/python/tvm_ffi/dataclasses/py_class.py
index 75e7fc39..1109cbff 100644
--- a/python/tvm_ffi/dataclasses/py_class.py
+++ b/python/tvm_ffi/dataclasses/py_class.py
@@ -66,6 +66,8 @@ _FFI_TYPE_ATTR_NAMES: frozenset[str] = frozenset(
"__s_equal__",
"__s_hash__",
"__s_visit__",
+ "__s_mutate__",
+ "__s_maybe_inplace_mutate__",
"__data_to_json__",
"__data_from_json__",
}
diff --git a/python/tvm_ffi/structural.py b/python/tvm_ffi/structural.py
index ed87dc7d..7a91ef6d 100644
--- a/python/tvm_ffi/structural.py
+++ b/python/tvm_ffi/structural.py
@@ -33,6 +33,7 @@ from .registry import register_object
__all__ = [
"DefRegionKind",
"StructuralKey",
+ "StructuralMutator",
"StructuralVisitor",
"VisitInterrupt",
"WalkOrder",
@@ -40,17 +41,20 @@ __all__ = [
"get_first_structural_mismatch",
"structural_equal",
"structural_hash",
+ "structural_map",
"structural_walk",
]
class WalkOrder(IntEnum):
- """Callback placement before or after visiting children for structural
walks.
+ """Callback placement before or after recursively traversing children.
See Also
--------
:py:func:`tvm_ffi.structural_walk`
Walk an object graph, invoke matching callbacks.
+ :py:func:`tvm_ffi.structural_map`
+ Structurally map an object graph with replacement callbacks.
"""
@@ -87,6 +91,8 @@ class DefRegionKind(IntEnum):
:py:class:`tvm_ffi.StructuralVisitor`
Structural traversal visitor that carries object dispatch and
def-region
state across recursive visits.
+ :py:class:`tvm_ffi.StructuralMutator`
+ Structural mutator that recursively mutates values.
"""
@@ -315,7 +321,7 @@ class VisitInterrupt(Object):
return None
- result = tvm_ffi.structural_walk(root, (object, on_node))
+ result = tvm_ffi.structural_walk(root, (object, on_node), order="pre")
if result is not None:
found = result.value
@@ -408,11 +414,124 @@ class StructuralVisitor(Object):
return _ffi_api.StructuralVisitorWithDefRegionKind(self, kind,
callback)
+@register_object("ffi.StructuralMutator")
+class StructuralMutator(Object):
+ """Low-level structural mutator.
+
+ This class exposes the mutator object used by structural-map and custom
+ mutation hooks.
+ """
+
+ def maybe_inplace_mutate(self, value: Any) -> Any:
+ """Mutate ``value``, permitting an in-place implementation when safe.
+
+ The caller must ensure that an object-backed ``value`` is safe to
mutate
+ in place and use :meth:`mutate` for a shared object.
+
+ Parameters
+ ----------
+ value
+ Value to mutate.
+
+ Returns
+ -------
+ result
+ The mutated owning value. It may refer to the same object as
``value``.
+
+ """
+ return _ffi_api.StructuralMutatorMaybeInplaceMutate(self, value)
+
+ def mutate(self, value: Any) -> Any:
+ """Mutate ``value`` without modifying it in place.
+
+ The original value is returned when none of its structural fields
+ change; otherwise, the result is a mutated copy.
+
+ Parameters
+ ----------
+ value
+ Value to mutate.
+
+ Returns
+ -------
+ result
+ The mutated owning value.
+
+ """
+ return _ffi_api.StructuralMutatorMutate(self, value)
+
+ def var_remap_get(self, var: Object) -> Any | None:
+ """Return the replacement recorded for a variable identity.
+
+ Parameters
+ ----------
+ var
+ Variable identity to look up. It must be an object-backed value
+ with free-variable or DAG-node structural metadata.
+
+ Returns
+ -------
+ result
+ The recorded replacement, or ``None`` if ``var`` has no
replacement.
+
+ """
+ return _ffi_api.StructuralMutatorVarRemapGet(self, var)
+
+ def var_remap_set(self, var: Object, mapped_value: Any) -> None:
+ """Record or replace the substitution for a variable identity.
+
+ Parameters
+ ----------
+ var
+ Variable identity to bind. It must be an object-backed value with
+ free-variable or DAG-node structural metadata.
+
+ mapped_value
+ Replacement returned for subsequent lookups of ``var``.
+
+ """
+ _ffi_api.StructuralMutatorVarRemapSet(self, var, mapped_value)
+
+ def def_region_kind(self) -> DefRegionKind:
+ """Return the currently active structural def-region kind.
+
+ Returns
+ -------
+ kind
+ The active :class:`DefRegionKind`.
+
+ """
+ return DefRegionKind(_ffi_api.StructuralMutatorDefRegionKind(self))
+
+ def with_def_region_kind(
+ self,
+ kind: int,
+ callback: Callable[[], Any],
+ ) -> Any:
+ """Run ``callback`` with a temporarily active def-region kind.
+
+ Parameters
+ ----------
+ kind
+ Def-region kind to use while running ``callback``.
+
+ callback
+ Nullary callable to execute inside the scoped region.
+
+ Returns
+ -------
+ result
+ The value returned by ``callback``.
+
+ """
+ return _ffi_api.StructuralMutatorWithDefRegionKind(self, kind,
callback)
+
+
def structural_walk(
root: Any,
callbacks: tuple | Sequence | Callable = (),
with_def_region_kind: tuple | Sequence | Callable = (),
- order: str | WalkOrder = "pre",
+ order: str | WalkOrder = "post",
) -> VisitInterrupt | None:
"""Walk a value structurally and invoke the first matching typed callback.
@@ -441,8 +560,9 @@ def structural_walk(
as ``callbacks``.
order
- ``"pre"``/``WalkOrder.PREORDER`` to invoke callbacks before children,
or
- ``"post"``/``WalkOrder.POSTORDER`` to invoke callbacks after children.
+ ``"post"``/``WalkOrder.POSTORDER`` (the default) to invoke callbacks
+ after children, or ``"pre"``/``WalkOrder.PREORDER`` to invoke callbacks
+ before children.
Returns
-------
@@ -477,68 +597,163 @@ def structural_walk(
else:
raise ValueError(f"Unknown structural walk order: {order!r}")
- def normalize_callbacks(
- callbacks: tuple | Sequence | Callable,
- ) -> list[tuple[object, Callable]]:
- callback_entries = []
-
- def add_callback_entry(callback_entry: tuple) -> None:
- callback_type, fn = callback_entry
- callback_types = callback_type if isinstance(callback_type, tuple)
else (callback_type,)
- callback_entries.extend((t, fn) for t in callback_types)
-
- if callable(callbacks):
- callback_entries.append((Any, callbacks))
- elif isinstance(callbacks, tuple) and len(callbacks) == 2 and
callable(callbacks[1]):
- add_callback_entry(callbacks)
- elif isinstance(callbacks, Sequence) and not isinstance(callbacks,
(str, bytes)):
- for callback in callbacks:
- if (
- not isinstance(callback, tuple)
- or len(callback) != 2
- or not callable(callback[1])
- ):
- raise TypeError(
- "structural_walk callbacks within a sequence must be "
- "(type, callback) tuples"
- )
- add_callback_entry(callback)
- else:
- raise TypeError(
- "structural_walk callbacks must be callbacks, (type, callback)
entries, "
- "((type1, type2, ...), callback) entries, or sequences of
tuple entries"
- )
- return callback_entries
-
def wrap_callback_with_def_region_kind(fn: Callable[..., Any]) ->
Callable[[Any, int], Any]:
return lambda value, kind: fn(value, DefRegionKind(kind))
- callback_entries = normalize_callbacks(callbacks)
- callback_entries_with_def_region_kind =
normalize_callbacks(with_def_region_kind)
+ callback_entries = _normalize_callbacks(callbacks,
api_name="structural_walk")
+ callback_entries_with_def_region_kind = _normalize_callbacks(
+ with_def_region_kind, api_name="structural_walk"
+ )
entries: list[tuple[int, Callable[[Any], Any]]] = [
- (_callback_type_to_type_index(t), fn) for t, fn in callback_entries
+ (_callback_type_to_type_index(t, api_name="structural_walk"), fn)
+ for t, fn in callback_entries
]
entries_with_def_region_kind: list[tuple[int, Callable[[Any, int], Any]]]
= [
- (_callback_type_to_type_index(t),
wrap_callback_with_def_region_kind(fn))
+ (
+ _callback_type_to_type_index(t, api_name="structural_walk"),
+ wrap_callback_with_def_region_kind(fn),
+ )
for t, fn in callback_entries_with_def_region_kind
]
return _ffi_api.StructuralWalk(root, entries,
entries_with_def_region_kind, order_int)
-def _callback_type_to_type_index(callback_type: type[Any] | Any) -> int:
+def structural_map(
+ root: Any,
+ callbacks: tuple | Sequence | Callable = (),
+ with_def_region_kind: tuple | Sequence | Callable = (),
+ order: str | WalkOrder = "post",
+) -> Any:
+ """Structurally map a value using typed replacement callbacks.
+
+ Each callback must follow map semantics: it returns the unchanged input or
+ a replacement value and must not mutate its input in place.
+
+ Parameters
+ ----------
+ root
+ Root value to map.
+
+ callbacks
+ Normal callbacks. These callbacks receive one argument, ``value``, and
+ return its mapped value. Callback entries are tried in order.
+
+ May be one of:
+
+ - A single callback, used as a ``typing.Any`` catch-all.
+ - A ``(type, callback)`` entry.
+ - A grouped ``((type1, type2, ...), callback)`` entry.
+ - A sequence of entries.
+
+ Types may be builtins, registered FFI object classes, or
+ ``typing.Any``/``object`` as a catch-all.
+
+ with_def_region_kind
+ Def-region-aware callbacks. These callbacks receive
+ ``(value, def_region_kind)`` and return the mapped value. They accept
+ the same callback entry forms as ``callbacks``.
+
+ order
+ ``"post"``/``WalkOrder.POSTORDER`` (the default) to invoke callbacks
+ after children, or ``"pre"``/``WalkOrder.PREORDER`` to invoke callbacks
+ before children.
+
+ Returns
+ -------
+ result
+ The mapped owning value.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ def fold_add(expr):
+ if isinstance(expr.lhs, IntImm) and isinstance(expr.rhs, IntImm):
+ return IntImm(expr.lhs.value + expr.rhs.value)
+ return expr
+
+
+ optimized = tvm_ffi.structural_map(
+ function,
+ (Add, fold_add),
+ )
+
+ """
+ if isinstance(order, WalkOrder):
+ order_int = int(order)
+ elif order in ("pre", "post"):
+ order_int = int(WalkOrder.PREORDER if order == "pre" else
WalkOrder.POSTORDER)
+ else:
+ raise ValueError(f"Unknown structural map order: {order!r}")
+
+ def wrap_callback_with_def_region_kind(fn: Callable[..., Any]) ->
Callable[[Any, int], Any]:
+ return lambda value, kind: fn(value, DefRegionKind(kind))
+
+ callback_entries = _normalize_callbacks(callbacks,
api_name="structural_map")
+ callback_entries_with_def_region_kind = _normalize_callbacks(
+ with_def_region_kind, api_name="structural_map"
+ )
+
+ entries: list[tuple[int, Callable[[Any], Any]]] = [
+ (_callback_type_to_type_index(t, api_name="structural_map"), fn)
+ for t, fn in callback_entries
+ ]
+ entries_with_def_region_kind: list[tuple[int, Callable[[Any, int], Any]]]
= [
+ (
+ _callback_type_to_type_index(t, api_name="structural_map"),
+ wrap_callback_with_def_region_kind(fn),
+ )
+ for t, fn in callback_entries_with_def_region_kind
+ ]
+ return _ffi_api.StructuralMap(root, entries, entries_with_def_region_kind,
order_int)
+
+
+def _normalize_callbacks(
+ callbacks: tuple | Sequence | Callable,
+ *,
+ api_name: str,
+) -> list[tuple[object, Callable]]:
+ """Normalize typed callback shorthand into individual callback entries."""
+ callback_entries = []
+
+ def add_callback_entry(callback_entry: tuple) -> None:
+ callback_type, fn = callback_entry
+ callback_types = callback_type if isinstance(callback_type, tuple)
else (callback_type,)
+ callback_entries.extend((t, fn) for t in callback_types)
+
+ if callable(callbacks):
+ callback_entries.append((Any, callbacks))
+ elif isinstance(callbacks, tuple) and len(callbacks) == 2 and
callable(callbacks[1]):
+ add_callback_entry(callbacks)
+ elif isinstance(callbacks, Sequence) and not isinstance(callbacks, (str,
bytes)):
+ for callback in callbacks:
+ if not isinstance(callback, tuple) or len(callback) != 2 or not
callable(callback[1]):
+ raise TypeError(
+ f"{api_name} callbacks within a sequence must be (type,
callback) tuples"
+ )
+ add_callback_entry(callback)
+ else:
+ raise TypeError(
+ f"{api_name} callbacks must be callbacks, (type, callback)
entries, "
+ "((type1, type2, ...), callback) entries, or sequences of tuple
entries"
+ )
+ return callback_entries
+
+
+def _callback_type_to_type_index(callback_type: type[Any] | Any, *, api_name:
str) -> int:
"""Convert a callback arg type to a type index."""
annotation = Any if callback_type is object else callback_type
try:
type_index =
core.TypeSchema.from_annotation(annotation).origin_type_index
except TypeError as err:
raise TypeError(
- "structural_walk callback type must be a supported builtin, "
+ f"{api_name} callback type must be a supported builtin, "
"typing.Any/object, or an FFI-registered object class"
) from err
if type_index < 0 and annotation is not Any:
raise TypeError(
- "structural_walk callback type_index is negative, the only"
+ f"{api_name} callback type_index is negative, the only "
"acceptable negative type_index is -1 for Any"
)
return type_index
diff --git a/rust/tvm-ffi/src/extra/structural_visit.rs
b/rust/tvm-ffi/src/extra/structural_visit.rs
index 1d15efd8..6d5f9f76 100644
--- a/rust/tvm-ffi/src/extra/structural_visit.rs
+++ b/rust/tvm-ffi/src/extra/structural_visit.rs
@@ -953,7 +953,7 @@ fn visit_children_raw<C: ChildVisit>(
{
// Fast path: read the MapBaseObj storage layout directly, like
// the SeqPrefix path for arrays — zero FFI calls per entry.
- // Dict entries are snapshotted first to keep the re-entrant
+ // Dict values are snapshotted first to keep the re-entrant
// mutation guard. If the one-time layout validation fails
// (e.g. an ABI-debug build), fall back to the packed-functor
// iteration protocol.
@@ -1025,9 +1025,10 @@ fn visit_sequence<C: ChildVisit>(
Ok(())
}
-/// Walk map/dict entries by reading the `MapBaseObj` storage directly —
-/// the map analog of the `SeqPrefix` array fast path. `snapshot` first
-/// takes owned copies of all entries (Dict re-entrant mutation guard).
+/// Walk map/dict values by reading the `MapBaseObj` storage directly —
+/// the map analog of the `SeqPrefix` array fast path. Keys are structural
+/// anchors and are skipped. `snapshot` first takes owned copies of all values
+/// (Dict re-entrant mutation guard).
#[inline(never)]
fn visit_map_layout<C: ChildVisit>(
value: TVMFFIAny,
@@ -1043,22 +1044,16 @@ fn visit_map_layout<C: ChildVisit>(
let mut cursor = unsafe { MapCursor::new(map) };
if snapshot {
- let mut entries: Vec<(Any, Any)> = Vec::with_capacity(size);
+ let mut values: Vec<Any> = Vec::with_capacity(size);
for _ in 0..size {
- let Some((key, val)) = (unsafe { cursor.next() }) else {
+ let Some((_, value)) = (unsafe { cursor.next() }) else {
return Err(runtime_error("native visitor: map iteration ended
early").into());
};
- entries.push((
- Any::from(unsafe { view_of(&key) }),
- Any::from(unsafe { view_of(&val) }),
- ));
+ values.push(Any::from(unsafe { view_of(&value) }));
}
- for (index, (key, val)) in entries.into_iter().enumerate() {
+ for (index, value) in values.into_iter().enumerate() {
visitor
- .visit_child(raw_of_owned(&key), def_region_kind)
- .map_err(|halt| with_error_context(halt, &format!("dict key
[{index}]")))?;
- visitor
- .visit_child(raw_of_owned(&val), def_region_kind)
+ .visit_child(raw_of_owned(&value), def_region_kind)
.map_err(|halt| with_error_context(halt, &format!("dict value
[{index}]")))?;
}
return Ok(());
@@ -1068,14 +1063,11 @@ fn visit_map_layout<C: ChildVisit>(
// callbacks, so visit them in place. The `size` bound also guards the
// dense iteration list against corruption-induced cycles.
for index in 0..size {
- let Some((key, val)) = (unsafe { cursor.next() }) else {
+ let Some((_, value)) = (unsafe { cursor.next() }) else {
return Err(runtime_error("native visitor: map iteration ended
early").into());
};
visitor
- .visit_child(key, def_region_kind)
- .map_err(|halt| with_error_context(halt, &format!("map key
[{index}]")))?;
- visitor
- .visit_child(val, def_region_kind)
+ .visit_child(value, def_region_kind)
.map_err(|halt| with_error_context(halt, &format!("map value
[{index}]")))?;
}
Ok(())
@@ -1084,10 +1076,10 @@ fn visit_map_layout<C: ChildVisit>(
/// Cold fallback used when the mirrored layout fails validation (e.g. an
/// ABI-debug build): iterate through the public packed functors. Map storage
/// is private C++; the Rust binding itself uses these iterator functors, so
-/// no structural visiting or traversal control leaves Rust. Entries are
-/// snapshotted before user callbacks run — required for Dict, whose mutation
-/// invalidates the iterator, and harmless for immutable Map on this
-/// non-performance path.
+/// no structural visiting or traversal control leaves Rust. Keys are
structural
+/// anchors and are skipped. Values are snapshotted before user callbacks run —
+/// required for Dict, whose mutation invalidates the iterator, and harmless
for
+/// immutable Map on this non-performance path.
fn visit_map<C: ChildVisit>(
value: TVMFFIAny,
visitor: &mut C,
@@ -1114,20 +1106,16 @@ fn visit_map<C: ChildVisit>(
let iter_any = Function::get_global(iter_name)?.call_packed(&[unsafe {
view_of(&value) }])?;
let iter = Function::try_from(iter_any)?;
- let mut entries = Vec::with_capacity(size);
+ let mut values = Vec::with_capacity(size);
for index in 0..size {
- let key = iter.call_packed(&[AnyView::from(&0i64)])?;
let map_value = iter.call_packed(&[AnyView::from(&1i64)])?;
- entries.push((key, map_value));
+ values.push(map_value);
if index + 1 != size {
iter.call_packed(&[AnyView::from(&2i64)])?;
}
}
- for (index, (key, map_value)) in entries.into_iter().enumerate() {
- visitor
- .visit_child(raw_of_owned(&key), def_region_kind)
- .map_err(|halt| with_error_context(halt, &format!("{kind} key
[{index}]")))?;
+ for (index, map_value) in values.into_iter().enumerate() {
visitor
.visit_child(raw_of_owned(&map_value), def_region_kind)
.map_err(|halt| with_error_context(halt, &format!("{kind} value
[{index}]")))?;
diff --git a/rust/tvm-ffi/tests/test_structural_visit.rs
b/rust/tvm-ffi/tests/test_structural_visit.rs
index 723de807..7fa04585 100644
--- a/rust/tvm-ffi/tests/test_structural_visit.rs
+++ b/rust/tvm-ffi/tests/test_structural_visit.rs
@@ -48,16 +48,19 @@ fn plain_walk_uses_native_sequence_fallback() {
}
#[test]
-fn plain_walk_uses_native_map_fallback() {
+fn plain_walk_visits_map_values_without_visiting_keys() {
let root: Map<FfiString, i64> = [(FfiString::from("a"), 1i64),
(FfiString::from("b"), 2i64)]
.into_iter()
.collect();
let mut integers = 0;
+ let mut strings = 0;
assert!(structural_walk(
&root,
|value: &VisitValue| {
if value.cast::<i64>().is_some() {
integers += 1;
+ } else if value.cast::<FfiString>().is_some() {
+ strings += 1;
}
WalkResult::Advance
},
@@ -66,6 +69,7 @@ fn plain_walk_uses_native_map_fallback() {
.unwrap()
.is_none());
assert_eq!(integers, 2);
+ assert_eq!(strings, 0);
}
#[derive(Default)]
@@ -229,7 +233,7 @@ fn mutable_dict_is_snapshotted_before_callbacks() {
}
#[test]
-fn dense_map_layout_is_traversed_completely() {
+fn dense_map_layout_visits_all_values_without_visiting_keys() {
// More than 4 entries forces the dense (block + iteration list) layout.
let root: Map<FfiString, i64> = (0..9)
.map(|i| (FfiString::from(format!("k{i}")), i as i64))
@@ -251,7 +255,7 @@ fn dense_map_layout_is_traversed_completely() {
.unwrap()
.is_none());
assert_eq!(sum, (0..9).sum::<i64>());
- assert_eq!(strings, 9);
+ assert_eq!(strings, 0);
}
#[test]
diff --git a/src/ffi/extra/structural_mutate.cc
b/src/ffi/extra/structural_mutate.cc
new file mode 100644
index 00000000..c240392c
--- /dev/null
+++ b/src/ffi/extra/structural_mutate.cc
@@ -0,0 +1,378 @@
+/*
+ * 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.
+ */
+/*!
+ * \file src/ffi/extra/structural_mutate.cc
+ * \brief Structural mutator and structural map registration.
+ */
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/dict.h>
+#include <tvm/ffi/container/list.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/accessor.h>
+#include <tvm/ffi/reflection/registry.h>
+
+#include <utility>
+
+namespace tvm {
+namespace ffi {
+
+namespace details {
+
+/*!
+ * \brief Runtime structural map for callback arrays.
+ *
+ * \param root The root value to map.
+ * \param callbacks Runtime callback entries of ``(type_index,
ffi::Function)`` invoked as
+ * ``callback(value)``.
+ * \param callbacks_with_def_region_kind Runtime callback entries of
``(type_index, ffi::Function)``
+ * invoked as ``callback(value,
def_region_kind)``.
+ * \param order Integer value of \ref WalkOrder.
+ * \return The mapped owning value, or an Error.
+ */
+Expected<Any> StructuralMapExpected(
+ AnyView root, const Array<Tuple<int32_t, Function>>& callbacks,
+ const Array<Tuple<int32_t, Function>>& callbacks_with_def_region_kind, int
order) noexcept {
+ auto dispatch = [callbacks, callbacks_with_def_region_kind](
+ AnyView x, TVMFFIDefRegionKind kind) -> Expected<Any> {
+ for (const auto& entry : callbacks) {
+ int32_t type_index = entry.template get<0>();
+ if (!RuntimeTypeIndexMatch(x.type_index(), type_index)) {
+ continue;
+ }
+ Function fn = entry.template get<1>();
+ return fn.CallExpected<Any>(x);
+ }
+ for (const auto& entry : callbacks_with_def_region_kind) {
+ int32_t type_index = entry.template get<0>();
+ if (!RuntimeTypeIndexMatch(x.type_index(), type_index)) {
+ continue;
+ }
+ Function fn = entry.template get<1>();
+ return fn.CallExpected<Any>(x, kind);
+ }
+ return Any(x);
+ };
+
+ if (order == static_cast<int>(WalkOrder::kPreOrder)) {
+ using Mutator = StructuralMapMutatorObj<WalkOrder::kPreOrder,
decltype(dispatch)>;
+ StructuralMutator mutator(make_object<Mutator>(std::move(dispatch)));
+ return mutator->MaybeInplaceMutateIfUniqueExpected(root);
+ } else {
+ using Mutator = StructuralMapMutatorObj<WalkOrder::kPostOrder,
decltype(dispatch)>;
+ StructuralMutator mutator(make_object<Mutator>(std::move(dispatch)));
+ return mutator->MaybeInplaceMutateIfUniqueExpected(root);
+ }
+}
+
+// ---------------------------------------------------------------------------
+// Built-in container structural mutation.
+// ---------------------------------------------------------------------------
+
+/*!
+ * \brief Structurally mutate the elements of a sequence container.
+ *
+ * \tparam SeqObj The underlying sequence object type.
+ * \param mutator The active structural mutator.
+ * \param value The borrowed sequence container.
+ * \param source The sequence object stored in \p value.
+ * \return The mutated sequence, or an Error.
+ */
+template <typename SeqObj>
+Expected<Any> MutateSeqContainerExpected(StructuralMutatorObj* mutator,
AnyView value,
+ const SeqObj* source) noexcept {
+ try {
+ int64_t size = static_cast<int64_t>(source->size());
+ ObjectPtr<SeqObj> output = nullptr;
+
+ for (int64_t i = 0; i < size; ++i) {
+ const Any& item = source->at(i);
+ Expected<Any> mapped_item = mutator->MutateExpected(item);
+ if (TVM_FFI_PREDICT_FALSE(mapped_item.is_err())) {
+ return Unexpected(std::move(mapped_item).error());
+ }
+ const Any& mapped_value = details::ExpectedUnsafe::GetData(mapped_item);
+
+ if (output == nullptr) {
+ if (item.same_as(mapped_value)) {
+ continue;
+ }
+ output = SeqObj::CreateRepeated(size, Any());
+ for (int64_t j = 0; j < i; ++j) {
+ output->SetItem(j, source->at(j));
+ }
+ }
+ output->SetItem(i, mapped_value);
+ }
+
+ if (output == nullptr) {
+ return Any(value);
+ }
+ return Any(ObjectRef(std::move(output)));
+ } catch (const Error& err) {
+ return Unexpected(err);
+ }
+}
+
+/*!
+ * \brief Structurally mutate the elements of a sequence container in place
when safe.
+ *
+ * \tparam SeqObj The underlying sequence object type.
+ * \param mutator The active structural mutator.
+ * \param value The borrowed sequence container, which must be safe to mutate
in place.
+ * \param target The sequence object stored in \p value.
+ * \return The mutated sequence, or an Error.
+ */
+template <typename SeqObj>
+Expected<Any> MaybeInplaceMutateSeqContainerExpected(StructuralMutatorObj*
mutator, AnyView value,
+ SeqObj* target) noexcept {
+ try {
+ for (int64_t i = 0; i < static_cast<int64_t>(target->size()); ++i) {
+ const Any& item = target->at(i);
+ Expected<Any> mapped_item =
mutator->MaybeInplaceMutateIfUniqueExpected(item);
+ if (TVM_FFI_PREDICT_FALSE(mapped_item.is_err())) {
+ return Unexpected(std::move(mapped_item).error());
+ }
+ const Any& mapped_value = details::ExpectedUnsafe::GetData(mapped_item);
+
+ if (!item.same_as(mapped_value)) {
+ target->SetItem(i, mapped_value);
+ }
+ }
+ return Any(value);
+ } catch (const Error& err) {
+ return Unexpected(err);
+ }
+}
+
+/*!
+ * \brief Structurally mutate the values of a map container.
+ *
+ * \tparam MapObjType The underlying map object type.
+ * \param mutator The active structural mutator.
+ * \param value The borrowed map container.
+ * \param source The map object stored in \p value.
+ * \return The mutated map, or an Error.
+ */
+template <typename MapObjType>
+Expected<Any> MutateMapValuesExpected(StructuralMutatorObj* mutator, AnyView
value,
+ const MapObjType* source) noexcept {
+ try {
+ ObjectPtr<Object> output = nullptr;
+ MapBaseObj::iterator output_it;
+ size_t index = 0;
+
+ for (auto source_it = source->begin(); source_it != source->end();
++source_it, ++index) {
+ const Any& old_value = source_it->second;
+ Expected<Any> mapped_value = mutator->MutateExpected(old_value);
+ if (TVM_FFI_PREDICT_FALSE(mapped_value.is_err())) {
+ return Unexpected(std::move(mapped_value).error());
+ }
+
+ const Any& new_value = details::ExpectedUnsafe::GetData(mapped_value);
+ bool changed = !old_value.same_as(new_value);
+ if (output == nullptr) {
+ if (!changed) {
+ continue;
+ }
+ output = MapObjType::ShallowCopy(source);
+ output_it = static_cast<MapBaseObj*>(output.get())->begin();
+ for (size_t i = 0; i < index; ++i) {
+ ++output_it;
+ }
+ }
+ if (changed) {
+ output_it->second = new_value;
+ }
+ ++output_it;
+ }
+
+ if (output == nullptr) {
+ return Any(value);
+ }
+ return Any(ObjectRef(std::move(output)));
+ } catch (const Error& err) {
+ return Unexpected(err);
+ }
+}
+
+/*!
+ * \brief Structurally mutate the values of a map container in place when safe.
+ *
+ * \tparam MapObjType The underlying map object type.
+ * \param mutator The active structural mutator.
+ * \param value The borrowed map container, which must be safe to mutate in
place.
+ * \param target The map object stored in \p value.
+ * \return The mutated map, or an Error.
+ */
+template <typename MapObjType>
+Expected<Any> MaybeInplaceMutateMapValuesExpected(StructuralMutatorObj*
mutator, AnyView value,
+ MapObjType* target) noexcept
{
+ try {
+ for (auto it = target->begin(); it != target->end(); ++it) {
+ const Any& old_value = it->second;
+ Expected<Any> mapped_value =
mutator->MaybeInplaceMutateIfUniqueExpected(old_value);
+ if (TVM_FFI_PREDICT_FALSE(mapped_value.is_err())) {
+ return Unexpected(std::move(mapped_value).error());
+ }
+ const Any& new_value = details::ExpectedUnsafe::GetData(mapped_value);
+
+ if (!old_value.same_as(new_value)) {
+ it->second = new_value;
+ }
+ }
+ return Any(value);
+ } catch (const Error& err) {
+ return Unexpected(err);
+ }
+}
+
+/*! \brief Identity structural mutation hook for immutable String and Bytes
leaves. */
+TVMFFIAny MutateImmutableLeaf(StructuralMutatorObj*, AnyView value) noexcept {
+ Expected<Any> result = Any(value);
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+}
+
+/*! \brief Structural mutation hook for ArrayObj. */
+TVMFFIAny MutateArray(StructuralMutatorObj* mutator, AnyView value) noexcept {
+ Expected<Any> result = MutateSeqContainerExpected(mutator, value,
value.cast<const ArrayObj*>());
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+}
+
+/*! \brief Maybe-in-place structural mutation hook for ArrayObj. */
+TVMFFIAny MaybeInplaceMutateArray(StructuralMutatorObj* mutator, AnyView
value) noexcept {
+ Expected<Any> result = MaybeInplaceMutateSeqContainerExpected(
+ mutator, value, const_cast<ArrayObj*>(value.cast<const ArrayObj*>()));
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+}
+
+/*! \brief Structural mutation hook for ListObj. */
+TVMFFIAny MutateList(StructuralMutatorObj* mutator, AnyView value) noexcept {
+ Expected<Any> result = MutateSeqContainerExpected(mutator, value,
value.cast<const ListObj*>());
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+}
+
+/*! \brief Maybe-in-place structural mutation hook for ListObj. */
+TVMFFIAny MaybeInplaceMutateList(StructuralMutatorObj* mutator, AnyView value)
noexcept {
+ Expected<Any> result = MaybeInplaceMutateSeqContainerExpected(
+ mutator, value, const_cast<ListObj*>(value.cast<const ListObj*>()));
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+}
+
+/*! \brief Structural mutation hook for MapObj. */
+TVMFFIAny MutateMap(StructuralMutatorObj* mutator, AnyView value) noexcept {
+ Expected<Any> result = MutateMapValuesExpected(mutator, value,
value.cast<const MapObj*>());
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+}
+
+/*! \brief Maybe-in-place structural mutation hook for MapObj. */
+TVMFFIAny MaybeInplaceMutateMap(StructuralMutatorObj* mutator, AnyView value)
noexcept {
+ Expected<Any> result = MaybeInplaceMutateMapValuesExpected(
+ mutator, value, const_cast<MapObj*>(value.cast<const MapObj*>()));
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+}
+
+/*! \brief Structural mutation hook for DictObj. */
+TVMFFIAny MutateDict(StructuralMutatorObj* mutator, AnyView value) noexcept {
+ Expected<Any> result = MutateMapValuesExpected(mutator, value,
value.cast<const DictObj*>());
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+}
+
+/*! \brief Maybe-in-place structural mutation hook for DictObj. */
+TVMFFIAny MaybeInplaceMutateDict(StructuralMutatorObj* mutator, AnyView value)
noexcept {
+ Expected<Any> result = MaybeInplaceMutateMapValuesExpected(
+ mutator, value, const_cast<DictObj*>(value.cast<const DictObj*>()));
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+}
+} // namespace details
+
+// ---------------------------------------------------------------------------
+// Static registration.
+// ---------------------------------------------------------------------------
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<StructuralMutatorObj>(); // NOLINT(bugprone-unused-raii)
+ refl::GlobalDef()
+ .def_method("ffi.StructuralMutatorMaybeInplaceMutate",
+ &StructuralMutatorObj::MaybeInplaceMutate)
+ .def_method("ffi.StructuralMutatorMutate", &StructuralMutatorObj::Mutate)
+ .def_method("ffi.StructuralMutatorVarRemapGet",
+ [](const StructuralMutator& mutator, AnyView var) {
+ return mutator->VarRemapGetExpected(var).value();
+ })
+ .def_method("ffi.StructuralMutatorVarRemapSet",
+ [](const StructuralMutator& mutator, AnyView var, AnyView
mapped_value) {
+ mutator->VarRemapSetExpected(var, mapped_value).value();
+ })
+ .def_method("ffi.StructuralMutatorDefRegionKind",
&StructuralMutatorObj::def_region_kind)
+ .def_method(
+ "ffi.StructuralMutatorWithDefRegionKind",
+ [](const StructuralMutator& mutator, TVMFFIDefRegionKind kind, const
Function& callback) {
+ return mutator->WithDefRegionKind(kind, callback);
+ })
+ .def("ffi.StructuralMap",
+ [](AnyView root, const Array<Tuple<int32_t, Function>>& callbacks,
+ const Array<Tuple<int32_t, Function>>&
callbacks_with_def_region_kind,
+ int32_t order) -> Any {
+ return details::StructuralMapExpected(root, callbacks,
callbacks_with_def_region_kind,
+ order)
+ .value();
+ });
+ refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralMutate);
+ refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralMaybeInplaceMutate);
+ refl::TypeAttrDef<details::StringObj>()
+ .attr(refl::type_attr::kStructuralMutate,
+
reinterpret_cast<void*>(static_cast<FStructuralMutate>(&details::MutateImmutableLeaf)))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+
reinterpret_cast<void*>(static_cast<FStructuralMutate>(&details::MutateImmutableLeaf)));
+ refl::TypeAttrDef<details::BytesObj>()
+ .attr(refl::type_attr::kStructuralMutate,
+
reinterpret_cast<void*>(static_cast<FStructuralMutate>(&details::MutateImmutableLeaf)))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+
reinterpret_cast<void*>(static_cast<FStructuralMutate>(&details::MutateImmutableLeaf)));
+ refl::TypeAttrDef<ArrayObj>()
+ .attr(refl::type_attr::kStructuralMutate,
+
reinterpret_cast<void*>(static_cast<FStructuralMutate>(&details::MutateArray)))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(
+
static_cast<FStructuralMutate>(&details::MaybeInplaceMutateArray)));
+ refl::TypeAttrDef<ListObj>()
+ .attr(refl::type_attr::kStructuralMutate,
+
reinterpret_cast<void*>(static_cast<FStructuralMutate>(&details::MutateList)))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(
+
static_cast<FStructuralMutate>(&details::MaybeInplaceMutateList)));
+ refl::TypeAttrDef<MapObj>()
+ .attr(refl::type_attr::kStructuralMutate,
+
reinterpret_cast<void*>(static_cast<FStructuralMutate>(&details::MutateMap)))
+ .attr(
+ refl::type_attr::kStructuralMaybeInplaceMutate,
+
reinterpret_cast<void*>(static_cast<FStructuralMutate>(&details::MaybeInplaceMutateMap)));
+ refl::TypeAttrDef<DictObj>()
+ .attr(refl::type_attr::kStructuralMutate,
+
reinterpret_cast<void*>(static_cast<FStructuralMutate>(&details::MutateDict)))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(
+
static_cast<FStructuralMutate>(&details::MaybeInplaceMutateDict)));
+}
+
+} // namespace ffi
+} // namespace tvm
diff --git a/src/ffi/extra/structural_visit.cc
b/src/ffi/extra/structural_visit.cc
index 913bb4e4..bcd4ce41 100644
--- a/src/ffi/extra/structural_visit.cc
+++ b/src/ffi/extra/structural_visit.cc
@@ -74,11 +74,11 @@ Expected<Optional<VisitInterrupt>> StructuralWalkExpected(
};
if (order == static_cast<int>(WalkOrder::kPreOrder)) {
- using Visitor = StructuralWalkCallbackVisitorObj<WalkOrder::kPreOrder,
decltype(dispatch)>;
+ using Visitor = StructuralWalkVisitorObj<WalkOrder::kPreOrder,
decltype(dispatch)>;
StructuralVisitor visitor(make_object<Visitor>(std::move(dispatch)));
return visitor->VisitExpected(root);
} else {
- using Visitor = StructuralWalkCallbackVisitorObj<WalkOrder::kPostOrder,
decltype(dispatch)>;
+ using Visitor = StructuralWalkVisitorObj<WalkOrder::kPostOrder,
decltype(dispatch)>;
StructuralVisitor visitor(make_object<Visitor>(std::move(dispatch)));
return visitor->VisitExpected(root);
}
@@ -92,10 +92,9 @@ TVMFFIAny VisitSeqContainer(StructuralVisitorObj* visitor,
const SeqBaseObj* seq
return
ExpectedUnsafe::MoveToTVMFFIAny(Expected<Optional<VisitInterrupt>>(std::nullopt));
}
-/*! \brief Visit keys and values in a map container. */
+/*! \brief Visit values in a map container while treating keys as structural
anchors. */
TVMFFIAny VisitMapContainer(StructuralVisitorObj* visitor, const MapBaseObj*
map) noexcept {
for (const auto& kv : *map) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(kv.first));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(kv.second));
}
return
ExpectedUnsafe::MoveToTVMFFIAny(Expected<Optional<VisitInterrupt>>(std::nullopt));
diff --git a/tests/cpp/extra/test_structural_mutate.cc
b/tests/cpp/extra/test_structural_mutate.cc
new file mode 100644
index 00000000..2066d074
--- /dev/null
+++ b/tests/cpp/extra/test_structural_mutate.cc
@@ -0,0 +1,259 @@
+/*
+ * 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.
+ */
+#include <gtest/gtest.h>
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/string.h>
+
+#include <cstdint>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "../testing_object.h"
+
+namespace {
+
+using namespace tvm::ffi;
+using namespace tvm::ffi::testing;
+
+using AnyArray = Array<Any>;
+using StringMap = Map<String, Any>;
+
+Expected<Any> Increment(int64_t value) { return Any(value + 1); }
+
+template <WalkOrder order>
+void CheckNestedArrayMapOrder(const std::vector<std::string>& expected_trace) {
+ AnyArray inner_array{int64_t{1}};
+ const Object* inner_array_address = inner_array.get();
+ StringMap map{{"value", Any(std::move(inner_array))}};
+ const Object* map_address = map.get();
+ AnyArray root{Any(std::move(map))};
+ const Object* root_address = root.get();
+ std::vector<std::string> trace;
+
+ AnyArray mapped =
+ StructuralMap<order>(
+ root,
+ [&](const AnyArray& array) -> Expected<Any> {
+ trace.emplace_back(array.get() == root_address ? "outer-array" :
"inner-array");
+ return Any(array);
+ },
+ [&](const StringMap& value) -> Expected<Any> {
+ trace.emplace_back("map");
+ return Any(value);
+ },
+ [&](const String&) -> Expected<Any> {
+ trace.emplace_back("map-key");
+ return Any(String("renamed"));
+ },
+ [&](int64_t value) -> Expected<Any> {
+ trace.emplace_back("int");
+ return Any(value + 1);
+ })
+ .template cast<AnyArray>();
+
+ StringMap mapped_map = mapped[0].cast<StringMap>();
+ AnyArray mapped_inner_array = mapped_map["value"].cast<AnyArray>();
+
+ EXPECT_EQ(trace, expected_trace);
+ EXPECT_EQ(mapped.get(), root_address);
+ EXPECT_EQ(mapped_map.get(), map_address);
+ EXPECT_EQ(mapped_inner_array.get(), inner_array_address);
+ EXPECT_EQ(mapped_inner_array[0].cast<int64_t>(), 2);
+ EXPECT_EQ(mapped_map.count("value"), 1U);
+ EXPECT_EQ(mapped_map.count("renamed"), 0U);
+}
+
+TEST(StructuralMap, MapsNestedArrayAndMapInConfiguredOrder) {
+ CheckNestedArrayMapOrder<WalkOrder::kPreOrder>({"outer-array", "map",
"inner-array", "int"});
+ CheckNestedArrayMapOrder<WalkOrder::kPostOrder>({"int", "inner-array",
"map", "outer-array"});
+}
+
+TEST(StructuralMap, PreservesSharedArrayAndMapInputs) {
+ // A shared Array is copied when one of its elements changes.
+ {
+ AnyArray child{int64_t{1}};
+ const Object* child_address = child.get();
+ AnyArray root{Any(std::move(child))};
+ AnyArray owner = root; //
NOLINT(performance-unnecessary-copy-initialization)
+ const Object* root_address = root.get();
+
+ AnyArray mapped = StructuralMap<WalkOrder::kPostOrder>(root,
Increment).cast<AnyArray>();
+ AnyArray original_child = root[0].cast<AnyArray>();
+ AnyArray mapped_child = mapped[0].cast<AnyArray>();
+
+ EXPECT_NE(mapped.get(), root_address);
+ EXPECT_EQ(owner.get(), root_address);
+ EXPECT_NE(mapped_child.get(), child_address);
+ EXPECT_EQ(original_child[0].cast<int64_t>(), 1);
+ EXPECT_EQ(mapped_child[0].cast<int64_t>(), 2);
+ }
+
+ // A shared Map and its changed value path are also copied.
+ {
+ AnyArray value{int64_t{1}};
+ const Object* value_address = value.get();
+ StringMap root{{"value", Any(std::move(value))}};
+ StringMap owner = root; //
NOLINT(performance-unnecessary-copy-initialization)
+ const Object* root_address = root.get();
+
+ StringMap mapped = StructuralMap<WalkOrder::kPostOrder>(root,
Increment).cast<StringMap>();
+ AnyArray original_value = root["value"].cast<AnyArray>();
+ AnyArray mapped_value = mapped["value"].cast<AnyArray>();
+
+ EXPECT_NE(mapped.get(), root_address);
+ EXPECT_EQ(owner.get(), root_address);
+ EXPECT_NE(mapped_value.get(), value_address);
+ EXPECT_EQ(original_value[0].cast<int64_t>(), 1);
+ EXPECT_EQ(mapped_value[0].cast<int64_t>(), 2);
+ }
+
+ // Copy-on-write remains lazy: an unchanged shared Map is returned directly.
+ {
+ StringMap root{{"value", AnyArray{int64_t{1}}}};
+ StringMap owner = root; //
NOLINT(performance-unnecessary-copy-initialization)
+
+ StringMap mapped =
+ StructuralMap<WalkOrder::kPostOrder>(root, [](int64_t value) ->
Expected<Any> {
+ return Any(value);
+ }).cast<StringMap>();
+
+ EXPECT_TRUE(mapped.same_as(root));
+ EXPECT_TRUE(owner.same_as(root));
+ EXPECT_TRUE(mapped["value"].same_as(root["value"]));
+ }
+}
+
+TEST(StructuralMap, PreOrderRecursivelyMapsCallbackResult) {
+ StringMap root{{"value", AnyArray{int64_t{1}}}};
+ AnyArray replacement{int64_t{10}};
+
+ StringMap mapped =
+ StructuralMap<WalkOrder::kPreOrder>(
+ root, [&](const AnyArray&) -> Expected<Any> { return
Any(replacement); }, Increment)
+ .cast<StringMap>();
+ AnyArray mapped_value = mapped["value"].cast<AnyArray>();
+
+ EXPECT_FALSE(mapped_value.same_as(replacement));
+ EXPECT_EQ(replacement[0].cast<int64_t>(), 10);
+ EXPECT_EQ(mapped_value[0].cast<int64_t>(), 11);
+}
+
+template <WalkOrder order>
+void CheckRepeatedVarRemap() {
+ TVar var("n");
+ StringMap use{{"use", var}};
+ AnyArray root{var, Any(std::move(use))};
+ int callback_count = 0;
+
+ AnyArray mapped = StructuralMap<order>(root, [&](const TVarObj* value) ->
Expected<Any> {
+ ++callback_count;
+ return Any(TVar(value->name + "-mapped"));
+ }).template cast<AnyArray>();
+ TVar mapped_var = mapped[0].cast<TVar>();
+ StringMap mapped_uses = mapped[1].cast<StringMap>();
+ TVar mapped_use = mapped_uses["use"].cast<TVar>();
+
+ EXPECT_EQ(callback_count, 1);
+ EXPECT_TRUE(mapped_var.same_as(mapped_use));
+ EXPECT_EQ(mapped_var->name, "n-mapped");
+ EXPECT_EQ(var->name, "n");
+}
+
+TEST(StructuralMap, ReusesFinalCallbackResultForRepeatedVar) {
+ CheckRepeatedVarRemap<WalkOrder::kPreOrder>();
+ CheckRepeatedVarRemap<WalkOrder::kPostOrder>();
+}
+
+AnyArray MakeStringAndBytesLeaves() {
+ return AnyArray{int64_t{1}, String("1234567"), String("12345678"),
Bytes("1234567", 7),
+ Bytes("12345678", 8)};
+}
+
+template <WalkOrder order>
+void CheckStringAndBytesLeaves() {
+ // An unmatched callback leaves inline and heap-backed values untouched.
+ {
+ AnyArray root = MakeStringAndBytesLeaves();
+ EXPECT_EQ(root[1].type_index(), TypeIndex::kTVMFFISmallStr);
+ EXPECT_EQ(root[2].type_index(), TypeIndex::kTVMFFIStr);
+ EXPECT_EQ(root[3].type_index(), TypeIndex::kTVMFFISmallBytes);
+ EXPECT_EQ(root[4].type_index(), TypeIndex::kTVMFFIBytes);
+
+ AnyArray unmatched = StructuralMap<order>(root, [](int64_t value) ->
Expected<Any> {
+ return Any(value);
+ }).template cast<AnyArray>();
+
+ EXPECT_TRUE(unmatched.same_as(root));
+ }
+
+ // Identity callbacks return the original shared Array for both
representations.
+ {
+ AnyArray root = MakeStringAndBytesLeaves();
+ AnyArray owner = root; //
NOLINT(performance-unnecessary-copy-initialization)
+ int string_callback_count = 0;
+ int bytes_callback_count = 0;
+
+ AnyArray identity = StructuralMap<order>(
+ root,
+ [&](const String& value) -> Expected<Any> {
+ ++string_callback_count;
+ return Any(value);
+ },
+ [&](const Bytes& value) -> Expected<Any> {
+ ++bytes_callback_count;
+ return Any(value);
+ })
+ .template cast<AnyArray>();
+
+ EXPECT_TRUE(identity.same_as(root));
+ EXPECT_TRUE(owner.same_as(root));
+ EXPECT_EQ(string_callback_count, 2);
+ EXPECT_EQ(bytes_callback_count, 2);
+ }
+
+ // Matching callbacks can replace both representations without traversing
into them.
+ {
+ AnyArray root = MakeStringAndBytesLeaves();
+ AnyArray replaced = StructuralMap<order>(
+ root,
+ [](const String& value) -> Expected<Any> {
+ return Any(static_cast<int64_t>(value.size()));
+ },
+ [](const Bytes& value) -> Expected<Any> {
+ return Any(static_cast<int64_t>(value.size()));
+ })
+ .template cast<AnyArray>();
+
+ EXPECT_EQ(replaced[0].cast<int64_t>(), 1);
+ EXPECT_EQ(replaced[1].cast<int64_t>(), 7);
+ EXPECT_EQ(replaced[2].cast<int64_t>(), 8);
+ EXPECT_EQ(replaced[3].cast<int64_t>(), 7);
+ EXPECT_EQ(replaced[4].cast<int64_t>(), 8);
+ }
+}
+
+TEST(StructuralMap, HandlesInlineAndHeapStringAndBytesLeaves) {
+ CheckStringAndBytesLeaves<WalkOrder::kPreOrder>();
+ CheckStringAndBytesLeaves<WalkOrder::kPostOrder>();
+}
+
+} // namespace
diff --git a/tests/cpp/extra/test_structural_visit.cc
b/tests/cpp/extra/test_structural_visit.cc
index 6bff25ad..6ac34a43 100644
--- a/tests/cpp/extra/test_structural_visit.cc
+++ b/tests/cpp/extra/test_structural_visit.cc
@@ -192,7 +192,7 @@ TEST(StructuralVisitor, TraversesArray) {
EXPECT_TRUE(AsTestVisitor(visitor)->visited[1].same_as(rhs));
}
-TEST(StructuralVisitor, TraversesMap) {
+TEST(StructuralVisitor, TraversesMapValuesWithoutVisitingKeys) {
ObjectRef key = TVar("key");
ObjectRef value = TVar("value");
Map<Any, Any> root{{key, value}};
@@ -202,9 +202,8 @@ TEST(StructuralVisitor, TraversesMap) {
ASSERT_TRUE(result.is_ok());
EXPECT_FALSE(result.value().has_value());
- ASSERT_EQ(AsTestVisitor(visitor)->visited.size(), 2U);
- EXPECT_TRUE(AsTestVisitor(visitor)->visited[0].same_as(key));
- EXPECT_TRUE(AsTestVisitor(visitor)->visited[1].same_as(value));
+ ASSERT_EQ(AsTestVisitor(visitor)->visited.size(), 1U);
+ EXPECT_TRUE(AsTestVisitor(visitor)->visited[0].same_as(value));
}
TEST(StructuralVisitor, UsesFuncHook) {
diff --git a/tests/python/test_structural.py b/tests/python/test_structural.py
index 68779788..51d4e6f7 100644
--- a/tests/python/test_structural.py
+++ b/tests/python/test_structural.py
@@ -20,6 +20,7 @@ from __future__ import annotations
from typing import Any
import numpy as np
+import pytest
import tvm_ffi
import tvm_ffi.testing
from tvm_ffi.dataclasses import Object, field, py_class
@@ -184,6 +185,7 @@ def test_structural_walk_typed_callbacks() -> None:
((int, float), lambda value: trace.append(f"number:{value}")),
(str, lambda value: trace.append(f"str:{value}")),
],
+ order=tvm_ffi.WalkOrder.PREORDER,
)
assert result is None
@@ -232,6 +234,7 @@ def test_structural_walk_first_match_and_skip() -> None:
),
(object, lambda value: trace.append(type(value).__name__)),
],
+ order=tvm_ffi.WalkOrder.PREORDER,
)
assert result is None
@@ -246,13 +249,17 @@ def test_structural_walk_interrupt() -> None:
return tvm_ffi.VisitInterrupt({"found": value})
return None
- result = tvm_ffi.structural_walk(root, (int, on_int))
+ result = tvm_ffi.structural_walk(
+ root,
+ (int, on_int),
+ order=tvm_ffi.WalkOrder.PREORDER,
+ )
assert isinstance(result, tvm_ffi.VisitInterrupt)
assert tvm_ffi.structural_equal(result.value, {"found": 2})
-def test_structural_walk_nested_containers() -> None:
+def test_structural_walk_nested_containers_and_skips_map_keys() -> None:
root = tvm_ffi.Array(
[
tvm_ffi.Map(
@@ -284,7 +291,7 @@ def test_structural_walk_nested_containers() -> None:
assert ("map", 2) in containers
assert ("dict", 1) in containers
assert sorted(scalars) == [1, 1, 2, 3]
- assert set(strings) == {"numbers", "meta", "flag"}
+ assert strings == []
def test_structural_walk_object_and_any_callbacks() -> None:
@@ -297,6 +304,7 @@ def test_structural_walk_object_and_any_callbacks() -> None:
(tvm_ffi.Object, lambda value:
trace.append(f"object:{type(value).__name__}")),
(Any, lambda value: trace.append(f"any:{value}")),
],
+ order=tvm_ffi.WalkOrder.PREORDER,
)
assert result is None
@@ -306,27 +314,51 @@ def test_structural_walk_object_and_any_callbacks() ->
None:
result = tvm_ffi.structural_walk(
tvm_ffi.Array([1]),
(object, lambda value: alias_trace.append(type(value).__name__)),
+ order=tvm_ffi.WalkOrder.PREORDER,
)
assert result is None
assert alias_trace == ["Array", "int"]
-def test_structural_walk_post_order_enum() -> None:
[email protected](
+ ("order", "expected_trace"),
+ [
+ pytest.param(
+ None,
+ ["int:1", "array:1", "int:2", "array:2"],
+ id="default-postorder",
+ ),
+ pytest.param(
+ tvm_ffi.WalkOrder.PREORDER,
+ ["array:2", "array:1", "int:1", "int:2"],
+ id="preorder",
+ ),
+ pytest.param(
+ tvm_ffi.WalkOrder.POSTORDER,
+ ["int:1", "array:1", "int:2", "array:2"],
+ id="postorder",
+ ),
+ ],
+)
+def test_structural_walk_pre_and_post_order(
+ order: tvm_ffi.WalkOrder | None,
+ expected_trace: list[str],
+) -> None:
root = tvm_ffi.Array([tvm_ffi.Array([1]), 2])
trace: list[str] = []
- result = tvm_ffi.structural_walk(
- root,
- [
- (tvm_ffi.Array, lambda value: trace.append(f"array:{len(value)}")),
- (int, lambda value: trace.append(f"int:{value}")),
- ],
- order=tvm_ffi.WalkOrder.POSTORDER,
- )
+ callbacks = [
+ (tvm_ffi.Array, lambda value: trace.append(f"array:{len(value)}")),
+ (int, lambda value: trace.append(f"int:{value}")),
+ ]
+ if order is None:
+ result = tvm_ffi.structural_walk(root, callbacks)
+ else:
+ result = tvm_ffi.structural_walk(root, callbacks, order=order)
assert result is None
- assert trace == ["int:1", "array:1", "int:2", "array:2"]
+ assert trace == expected_trace
def test_structural_walk_mixed_callback_forms() -> None:
@@ -360,7 +392,172 @@ def test_structural_walk_mixed_callback_forms() -> None:
),
),
],
+ order=tvm_ffi.WalkOrder.PREORDER,
)
assert result is None
assert trace == ["array:2", "array:1", "array:2", "use:x", "use:y",
"str:tag"]
+
+
+def test_structural_map_nested_array_map_order_and_keys() -> None:
+ def run(order: tvm_ffi.WalkOrder | None) -> list[str]:
+ root = tvm_ffi.Array([tvm_ffi.Map({"value": tvm_ffi.Array([1])})])
+ root_handle = root.__chandle__()
+ map_handle = root[0].__chandle__()
+ inner_array_handle = root[0]["value"].__chandle__()
+ trace: list[str] = []
+
+ def map_array(value: tvm_ffi.Array) -> tvm_ffi.Array:
+ trace.append("outer-array" if value.same_as(root) else
"inner-array")
+ return value
+
+ def map_map(value: tvm_ffi.Map) -> tvm_ffi.Map:
+ trace.append("map")
+ return value
+
+ def map_string(_: str) -> str:
+ trace.append("map-key")
+ return "renamed"
+
+ def map_int(value: int) -> int:
+ trace.append("int")
+ return value + 1
+
+ callbacks = [
+ (tvm_ffi.Array, map_array),
+ (tvm_ffi.Map, map_map),
+ (str, map_string),
+ (int, map_int),
+ ]
+ if order is None:
+ mapped = tvm_ffi.structural_map(root, callbacks)
+ else:
+ mapped = tvm_ffi.structural_map(root, callbacks, order=order)
+
+ assert mapped.__chandle__() == root_handle
+ assert mapped[0].__chandle__() == map_handle
+ assert mapped[0]["value"].__chandle__() == inner_array_handle
+ assert list(mapped[0]["value"]) == [2]
+ assert "value" in mapped[0]
+ assert "renamed" not in mapped[0]
+ return trace
+
+ assert run(tvm_ffi.WalkOrder.PREORDER) == [
+ "outer-array",
+ "map",
+ "inner-array",
+ "int",
+ ]
+ assert run(None) == ["int", "inner-array", "map", "outer-array"]
+
+
+def test_structural_map_array_ownership() -> None:
+ # A unique outer Array is reused, but its externally shared child is
copied.
+ shared_child = tvm_ffi.Array([1])
+ root = tvm_ffi.Array([shared_child])
+ root_handle = root.__chandle__()
+ mapped = tvm_ffi.structural_map(root, (int, lambda value: value + 1))
+
+ assert mapped.__chandle__() == root_handle
+ assert not mapped[0].same_as(shared_child)
+ assert list(shared_child) == [1]
+ assert list(mapped[0]) == [2]
+
+ # Sharing the outer Array preserves its complete original path.
+ shared_root = tvm_ffi.Array([tvm_ffi.Array([1])])
+ owner = tvm_ffi.Array([shared_root])
+ child_handle = shared_root[0].__chandle__()
+ mapped = tvm_ffi.structural_map(shared_root, (int, lambda value: value +
1))
+
+ assert not mapped.same_as(shared_root)
+ assert owner[0].same_as(shared_root)
+ assert mapped[0].__chandle__() != child_handle
+ assert list(shared_root[0]) == [1]
+ assert list(mapped[0]) == [2]
+
+
+def test_structural_map_map_value_ownership() -> None:
+ # A unique Map is reused, but its externally shared value is copied.
+ shared_value = tvm_ffi.Array([1])
+ root = tvm_ffi.Map({"value": shared_value})
+ root_handle = root.__chandle__()
+ mapped = tvm_ffi.structural_map(root, (int, lambda value: value + 1))
+
+ assert mapped.__chandle__() == root_handle
+ assert not mapped["value"].same_as(shared_value)
+ assert list(shared_value) == [1]
+ assert list(mapped["value"]) == [2]
+
+ # Sharing the Map copies both the Map and its changed value path.
+ shared_root = tvm_ffi.Map({"value": tvm_ffi.Array([1])})
+ owner = tvm_ffi.Array([shared_root])
+ value_handle = shared_root["value"].__chandle__()
+ mapped = tvm_ffi.structural_map(shared_root, (int, lambda value: value +
1))
+
+ assert not mapped.same_as(shared_root)
+ assert owner[0].same_as(shared_root)
+ assert mapped["value"].__chandle__() != value_handle
+ assert list(shared_root["value"]) == [1]
+ assert list(mapped["value"]) == [2]
+
+ # A shared Map is returned directly when no value changes.
+ shared_root = tvm_ffi.Map({"value": tvm_ffi.Array([1])})
+ owner = tvm_ffi.Array([shared_root])
+ mapped = tvm_ffi.structural_map(shared_root, (int, lambda value: value))
+
+ assert mapped.same_as(shared_root)
+ assert owner[0].same_as(shared_root)
+ assert mapped["value"].same_as(shared_root["value"])
+
+
+def test_structural_map_reuses_var_and_dag_callback_results() -> None:
+ @py_class(structural_eq="var")
+ class PyMapVar(Object):
+ value: int = field(structural_eq="ignore")
+
+ @py_class(structural_eq="dag")
+ class PyMapDAG(Object):
+ value: int
+
+ for order in (tvm_ffi.WalkOrder.PREORDER, tvm_ffi.WalkOrder.POSTORDER):
+ for node_type in (PyMapVar, PyMapDAG):
+ node = node_type(1)
+ root = tvm_ffi.Array([node, tvm_ffi.Map({"use": node})])
+ callback_count = 0
+
+ def replace(value: Any) -> Any:
+ nonlocal callback_count
+ callback_count += 1
+ return node_type(value.value + 1)
+
+ mapped = tvm_ffi.structural_map(root, (node_type, replace),
order=order)
+
+ assert callback_count == 1
+ assert mapped[0].same_as(mapped[1]["use"])
+ assert not mapped[0].same_as(node)
+ assert mapped[0].value == 2
+
+
+def test_structural_map_handles_inline_and_heap_strings_and_bytes() -> None:
+ values = [1, "1234567", "12345678", b"1234567", b"12345678"]
+
+ for order in (tvm_ffi.WalkOrder.PREORDER, tvm_ffi.WalkOrder.POSTORDER):
+ root = tvm_ffi.Array(values)
+ unmatched = tvm_ffi.structural_map(root, (int, lambda value: value),
order=order)
+ assert unmatched.same_as(root)
+ assert list(unmatched) == values
+
+ root = tvm_ffi.Array(values)
+ owner = tvm_ffi.Array([root])
+ identity = tvm_ffi.structural_map(
+ root,
+ [(str, lambda value: value), (bytes, lambda value: value)],
+ order=order,
+ )
+ assert identity.same_as(root)
+ assert owner[0].same_as(root)
+ assert list(identity) == values
+
+ root = tvm_ffi.Array(values)
+ replaced = tvm_ffi.structural_map(root, [(str, len), (bytes, len)],
order=order)
+ assert list(replaced) == [1, 7, 8, 7, 8]
