Skip to content

[FEAT][RUST] Add rust structural_visit and structural_walk. - #693

Merged
tqchen merged 20 commits into
apache:mainfrom
Seven-Streams:main-dev/2026-07-29/rust_visitor
Aug 3, 2026
Merged

[FEAT][RUST] Add rust structural_visit and structural_walk.#693
tqchen merged 20 commits into
apache:mainfrom
Seven-Streams:main-dev/2026-07-29/rust_visitor

Conversation

@Seven-Streams

@Seven-Streams Seven-Streams commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

A refreshed summary of the API as it now stands, since it has moved since the description above was written.

structural_walk(root, walker, order) walks a value graph and dispatches typed handlers in pre- or post-order; handlers steer the traversal through the returned WalkResult (Advance, Skip, Interrupt/InterruptWith(payload)) and may return Result<WalkResult> to propagate errors with ?. The walker can be:

  • a &mut #[dispatch(visit)] visitor (or a hand-written VisitDispatch) — its visit_* methods dispatch on argument type in source order;
  • a single closure over an FFI value type T (value cast), a borrowed object node &N (refcount-free subtype check), or the &VisitValue catch-all;
  • a tuple of up to 8 such callbacks — closures and &mut visitors mixed freely — tried in order with the first matching argument type winning, the analog of the variadic C++ StructuralWalk(root, callbacks...) chain.

Every handler shape may declare a trailing DefRegionKind argument to receive the definition-region state:

use tvm_ffi::{structural_walk, Array, DefRegionKind, Object, WalkOrder, WalkResult};

let values = Array::new(vec![10_i64, 2]);
let mut total = 0_i64;
let mut objects = 0;
assert!(structural_walk(
    &values,
    (
        |value: i64| {
            total += value;
            WalkResult::Advance
        },
        |_object: &Object, _kind: DefRegionKind| {
            objects += 1;
            WalkResult::Advance
        },
    ),
    WalkOrder::PreOrder,
)
.unwrap()
.is_none());
assert_eq!((total, objects), (12, 1));

structural_visit(root, visitor) drives a hand-written StructuralVisitor when recursion itself is part of the analysis, mirroring a C++ StructuralVisitorObj: visit runs for each value and controls all descent through default_visit_children or visit_child(child, kind), whose explicit kind argument plays the role of WithDefRegionKind. Both entry points return Result<Option<VisitInterrupt>>: Ok(None) means the whole graph was visited, and an interrupting handler's payload comes back in Ok(Some(interrupt)):

use tvm_ffi::{
    structural_visit, Array, DefRegionKind, Result, StructuralVisitor, VisitInterrupt, VisitValue,
};

#[derive(Default)]
struct Depth {
    max: usize,
    current: usize,
}

impl StructuralVisitor for Depth {
    fn visit(
        &mut self,
        value: &VisitValue,
        def_region_kind: DefRegionKind,
    ) -> Result<Option<VisitInterrupt>> {
        self.current += 1;
        self.max = self.max.max(self.current);
        let interrupt = self.default_visit_children(value, def_region_kind)?;
        self.current -= 1;
        Ok(interrupt)
    }
}

let values = Array::new(vec![1_i64, 2]);
let mut depth = Depth::default();
structural_visit(&values, &mut depth)?;
assert_eq!(depth.max, 2);

Performance

Re-measured on the current branch with the shared benchmark: Rust structural_visit (pre-order) and structural_walk (post-order) against the equivalent C++ StructuralWalk traversals on identical object graphs, all language x style combinations checksum-asserted before timing. Medians of 30 pinned reps.

Rust / C++ time ratio (lower is better; < 1 means Rust is faster):

shape structural_visit structural_walk
reflected-object tree (196k values) 0.98x 1.04x
nested arrays (87k values) 1.10x 1.16x
map trees, small & dense layouts (31k values) 0.90x 0.96x

Every combination is within 1.16x of C++, at an absolute cost of ~7-21 ns per visited value, matching C++.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@Seven-Streams
Seven-Streams force-pushed the main-dev/2026-07-29/rust_visitor branch 2 times, most recently from 600aa48 to 3163906 Compare July 30, 2026 17:43
@Seven-Streams
Seven-Streams force-pushed the main-dev/2026-07-29/rust_visitor branch 2 times, most recently from baca00f to d8a7452 Compare July 31, 2026 01:10
@tlopex
tlopex force-pushed the main-dev/2026-07-29/rust_visitor branch from 8c2314f to bc5d8f9 Compare July 31, 2026 01:37
@Seven-Streams
Seven-Streams force-pushed the main-dev/2026-07-29/rust_visitor branch 4 times, most recently from 1ee5724 to 00a3229 Compare August 1, 2026 02:19
Seven-Streams and others added 10 commits August 2, 2026 12:44
Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>
Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>

optim the effciciency.

Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>
Signed-off-by: tlopex <820958424@qq.com>
Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>
@Seven-Streams
Seven-Streams force-pushed the main-dev/2026-07-29/rust_visitor branch from 00a3229 to 65c5a68 Compare August 2, 2026 16:46
Seven-Streams and others added 6 commits August 2, 2026 12:53
Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>
Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>
Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>
Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>
The walk-chain tests asserted lossless numeric matching, but VisitValue::cast
matches on the FFI type tag and converts with `as` semantics. Remove the two
tests and fix the WalkChainLink rustdoc that claimed exact-value matching,
recommending i64/f64 links unless a deliberate narrowing is wanted.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…apache#693)

Callgrind on the nested-array shape showed two per-visit costs C++ does not
pay: reject_foreign_structural_visit stayed a real call (~20% of the
container fast path) because its cold error-formatting body defeated the
inline hint, and the split Result/WalkResult matches in visit_raw left a
partially-moved temporary whose drop glue could not fold away. Split the
cold body behind #[cold] #[inline(never)] and match the handler results
by value in one step. Nested-array walks drop from ~1.26-1.30x of C++ to
~1.09-1.17x; reflected-object and map shapes stay at or below parity.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Seven-Streams
Seven-Streams marked this pull request as ready for review August 2, 2026 22:28
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Comment thread rust/tvm-ffi/tests/test_walk_chain.rs Outdated
Comment thread rust/tvm-ffi/tests/test_walk_chain.rs Outdated
Seven-Streams and others added 4 commits August 2, 2026 18:36
Remove test_dispatch.rs (its downstream-path and cfg checks duplicate
coverage in test_structural_visit.rs; the mixed-arity handler case moves
into GenericDispatchProbe as a trailing DefRegionKind argument) and drop
two closure-walk tests whose behavior is already asserted by the sequence
fallback, interrupt-payload, and error-context tests.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>
Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>
Signed-off-by: yuchuan <yuchuan.7streams@gmail.com>
@tqchen
tqchen merged commit b98be1b into apache:main Aug 3, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants