Kathryn-cat commented on code in PR #601:
URL: https://github.com/apache/tvm-ffi/pull/601#discussion_r3384386888
##########
python/tvm_ffi/structural.py:
##########
@@ -231,3 +290,260 @@ def __hash__(self) -> int:
def __eq__(self, other: Any) -> bool:
"""Compare by structural equality."""
return isinstance(other, StructuralKey) and
_ffi_api.StructuralKeyEqual(self, other)
+
+
+@register_object("ffi.VisitInterrupt")
+class VisitInterrupt(Object):
+ """Payload-carrying signal that stops a structural visit.
+
+ This object can be returned from structural walk callbacks and structural
+ visit hooks to halt traversal early. The optional payload is preserved in
+ :py:attr:`value` and returned to the caller.
+
+ Examples
+ --------
+ Use ``VisitInterrupt`` to stop a structural walk when a target node is
found:
+
+ .. code-block:: python
+
+ import tvm_ffi
+
+
+ def on_node(node):
+ if is_target(node):
+ return tvm_ffi.VisitInterrupt(node)
+ return None
+
+
+ result = tvm_ffi.structural_walk(root, (object, on_node))
+ if result is not None:
+ found = result.value
+
+ See Also
+ --------
+ :py:func:`tvm_ffi.structural_walk`
+ Structural walk API whose callbacks may return ``VisitInterrupt``.
+
+ """
+
+ # tvm-ffi-stubgen(begin): object/ffi.VisitInterrupt
+ # fmt: off
+ value: Any
+ if TYPE_CHECKING:
+ def __init__(self, value: Any = ...) -> None: ...
+ def __ffi_init__(self, _0: Any, /) -> None: ... # ty:
ignore[invalid-method-override]
+ # fmt: on
+ # tvm-ffi-stubgen(end)
+
+ def __init__(self, value: Any = None) -> None:
+ """Create an interrupt with an optional payload.
+
+ Parameters
+ ----------
+ value
+ Payload returned to the caller when traversal stops.
+
+ """
+ self.__init_handle_by_constructor__(_ffi_api.VisitInterrupt, value)
+
+
+@register_object("ffi.StructuralVisitor")
+class StructuralVisitor(Object):
+ """Low-level structural traversal visitor.
+
+ This class exposes the low-level visitor object used by structural
+ traversal hooks.
+ """
+
+ # tvm-ffi-stubgen(begin): object/ffi.StructuralVisitor
+ # fmt: off
+ if TYPE_CHECKING:
+ def __init__(self) -> None: ...
+ def __ffi_init__(self) -> None: ... # ty:
ignore[invalid-method-override]
+ # fmt: on
+ # tvm-ffi-stubgen(end)
+
+ def __init__(self) -> None:
+ """Create a default structural visitor."""
+ self.__init_handle_by_constructor__(_ffi_api.StructuralVisitor)
+
+ def visit(self, value: Any) -> VisitInterrupt | None:
+ """Low-level API to visit ``value`` using this visitor's dispatch
behavior.
+
+ Parameters
+ ----------
+ value
+ Value to visit.
+
+ Returns
+ -------
+ result
+ ``None`` if traversal should continue, otherwise a
+ :class:`VisitInterrupt` carrying the early-exit payload.
+
+ """
+ return _ffi_api.StructuralVisitorVisit(self, value)
+
+ def def_region_kind(self) -> DefRegionKind:
+ """Low-level API to return the currently active structural def-region
kind.
+
+ Returns
+ -------
+ kind
+ The active :class:`DefRegionKind`.
+
+ """
+ return DefRegionKind(_ffi_api.StructuralVisitorDefRegionKind(self))
+
+ def with_def_region_kind(
+ self,
+ kind: int,
+ callback: Callable[[], Any],
+ ) -> Any:
+ """Low-level API to 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.StructuralVisitorWithDefRegionKind(self, kind,
callback)
+
+
+def structural_walk(
+ root: Any,
+ *callbacks: tuple | Sequence | Callable,
+ with_def_region_kind: tuple | Sequence | Callable = (),
+ order: str | WalkOrder = "pre",
+) -> VisitInterrupt | None:
+ """Walk a value structurally and invoke the first matching typed callback.
+
+ Parameters
+ ----------
+ root
+ Root value to traverse.
+
+ callbacks
+ Normal callbacks. These callbacks receive one argument: ``value``.
+ Callback entries are tried in order.
+
+ Each positional callback slot 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 two arguments:
+ ``(value, def_region_kind)``. They accept the same callback entry forms
+ as ``callbacks``.
+
+ order
+ ``"pre"``/``WalkOrder.PREORDER`` to invoke callbacks before children,
or
+ ``"post"``/``WalkOrder.POSTORDER`` to invoke callbacks after children.
+
+ Returns
+ -------
+ result
+ ``None`` if traversal completed, otherwise a :class:`VisitInterrupt`
+ returned by a callback.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ 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
+ ),
+ ),
+ )
+
+ """
+ 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 walk order: {order!r}")
+
+ def normalize_callback_slots(
+ callback_slots: tuple[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)
+
+ for slot in callback_slots:
+ if callable(slot):
+ callback_entries.append((Any, slot))
+ elif isinstance(slot, tuple) and len(slot) == 2 and
callable(slot[1]):
+ add_callback_entry(slot)
+ elif isinstance(slot, Sequence):
+ for callback in slot:
+ assert isinstance(callback, tuple)
+ assert len(callback) == 2 and callable(callback[1])
+ add_callback_entry(callback)
Review Comment:
resolved
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
