Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 148 additions & 0 deletions docs/guides/rust_lang_guide.md
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,154 @@ fn may_fail(value: i32) -> Result<()> {
}
```

### Structural Walk and Visit

Rust provides equivalents of the C++ `StructuralWalk`/`StructuralVisitor`
APIs. Put `#[dispatch(visit)]` on an impl to turn its `visit_*` methods into
typed handlers, then pass it to `structural_walk`; each handler returns a
`WalkResult` (`Advance`, `Skip`, or `Interrupt`) to steer the traversal.
Handlers dispatch on their argument type and may take an optional trailing
`DefRegionKind` argument:

```rust
use tvm_ffi::{dispatch, structural_walk, Array, DefRegionKind, WalkOrder, WalkResult};

#[derive(Default)]
struct Probe {
total: i64,
floats: usize,
}

#[dispatch(visit)]
impl Probe {
fn visit_integer(&mut self, value: i64) -> WalkResult {
self.total += value;
WalkResult::Advance
}

fn visit_float(&mut self, _value: f64, _kind: DefRegionKind) -> WalkResult {
self.floats += 1;
WalkResult::Advance
}
}

let values = Array::new(vec![1_i64, 2, 3]);
let mut probe = Probe::default();
structural_walk(&values, &mut probe, WalkOrder::PreOrder)?;
assert_eq!(probe.total, 6);
```

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
simply advance; a `&VisitValue` lambda acts as a catch-all and must come
last, since links after an always-matching one never run. Each lambda may
take a trailing `DefRegionKind` argument:

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

let values = Array::new(vec![1_i64, 2, 3]);

let mut total = 0;
structural_walk(
&values,
|value: i64| {
total += value;
WalkResult::Advance
},
WalkOrder::PreOrder,
)?;
assert_eq!(total, 6);

let mut evens = 0;
let mut objects = 0;
structural_walk(
&values,
(
|value: i64| {
if value % 2 == 0 {
evens += 1;
}
WalkResult::Advance
},
|_object: &Object, _kind: DefRegionKind| {
objects += 1;
WalkResult::Advance
},
),
WalkOrder::PreOrder,
)?;
assert_eq!((evens, objects), (1, 1));
```

Both entry points return `Result<Option<VisitInterrupt>>`: `Ok(None)` means
the whole graph was visited, and a handler stops the walk early by returning
`WalkResult::interrupt_with(payload)`, which comes back to the caller as
`Ok(Some(interrupt))`. Handlers may also return `Result<WalkResult>` and
propagate errors with `?`:

```rust
use tvm_ffi::{structural_walk, Array, WalkOrder, WalkResult};

let values = Array::new(vec![1_i64, 2, 3]);
let found = structural_walk(
&values,
|value: i64| {
if value == 2 {
return WalkResult::interrupt_with(value);
}
WalkResult::Advance
},
WalkOrder::PreOrder,
)?;
assert_eq!(found.map(|i| i64::try_from(i.value).unwrap()), Some(2));
```

To drive recursion yourself, implement `StructuralVisitor` and call
`structural_visit`; `visit` runs for each value and descends through
`default_visit_children`, or through `visit_child`, which visits one
selected child and can override the def-region state for it (e.g.
`DefRegionKind::Recursive` when descending into a binder's parameters):

```rust
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);
```

Two safety notes: mutable `List`/`Dict` contents are snapshotted before
callbacks run, so mutation during traversal cannot invalidate the walk; and
a non-container type with a foreign `__s_visit__` hook is rejected rather
than silently walked through reflection — visit such a type's children
explicitly from a `StructuralVisitor`, or skip it with a pre-order
`WalkResult::Skip`.

## Examples

The repository includes a complete example in `rust/tvm-ffi/examples/load_library.rs`.
Expand Down
9 changes: 9 additions & 0 deletions rust/tvm-ffi-macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,15 @@ use proc_macro_error::proc_macro_error;
mod match_any;
mod object_macros;
mod utils;
mod visit;

/// Generate typed structural-visit dispatch from the `visit_*` methods in an
/// inherent implementation.
#[proc_macro_error]
#[proc_macro_attribute]
pub fn dispatch(attr: TokenStream, item: TokenStream) -> TokenStream {
visit::dispatch(attr, item)
}

/// Match object-backed values carried by an Any-compatible scrutinee.
///
Expand Down
Loading