[FEAT] Add structural visitor and typed structural walk APIs

[Kathryn-cat]

Summary

This PR adds structural traversal APIs for TVM FFI values in C++ and Python.

The main user-facing addition is StructuralWalk / tvm_ffi.structural_walk, a pytree-style walk that can visit both object-backed nodes and POD leaves (int, float, bool, str, bytes). It supports typed callbacks, pre/post-order traversal, child skipping, early interruption, and def-region-aware callbacks.

Key Changes

• Added C++ StructuralVisitor, VisitInterrupt, WalkResult, and StructuralWalk.
• Added Python tvm_ffi.structural_walk.
• Traversal operates on AnyView, allowing callbacks to match objects, containers, and POD values.
• C++ callbacks dispatch on their first argument and may optionally accept a second TVMFFIDefRegionKind.
• Python callbacks are passed as (type, callback) entries, with support for grouped types like ((int, float), callback), a single callback catch-all, or a sequence of entries.
• Python def-region-aware callbacks are passed via with_def_region_kind= and receive (value, def_region_kind).
• Added runtime type-index matching for Python callbacks, including subclass matching, Any catch-all, Object matching, and string/bytes storage variants.
• Custom structural visit hooks now use AnyView for the visited value.
• Added visit error-context propagation for object-backed nodes when traversal fails.

Examples

Python:

visited = []

def on_int(value):
    visited.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")

Grouped callback types:

tvm_ffi.structural_walk(
    root,
    [
        ((int, float), on_number),
        (MyNode, on_node),
        (object, on_any),
    ],
)

Def-region-aware callbacks:

uses = []

tvm_ffi.structural_walk(
    func,
    with_def_region_kind=(
        Var,
        lambda var, kind: uses.append(var) if kind == tvm_ffi.DefRegionKind.NONE else None,
    ),
)

C++:

auto result = StructuralWalk<WalkOrder::kPreOrder>(
    root,
    [&](const Add& add) -> Expected<WalkResult> {
      ++num_adds;
      return WalkResult::Advance();
    },
    [&](const Mul& mul) -> Expected<WalkResult> {
      return WalkResult::Skip();
    });

C++ def-region-aware walk:

List<Var> uses;
auto result = StructuralWalk<WalkOrder::kPreOrder>(
    func,
    [&](const VarObj* var, TVMFFIDefRegionKind kind) -> Expected<WalkResult> {
      if (kind == kTVMFFIDefRegionKindNone) {
        uses.push_back(Var(GetObjectPtr<VarObj>(var)));
      }
      return WalkResult::Advance();
    });

Testing

• Added C++ tests for structural visitor traversal, def-region hooks, callback dispatch, POD leaves, object pointers, skip/interrupt behavior, def-region-aware callbacks, and error propagation.
• Added Python tests for typed callbacks, grouped callback types, mixed callback forms, def-region-aware callbacks, nested containers, Any/Object behavior, post-order traversal, skip, and interrupts.
• Verified C++ structural visitor and visit error-context tests.

[gemini-code-assist]

Code Review

This pull request introduces a structural visit and walk implementation for the TVM FFI, enabling traversal of object graphs with support for early interruption. It adds the StructuralVisitor and StructuralWalk APIs, implements built-in container visitors, and provides unsafe raw-storage helpers for Expected. The review feedback recommends adding a defensive check for undefined ObjectRef values in DefaultVisitExpected to prevent potential crashes, supporting void return types in StructuralWalk callbacks to reduce boilerplate, and adding a static assertion to verify that each callback takes exactly one argument to avoid obscure template compilation errors.

[tqchen]

This is likely not need atm, focus ontesting StructuralWalk

[Kathryn-cat]

I've significantly reduced the test_structural_visit.cc test cases, and moved objects to testing_object.h. I'd like to keep TestVisitorObj in the test file because it is useful for helping to check visited nodes, def region kinds, and interrupt behavior. This is a useful fixture I think.

[tqchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a structural walk and visit API to the TVM FFI, allowing users to traverse structured value graphs and invoke typed callbacks in both C++ and Python. It adds core classes and functions such as StructuralVisitor, VisitInterrupt, WalkOrder, WalkResult, and structural_walk, along with comprehensive unit tests and documentation. The review feedback highlights three improvement opportunities: refactoring the C++ macros TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN and TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT to prevent multiple evaluations of their arguments, and replacing assert statements in the Python structural_walk callback normalization with explicit type checks to ensure robust input validation.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Kathryn-cat]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a new structural_walk API and StructuralVisitor to the TVM FFI, enabling structural traversal of object graphs with support for custom callbacks and definition-region tracking. The implementation includes necessary C++ headers, source files, Python bindings, and comprehensive tests. A review comment suggests improving robustness by catching std::exception in the visitor's callback dispatcher to avoid potential process termination.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.