tvm_ffi.structural_walk

tvm_ffi.structural_walk#

tvm_ffi.structural_walk(root, callbacks=(), with_def_region_kind=(), order='post')#

Walk a value structurally and invoke the first matching typed callback.

The walk keeps no engine state and visits every occurrence as a tree. Deduplication can be composed with a pre-order callback whose own visited set returns WalkResult.SKIP on repeats.

At the outer error boundary, after recursive traversal collects visit context, errors are offered once to the original input root type’s __ffi_enrich_visit_error__(error, root) hook. Only the exact root dynamic type’s attribute is used; parent types are not searched. The hook receives the original input root and returns the original or enriched FFI error. An absent or failing hook leaves the original error unchanged. Recursive visitor/mutator methods only collect context.

Parameters:

  • root (Any) – Root value to traverse.

  • callbacks (tuple | Sequence | Callable) –

    Normal callbacks. These callbacks receive one argument: 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 (tuple | Sequence | Callable) – Def-region-aware callbacks. These callbacks receive two arguments: (value, def_region_kind). They accept the same callback entry forms as callbacks.

  • order (str | WalkOrder) – "post"/WalkOrder.POSTORDER (the default) to invoke callbacks after children, or "pre"/WalkOrder.PREORDER to invoke callbacks before children.

Returns:

result – None if traversal completed, otherwise a VisitInterrupt returned by a callback.

Return type:

VisitInterrupt | None

Examples

visited = []

uses = []
result = tvm_ffi.structural_walk(
    node,
    ((int, float), lambda value: visited.append(("leaf", value))),
    with_def_region_kind=(
        Var,
        lambda var, kind: (
            uses.append(var) if kind == tvm_ffi.DefRegionKind.NONE else None
        ),
    ),
)