From 36b4b09583db5f6656785c26da6c8d4cf30f4e23 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 05:24:17 +0000 Subject: [PATCH 1/5] Add walk_frontier and walk_frontier_edges Return only the descendants a NodeVisitor would dispatch on, in generic_visit order, so a visitor with a few visit_ methods can skip every other node without a Python loop. The edges variant carries the parent slot each match was read from, for a NodeTransformer. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hv7iDZqYjqWUscP5LcdwfB --- README.md | 17 ++++ fast_walk.pyi | 20 +++++ src/lib.rs | 194 ++++++++++++++++++++++++++++++++++++++++- tests/test_frontier.py | 112 ++++++++++++++++++++++++ 4 files changed, 342 insertions(+), 1 deletion(-) create mode 100644 tests/test_frontier.py diff --git a/README.md b/README.md index d0f5f96..91ba70a 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,23 @@ for node in walk_unordered(tree): depth-first pre-order visitation. `ast.walk` does not document an order, so most callers can safely use `walk_unordered`. +### Visiting only some node types + +`walk_frontier(node, kinds)` returns the descendants whose exact type is in +`kinds`, in depth-first pre-order, without descending below a match. That is +the set `ast.NodeVisitor.generic_visit` hands to `visit_` methods, so a +visitor whose only methods are `visit_` can descend with: + +```python +def generic_visit(self, node): + for child in walk_frontier(node, (ast.Call, ast.Attribute)): + self.visit(child) +``` + +`walk_frontier_edges` returns `(node, parent, field, index)` tuples instead +(`index` is `-1` for a non-list field), which is what a `NodeTransformer` +needs to put a replacement back. + ## Performance Benchmark on CPython 3.13, walking the AST of `difflib.py` (~2000 lines, diff --git a/fast_walk.pyi b/fast_walk.pyi index 8dd5e38..36e9b00 100644 --- a/fast_walk.pyi +++ b/fast_walk.pyi @@ -23,6 +23,26 @@ def walk_unordered(node: ast.AST) -> list[ast.AST]: 25% faster than :func:`walk_dfs` on real Python source. """ +def walk_frontier( + node: ast.AST, kinds: tuple[type[ast.AST], ...] +) -> list[ast.AST]: + """Return the descendants of `node` whose exact type is in `kinds`, in + depth-first pre-order, without descending below any match. `node` + itself is never included. + + These are the nodes ``ast.NodeVisitor.generic_visit`` dispatches to a + ``visit_`` method when every other node falls through to the + generic descent. + """ + +def walk_frontier_edges( + node: ast.AST, kinds: tuple[type[ast.AST], ...] +) -> list[tuple[ast.AST, ast.AST, str, int]]: + """Like :func:`walk_frontier`, but each match comes as + ``(node, parent, field, index)``: it was read from ``parent.`` + when ``index`` is ``-1``, else from ``parent.[index]``. + """ + def walk(node: ast.AST) -> list[ast.AST]: """Deprecated. Use :func:`walk_dfs` for explicit depth-first order or :func:`walk_unordered` for the faster order-agnostic variant. diff --git a/src/lib.rs b/src/lib.rs index 9837ae8..db195d2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,7 +17,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use pyo3::exceptions::PyDeprecationWarning; use pyo3::ffi::{self, PyListObject, PyObject, PyTypeObject}; -use pyo3::types::{PyList, PyModule, PyType}; +use pyo3::types::{PyList, PyModule, PyTuple, PyType}; use pyo3::{PyTypeInfo, prelude::*}; /// Open-addressed, direct-mapped lookup from `*mut PyTypeObject` to an @@ -514,6 +514,196 @@ fn walk_unordered<'py>(py: Python<'py>, node: Bound<'py, PyAny>) -> PyResult` when `index < 0`, else `parent.[index]`. +#[derive(Clone, Copy)] +struct Edge { + node: *mut PyObject, + parent: *mut PyObject, + key: *mut PyObject, + index: ffi::Py_ssize_t, +} + +/// Push every AST child of `parent` onto `stack` in reverse `_fields` +/// order, recording the slot each was read from. Same field discovery as +/// `process_node`, keeping the dict key alongside each value. +#[inline(always)] +unsafe fn push_child_edges( + parent: *mut PyObject, + base_ast_and_expr_type: (*mut PyTypeObject, *mut PyTypeObject), + py_list_type: *mut PyTypeObject, + field_table: &FieldTable, + stack: &mut Vec, +) { + let encoded = field_table.lookup(unsafe { ffi::Py_TYPE(parent) }); + if encoded <= 1 { + return; + } + let Some(dict) = get_instance_dict_fast(parent) else { + return; + }; + unsafe { + let keys = &*(*dict.cast::()) + .ma_keys + .cast::(); + let entries = keys.unicode_entries(); + let mut current = (keys.dk_nentries as usize).min((encoded - 1) as usize); + while current > 0 { + current -= 1; + let entry = *entries.add(current); + let value = entry.me_value; + if value.is_null() { + continue; + } + let value_type = ffi::Py_TYPE(value); + if value_type == py_list_type { + let list = value as *mut PyListObject; + let length = (*(list as *mut ffi::PyVarObject)).ob_size; + let ob_item = (*list).ob_item; + for index in (0..length).rev() { + let child = *ob_item.offset(index); + if issubclass_of_ast(ffi::Py_TYPE(child), base_ast_and_expr_type) { + stack.push(Edge { + node: child, + parent, + key: entry.me_key, + index, + }); + } + } + } else if issubclass_of_ast(value_type, base_ast_and_expr_type) { + stack.push(Edge { + node: value, + parent, + key: entry.me_key, + index: -1, + }); + } + } + } +} + +/// Descendants of `root` (excluding `root`) whose exact type is in `kinds`, +/// in depth-first pre-order, without descending below any of them. This is +/// the set of nodes `ast.NodeVisitor.generic_visit` would dispatch to a +/// `visit_` method, given that every other node falls through to the +/// generic descent. +fn walk_node_frontier( + root: *mut PyObject, + kinds: &[*mut PyTypeObject], + base_ast_and_expr_type: (*mut PyTypeObject, *mut PyTypeObject), + py_list_type: *mut PyTypeObject, + field_table: &FieldTable, + result: &mut Vec, +) { + let mut stack = Vec::new(); + unsafe { + push_child_edges( + root, + base_ast_and_expr_type, + py_list_type, + field_table, + &mut stack, + ); + } + while let Some(edge) = stack.pop() { + if kinds.contains(&unsafe { ffi::Py_TYPE(edge.node) }) { + result.push(edge); + } else { + unsafe { + push_child_edges( + edge.node, + base_ast_and_expr_type, + py_list_type, + field_table, + &mut stack, + ); + } + } + } +} + +fn type_pointers(kinds: &Bound<'_, PyTuple>) -> PyResult> { + kinds + .iter() + .map(|kind| Ok(kind.cast_into::()?.as_type_ptr())) + .collect() +} + +fn frontier_edges( + py: Python<'_>, + node: &Bound<'_, PyAny>, + kinds: &Bound<'_, PyTuple>, +) -> PyResult> { + let base = resolve_base_types(py)?; + let py_list_type = PyList::type_object_raw(py); + let kinds = type_pointers(kinds)?; + with_field_table(py, |table| { + let mut result = Vec::new(); + walk_node_frontier( + node.as_ptr(), + &kinds, + base, + py_list_type, + table, + &mut result, + ); + Ok(result) + }) +} + +/// Return the descendants of `node` whose exact type is one of `kinds`, in +/// depth-first pre-order, without descending below any match. `node` +/// itself is never included. +#[pyfunction] +fn walk_frontier<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, + kinds: Bound<'py, PyTuple>, +) -> PyResult> { + let edges = frontier_edges(py, &node, &kinds)?; + let nodes: Vec<*mut PyObject> = edges.iter().map(|edge| edge.node).collect(); + vec_into_pylist(py, &nodes) +} + +/// Like `walk_frontier`, but each match comes as `(node, parent, field, +/// index)`: the node was read from `parent.` when `index` is `-1`, +/// else from `parent.[index]`. +#[pyfunction] +fn walk_frontier_edges<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, + kinds: Bound<'py, PyTuple>, +) -> PyResult> { + let edges = frontier_edges(py, &node, &kinds)?; + unsafe { + let list_ptr = ffi::PyList_New(edges.len() as ffi::Py_ssize_t); + if list_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + let list = Bound::from_owned_ptr(py, list_ptr); + let ob_item = (*(list_ptr as *mut ffi::PyListObject)).ob_item; + for (position, edge) in edges.iter().enumerate() { + let index = ffi::PyLong_FromSsize_t(edge.index); + if index.is_null() { + return Err(PyErr::fetch(py)); + } + let item = ffi::PyTuple_New(4); + if item.is_null() { + ffi::Py_DECREF(index); + return Err(PyErr::fetch(py)); + } + for (slot, value) in [edge.node, edge.parent, edge.key].into_iter().enumerate() { + ffi::Py_INCREF(value); + ffi::PyTuple_SET_ITEM(item, slot as ffi::Py_ssize_t, value); + } + ffi::PyTuple_SET_ITEM(item, 3, index); + *ob_item.add(position) = item; + } + Ok(list) + } +} + static DEPRECATED_WALK_WARNED: AtomicBool = AtomicBool::new(false); /// Deprecated. Use `walk_dfs` for explicit depth-first order or @@ -556,6 +746,8 @@ fn fast_walk(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(walk, m)?)?; m.add_function(wrap_pyfunction!(walk_dfs, m)?)?; m.add_function(wrap_pyfunction!(walk_unordered, m)?)?; + m.add_function(wrap_pyfunction!(walk_frontier, m)?)?; + m.add_function(wrap_pyfunction!(walk_frontier_edges, m)?)?; m.add_function(wrap_pyfunction!(_walk_count, m)?)?; Ok(()) } diff --git a/tests/test_frontier.py b/tests/test_frontier.py new file mode 100644 index 0000000..1eede54 --- /dev/null +++ b/tests/test_frontier.py @@ -0,0 +1,112 @@ +"""`walk_frontier` and `walk_frontier_edges` must match a pure-Python reference: +the nodes `ast.NodeVisitor.generic_visit` dispatches to `visit_`, given +every other node falls through to the generic descent.""" + +from __future__ import annotations + +import ast +import gc +import sys +import textwrap + +import pytest + +from fast_walk import walk_frontier, walk_frontier_edges + +SOURCE = textwrap.dedent(""" + import x as y + def f(a, b=g(h.i)): + c = [k(l) for l in m.n.o] + if c: + return p.q(r, *s, t=u.v) + class W(X.Y): + z: Z = w.call() + lambda: q.attr +""") + +KINDS = [ + (ast.Call,), + (ast.Attribute,), + (ast.Call, ast.Attribute), + (ast.FunctionDef, ast.Name), + (ast.Load,), + (), +] + + +def reference(node: ast.AST, kinds: tuple[type, ...]): + found = [] + + def descend(parent: ast.AST): + for field, value in ast.iter_fields(parent): + items = value if isinstance(value, list) else [value] + for index, child in enumerate(items): + if not isinstance(child, ast.AST): + continue + if type(child) in kinds: + found.append(( + child, + parent, + field, + index if isinstance(value, list) else -1, + )) + else: + descend(child) + + descend(node) + return found + + +def identities(edges): + return [(id(n), id(p), f, i) for n, p, f, i in edges] + + +@pytest.mark.parametrize("kinds", KINDS) +def test_edges_match_the_generic_visit_order(kinds): + tree = ast.parse(SOURCE) + assert identities(walk_frontier_edges(tree, kinds)) == identities( + reference(tree, kinds) + ) + + +@pytest.mark.parametrize("kinds", KINDS) +def test_frontier_is_the_edge_nodes(kinds): + tree = ast.parse(SOURCE) + assert [id(n) for n in walk_frontier(tree, kinds)] == [ + id(edge[0]) for edge in walk_frontier_edges(tree, kinds) + ] + + +def test_the_root_is_never_included(): + tree = ast.parse("f(g())") + call = tree.body[0].value + assert [ast.unparse(n) for n in walk_frontier(call, (ast.Call,))] == ["g()"] + + +def test_a_match_is_not_descended(): + tree = ast.parse("a.b.c.d") + assert [ast.unparse(n) for n in walk_frontier(tree, (ast.Attribute,))] == [ + "a.b.c.d" + ] + + +def test_a_deep_chain_does_not_recurse(): + tree = ast.parse("x" + ".y" * 20_000) + assert [type(n) for n in walk_frontier(tree, (ast.Name,))] == [ast.Name] + + +def test_kinds_must_be_types(): + with pytest.raises(TypeError): + walk_frontier(ast.parse("x"), ("Name",)) + + +def test_refcount_neutral(): + tree = ast.parse(SOURCE) + nodes = list(ast.walk(tree)) + gc.collect() + before = [sys.getrefcount(n) for n in nodes] + for _ in range(50): + walk_frontier(tree, (ast.Call, ast.Attribute)) + walk_frontier_edges(tree, (ast.Call, ast.Attribute, ast.Name)) + gc.collect() + assert [sys.getrefcount(n) for n in nodes] == before From d88ee5bcedcb77397ef755aad31a8a2faa1f6ad8 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 17:25:00 +0000 Subject: [PATCH 2/5] Add walk_of_types, fix_missing_locations and subtree_hashes All three, and the frontier walks, read each child by its _fields name rather than by dict position, so they also handle nodes a transformer built with keyword arguments in another order. subtree_hashes mixes in each leaf value's type, so 1, 1.0 and True hash apart. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hv7iDZqYjqWUscP5LcdwfB --- README.md | 14 ++ fast_walk.pyi | 18 ++ src/fields.rs | 394 +++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 134 +------------- tests/test_fields.py | 158 +++++++++++++++++ tests/test_frontier.py | 7 + 6 files changed, 597 insertions(+), 128 deletions(-) create mode 100644 src/fields.rs create mode 100644 tests/test_fields.py diff --git a/README.md b/README.md index 91ba70a..b2988f9 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,20 @@ def generic_visit(self, node): (`index` is `-1` for a non-list field), which is what a `NodeTransformer` needs to put a replacement back. +### Other helpers + +- `walk_of_types(node, kinds)` — every node of the given exact types, in + pre-order, descending into matches. +- `fix_missing_locations(node)` — `ast.fix_missing_locations` without + recursion. +- `subtree_hashes(node)` — `{id(n): hash}` such that equal hashes mean equal + `ast.dump` output (leaf types included, so `1` and `True` differ). + +These, and both frontier walks, read each child by its `_fields` name, so +they also work on nodes a transformer built with keyword arguments in +another order. `walk_dfs` and `walk_unordered` read fields by position, +which is only guaranteed for parsed trees. + ## Performance Benchmark on CPython 3.13, walking the AST of `difflib.py` (~2000 lines, diff --git a/fast_walk.pyi b/fast_walk.pyi index 36e9b00..fc8bbe8 100644 --- a/fast_walk.pyi +++ b/fast_walk.pyi @@ -43,6 +43,24 @@ def walk_frontier_edges( when ``index`` is ``-1``, else from ``parent.[index]``. """ +def walk_of_types( + node: ast.AST, kinds: tuple[type[ast.AST], ...] +) -> list[ast.AST]: + """Return every node under `node` (`node` included) whose exact type is + in `kinds`, in depth-first pre-order. Unlike :func:`walk_frontier`, + matches are descended into. + """ + +def fix_missing_locations[T: ast.AST](node: T) -> T: + """``ast.fix_missing_locations`` without recursing once per tree level.""" + +def subtree_hashes(node: ast.AST) -> dict[int, int]: + """A structural hash per node under `node`, keyed by ``id(node)``. + + Subtrees with equal ``ast.dump`` output get equal hashes. A leaf value's + type is mixed in, so ``1``, ``1.0`` and ``True`` hash apart. + """ + def walk(node: ast.AST) -> list[ast.AST]: """Deprecated. Use :func:`walk_dfs` for explicit depth-first order or :func:`walk_unordered` for the faster order-agnostic variant. diff --git a/src/fields.rs b/src/fields.rs new file mode 100644 index 0000000..f0c3b20 --- /dev/null +++ b/src/fields.rs @@ -0,0 +1,394 @@ +//! Traversals that read a node's children by its type's `_fields` names. +//! +//! `walk_dfs` reads the first `len(_fields)` entries of the instance dict, +//! which holds for a parsed tree. A transformer can build a node whose dict +//! has another order (`ast.Call(lineno=1, func=f)` stores `lineno` first), +//! and these functions run on such trees, so each field is looked up by +//! name instead. + +use std::cell::RefCell; +use std::collections::HashMap; +use std::hash::{BuildHasherDefault, Hasher}; + +use pyo3::ffi::{self, PyObject, PyTypeObject}; +use pyo3::prelude::*; +use pyo3::types::{PyList, PyTuple, PyType}; +use pyo3::{PyTypeInfo, intern}; + +use crate::{Edge, get_instance_dict_fast, issubclass_of_ast, resolve_base_types, vec_into_pylist}; + +#[derive(Default)] +struct PointerHasher(u64); + +impl Hasher for PointerHasher { + fn finish(&self) -> u64 { + self.0 + } + + fn write(&mut self, bytes: &[u8]) { + for &byte in bytes { + self.0 = (self.0 << 8 | byte as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15); + } + } + + fn write_usize(&mut self, value: usize) { + self.0 = ((value as u64) >> 4).wrapping_mul(0x9E37_79B9_7F4A_7C15); + } +} + +type PointerMap = HashMap>; + +const LINENO: u8 = 1; +const COL_OFFSET: u8 = 2; +const END_LINENO: u8 = 4; +const END_COL_OFFSET: u8 = 8; + +/// An AST type's `_fields` names, and which location `_attributes` it has. +struct TypeInfo { + fields: Vec>, + locations: u8, +} + +thread_local! { + static TYPE_INFO: RefCell>> = RefCell::new(PointerMap::default()); +} + +fn names_of(py: Python<'_>, kind: &Bound<'_, PyAny>, attribute: &str) -> Vec> { + let Ok(names) = kind.getattr(attribute) else { + return Vec::new(); + }; + let Ok(names) = names.cast_into::() else { + return Vec::new(); + }; + names + .iter() + .filter_map(|name| { + let name = name.cast_into::().ok()?; + let interned = pyo3::types::PyString::intern(py, name.to_str().ok()?); + Some(interned.into_any().unbind()) + }) + .collect() +} + +/// The cached `TypeInfo` for `kind`. Entries are boxed and never removed, +/// so the pointer stays valid for the life of the thread. +fn type_info(py: Python<'_>, kind: *mut PyTypeObject) -> *const TypeInfo { + TYPE_INFO.with(|cache| { + let mut cache = cache.borrow_mut(); + let entry = cache.entry(kind as usize).or_insert_with(|| { + let kind = unsafe { Bound::from_borrowed_ptr(py, kind.cast::()) }; + let attributes = names_of(py, &kind, "_attributes"); + let mut locations = 0; + for name in &attributes { + locations |= match name.bind(py).extract::<&str>().unwrap_or_default() { + "lineno" => LINENO, + "col_offset" => COL_OFFSET, + "end_lineno" => END_LINENO, + "end_col_offset" => END_COL_OFFSET, + _ => 0, + }; + } + Box::new(TypeInfo { + fields: names_of(py, &kind, "_fields"), + locations, + }) + }); + &**entry as *const TypeInfo + }) +} + +/// Call `visit(field_index, value)` for each field `node` has set, in +/// `_fields` order. Values are borrowed from the node's dict. +unsafe fn for_each_field( + py: Python<'_>, + node: *mut PyObject, + info: &TypeInfo, + mut visit: impl FnMut(usize, *mut PyObject) -> PyResult<()>, +) -> PyResult<()> { + let Some(dict) = get_instance_dict_fast(node) else { + return Ok(()); + }; + for (index, name) in info.fields.iter().enumerate() { + let value = unsafe { ffi::PyDict_GetItemWithError(dict, name.as_ptr()) }; + if value.is_null() { + if unsafe { !ffi::PyErr_Occurred().is_null() } { + return Err(PyErr::fetch(py)); + } + continue; + } + visit(index, value)?; + } + Ok(()) +} + +struct Walk<'py> { + py: Python<'py>, + base: (*mut PyTypeObject, *mut PyTypeObject), + list_type: *mut PyTypeObject, +} + +impl<'py> Walk<'py> { + fn new(py: Python<'py>) -> PyResult { + Ok(Self { + py, + base: resolve_base_types(py)?, + list_type: PyList::type_object_raw(py), + }) + } + + fn is_ast(&self, value: *mut PyObject) -> bool { + issubclass_of_ast(unsafe { ffi::Py_TYPE(value) }, self.base) + } + + /// Push the AST children of `node` onto `stack` so they pop in + /// `ast.iter_child_nodes` order. + fn push_children(&self, node: *mut PyObject, stack: &mut Vec<*mut PyObject>) -> PyResult<()> { + let info = unsafe { &*type_info(self.py, ffi::Py_TYPE(node)) }; + let start = stack.len(); + unsafe { + for_each_field(self.py, node, info, |_, value| { + if ffi::Py_TYPE(value) == self.list_type { + let list = value as *mut ffi::PyListObject; + let length = (*(list as *mut ffi::PyVarObject)).ob_size; + for index in 0..length { + let item = *(*list).ob_item.offset(index); + if self.is_ast(item) { + stack.push(item); + } + } + } else if self.is_ast(value) { + stack.push(value); + } + Ok(()) + })?; + } + stack[start..].reverse(); + Ok(()) + } + + /// Every node under `root`, `root` included, in depth-first pre-order. + fn pre_order(&self, root: *mut PyObject) -> PyResult> { + let mut result = Vec::new(); + let mut stack = vec![root]; + while let Some(node) = stack.pop() { + result.push(node); + self.push_children(node, &mut stack)?; + } + Ok(result) + } +} + +fn type_pointers(kinds: &Bound<'_, PyTuple>) -> PyResult> { + kinds + .iter() + .map(|kind| Ok(kind.cast_into::()?.as_type_ptr())) + .collect() +} + +/// Return every node under `node` (`node` included) whose exact type is in +/// `kinds`, in depth-first pre-order. Unlike `walk_frontier`, matches are +/// descended into. +#[pyfunction] +pub fn walk_of_types<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, + kinds: Bound<'py, PyTuple>, +) -> PyResult> { + let kinds = type_pointers(&kinds)?; + let walk = Walk::new(py)?; + let mut result = Vec::new(); + let mut stack = vec![node.as_ptr()]; + while let Some(current) = stack.pop() { + if kinds.contains(&unsafe { ffi::Py_TYPE(current) }) { + result.push(current); + } + walk.push_children(current, &mut stack)?; + } + vec_into_pylist(py, &result) +} + +/// `ast.fix_missing_locations`, without recursing per tree level. Returns +/// `node`. +#[pyfunction] +pub fn fix_missing_locations<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, +) -> PyResult> { + let walk = Walk::new(py)?; + let names = [ + (LINENO, intern!(py, "lineno"), false), + (COL_OFFSET, intern!(py, "col_offset"), false), + (END_LINENO, intern!(py, "end_lineno"), true), + (END_COL_OFFSET, intern!(py, "end_col_offset"), true), + ]; + let one = 1i64.into_pyobject(py)?.into_any().unbind(); + let zero = 0i64.into_pyobject(py)?.into_any().unbind(); + let mut locations: Vec<[Py; 4]> = + vec![[one.clone_ref(py), zero.clone_ref(py), one, zero]]; + let mut stack: Vec<(*mut PyObject, usize)> = vec![(node.as_ptr(), 0)]; + let mut children = Vec::new(); + while let Some((current, inherited)) = stack.pop() { + let info = unsafe { &*type_info(py, ffi::Py_TYPE(current)) }; + let mut own = inherited; + if info.locations != 0 { + let dict = get_instance_dict_fast(current); + let mut values: [Py; 4] = + locations[inherited].each_ref().map(|v| v.clone_ref(py)); + for (slot, (flag, name, none_is_missing)) in names.iter().enumerate() { + if info.locations & flag == 0 { + continue; + } + let found = match dict { + Some(dict) => unsafe { ffi::PyDict_GetItemWithError(dict, name.as_ptr()) }, + None => std::ptr::null_mut(), + }; + if found.is_null() && unsafe { !ffi::PyErr_Occurred().is_null() } { + return Err(PyErr::fetch(py)); + } + let missing = + found.is_null() || (*none_is_missing && unsafe { found == ffi::Py_None() }); + if missing { + unsafe { Bound::from_borrowed_ptr(py, current) } + .setattr(name, values[slot].bind(py))?; + } else { + values[slot] = unsafe { Bound::from_borrowed_ptr(py, found) }.unbind(); + } + } + locations.push(values); + own = locations.len() - 1; + } + children.clear(); + walk.push_children(current, &mut children)?; + stack.extend(children.iter().map(|&child| (child, own))); + } + Ok(node) +} + +#[inline(always)] +fn mix(hash: u64, value: u64) -> u64 { + (hash.rotate_left(5) ^ value).wrapping_mul(0x9E37_79B9_7F4A_7C15) +} + +/// A structural hash per node, keyed by `id(node)`. Two subtrees that would +/// give the same `ast.dump` get the same hash; a leaf value's type is mixed +/// in, so `1`, `1.0` and `True` (equal hashes in Python) stay apart. +#[pyfunction] +pub fn subtree_hashes<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, +) -> PyResult> { + let walk = Walk::new(py)?; + let nodes = walk.pre_order(node.as_ptr())?; + let mut hashes: PointerMap = PointerMap::default(); + hashes.reserve(nodes.len()); + let leaf = |value: *mut PyObject, hashes: &PointerMap| -> PyResult { + if walk.is_ast(value) { + return Ok(hashes[&(value as usize)]); + } + let hashed = unsafe { ffi::PyObject_Hash(value) }; + if hashed == -1 && unsafe { !ffi::PyErr_Occurred().is_null() } { + return Err(PyErr::fetch(py)); + } + Ok(mix(unsafe { ffi::Py_TYPE(value) } as u64, hashed as u64)) + }; + for ¤t in nodes.iter().rev() { + if hashes.contains_key(&(current as usize)) { + continue; + } + let kind = unsafe { ffi::Py_TYPE(current) }; + let info = unsafe { &*type_info(py, kind) }; + let mut hash = mix(0, kind as u64); + unsafe { + for_each_field(py, current, info, |index, value| { + let value_hash = if ffi::Py_TYPE(value) == walk.list_type { + let list = value as *mut ffi::PyListObject; + let length = (*(list as *mut ffi::PyVarObject)).ob_size; + let mut items = mix(u64::MAX, length as u64); + for item in 0..length { + items = mix(items, leaf(*(*list).ob_item.offset(item), &hashes)?); + } + items + } else { + leaf(value, &hashes)? + }; + hash = mix(mix(hash, index as u64), value_hash); + Ok(()) + })?; + } + hashes.insert(current as usize, hash); + } + let result = pyo3::types::PyDict::new(py); + for (&pointer, &hash) in &hashes { + let key = unsafe { + Bound::from_owned_ptr_or_err( + py, + ffi::PyLong_FromVoidPtr(pointer as *mut std::ffi::c_void), + )? + }; + result.set_item(key, hash as i64)?; + } + Ok(result.into_any()) +} + +impl Walk<'_> { + /// Push the AST children of `parent` onto `stack` so they pop in + /// `ast.iter_child_nodes` order, each with the slot it was read from. + fn push_child_edges(&self, parent: *mut PyObject, stack: &mut Vec) -> PyResult<()> { + let info = unsafe { &*type_info(self.py, ffi::Py_TYPE(parent)) }; + let start = stack.len(); + unsafe { + for_each_field(self.py, parent, info, |index, value| { + let key = info.fields[index].as_ptr(); + if ffi::Py_TYPE(value) == self.list_type { + let list = value as *mut ffi::PyListObject; + let length = (*(list as *mut ffi::PyVarObject)).ob_size; + for item in 0..length { + let child = *(*list).ob_item.offset(item); + if self.is_ast(child) { + stack.push(Edge { + node: child, + parent, + key, + index: item, + }); + } + } + } else if self.is_ast(value) { + stack.push(Edge { + node: value, + parent, + key, + index: -1, + }); + } + Ok(()) + })?; + } + stack[start..].reverse(); + Ok(()) + } +} + +/// Descendants of `node` (excluding `node`) whose exact type is in `kinds`, +/// in depth-first pre-order, without descending below any of them. This is +/// the set `ast.NodeVisitor.generic_visit` dispatches to `visit_` +/// methods when every other node falls through to the generic descent. +pub(crate) fn frontier_edges( + py: Python<'_>, + node: &Bound<'_, PyAny>, + kinds: &Bound<'_, PyTuple>, +) -> PyResult> { + let kinds = type_pointers(kinds)?; + let walk = Walk::new(py)?; + let mut result = Vec::new(); + let mut stack = Vec::new(); + walk.push_child_edges(node.as_ptr(), &mut stack)?; + while let Some(edge) = stack.pop() { + if kinds.contains(&unsafe { ffi::Py_TYPE(edge.node) }) { + result.push(edge); + } else { + walk.push_child_edges(edge.node, &mut stack)?; + } + } + Ok(result) +} diff --git a/src/lib.rs b/src/lib.rs index db195d2..619eb6a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,6 +10,7 @@ //! `ast.walk` makes no ordering guarantee, so `walk_unordered` is a drop-in //! replacement wherever order doesn't matter. +mod fields; mod pydict; use std::cell::{Cell, RefCell}; @@ -18,6 +19,8 @@ use std::sync::atomic::{AtomicBool, Ordering}; use pyo3::exceptions::PyDeprecationWarning; use pyo3::ffi::{self, PyListObject, PyObject, PyTypeObject}; use pyo3::types::{PyList, PyModule, PyTuple, PyType}; + +use fields::frontier_edges; use pyo3::{PyTypeInfo, prelude::*}; /// Open-addressed, direct-mapped lookup from `*mut PyTypeObject` to an @@ -524,134 +527,6 @@ struct Edge { index: ffi::Py_ssize_t, } -/// Push every AST child of `parent` onto `stack` in reverse `_fields` -/// order, recording the slot each was read from. Same field discovery as -/// `process_node`, keeping the dict key alongside each value. -#[inline(always)] -unsafe fn push_child_edges( - parent: *mut PyObject, - base_ast_and_expr_type: (*mut PyTypeObject, *mut PyTypeObject), - py_list_type: *mut PyTypeObject, - field_table: &FieldTable, - stack: &mut Vec, -) { - let encoded = field_table.lookup(unsafe { ffi::Py_TYPE(parent) }); - if encoded <= 1 { - return; - } - let Some(dict) = get_instance_dict_fast(parent) else { - return; - }; - unsafe { - let keys = &*(*dict.cast::()) - .ma_keys - .cast::(); - let entries = keys.unicode_entries(); - let mut current = (keys.dk_nentries as usize).min((encoded - 1) as usize); - while current > 0 { - current -= 1; - let entry = *entries.add(current); - let value = entry.me_value; - if value.is_null() { - continue; - } - let value_type = ffi::Py_TYPE(value); - if value_type == py_list_type { - let list = value as *mut PyListObject; - let length = (*(list as *mut ffi::PyVarObject)).ob_size; - let ob_item = (*list).ob_item; - for index in (0..length).rev() { - let child = *ob_item.offset(index); - if issubclass_of_ast(ffi::Py_TYPE(child), base_ast_and_expr_type) { - stack.push(Edge { - node: child, - parent, - key: entry.me_key, - index, - }); - } - } - } else if issubclass_of_ast(value_type, base_ast_and_expr_type) { - stack.push(Edge { - node: value, - parent, - key: entry.me_key, - index: -1, - }); - } - } - } -} - -/// Descendants of `root` (excluding `root`) whose exact type is in `kinds`, -/// in depth-first pre-order, without descending below any of them. This is -/// the set of nodes `ast.NodeVisitor.generic_visit` would dispatch to a -/// `visit_` method, given that every other node falls through to the -/// generic descent. -fn walk_node_frontier( - root: *mut PyObject, - kinds: &[*mut PyTypeObject], - base_ast_and_expr_type: (*mut PyTypeObject, *mut PyTypeObject), - py_list_type: *mut PyTypeObject, - field_table: &FieldTable, - result: &mut Vec, -) { - let mut stack = Vec::new(); - unsafe { - push_child_edges( - root, - base_ast_and_expr_type, - py_list_type, - field_table, - &mut stack, - ); - } - while let Some(edge) = stack.pop() { - if kinds.contains(&unsafe { ffi::Py_TYPE(edge.node) }) { - result.push(edge); - } else { - unsafe { - push_child_edges( - edge.node, - base_ast_and_expr_type, - py_list_type, - field_table, - &mut stack, - ); - } - } - } -} - -fn type_pointers(kinds: &Bound<'_, PyTuple>) -> PyResult> { - kinds - .iter() - .map(|kind| Ok(kind.cast_into::()?.as_type_ptr())) - .collect() -} - -fn frontier_edges( - py: Python<'_>, - node: &Bound<'_, PyAny>, - kinds: &Bound<'_, PyTuple>, -) -> PyResult> { - let base = resolve_base_types(py)?; - let py_list_type = PyList::type_object_raw(py); - let kinds = type_pointers(kinds)?; - with_field_table(py, |table| { - let mut result = Vec::new(); - walk_node_frontier( - node.as_ptr(), - &kinds, - base, - py_list_type, - table, - &mut result, - ); - Ok(result) - }) -} - /// Return the descendants of `node` whose exact type is one of `kinds`, in /// depth-first pre-order, without descending below any match. `node` /// itself is never included. @@ -748,6 +623,9 @@ fn fast_walk(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(walk_unordered, m)?)?; m.add_function(wrap_pyfunction!(walk_frontier, m)?)?; m.add_function(wrap_pyfunction!(walk_frontier_edges, m)?)?; + m.add_function(wrap_pyfunction!(fields::walk_of_types, m)?)?; + m.add_function(wrap_pyfunction!(fields::fix_missing_locations, m)?)?; + m.add_function(wrap_pyfunction!(fields::subtree_hashes, m)?)?; m.add_function(wrap_pyfunction!(_walk_count, m)?)?; Ok(()) } diff --git a/tests/test_fields.py b/tests/test_fields.py new file mode 100644 index 0000000..0bb4a2b --- /dev/null +++ b/tests/test_fields.py @@ -0,0 +1,158 @@ +"""`walk_of_types`, `fix_missing_locations` and `subtree_hashes` read each +node's children by `_fields` name, so they must agree with the stdlib on +parsed trees and on nodes a transformer built in another attribute order.""" + +from __future__ import annotations + +import ast +import copy +import gc +import sys +import textwrap + +import pytest + +from fast_walk import fix_missing_locations, subtree_hashes, walk_of_types + +SOURCE = textwrap.dedent(""" + import os + from x import y as z + def f(a: int = 1, *b, c=2.0, **d) -> str: + e = [g(h) for h in i if h] + if e and not a: + return k.l.m(n, *o, p=q) + with r() as s, t: + del u[v:w] + class C(D, metaclass=E): + x: int = True + async def m(self): + async for a in b: + await c + match a: + case [1, *rest] | [2, *rest]: + pass + case {"k": v, **kw}: + pass + lambda q=None: f"{q!r:>{w}}" +""") + + +def pre_order(node: ast.AST): + yield node + for child in ast.iter_child_nodes(node): + yield from pre_order(child) + + +@pytest.mark.parametrize( + "kinds", + [(ast.Name,), (ast.Call, ast.Attribute), (ast.Module,), (ast.Load,), ()], +) +def test_walk_of_types_is_the_filtered_pre_order(kinds): + tree = ast.parse(SOURCE) + assert [id(n) for n in walk_of_types(tree, kinds)] == [ + id(n) for n in pre_order(tree) if type(n) in kinds + ] + + +def test_walk_of_types_reads_fields_by_name(): + call = ast.Call(lineno=1, func=ast.Name(id="f", ctx=ast.Load()), args=[]) + assert [n.id for n in walk_of_types(call, (ast.Name,))] == ["f"] + + +def stripped(source: str) -> ast.Module: + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, (ast.expr, ast.stmt)) and len(node._fields) % 2: + for name in ("lineno", "col_offset", "end_lineno", "end_col_offset"): + if hasattr(node, name): + delattr(node, name) + return tree + + +def locations(tree: ast.AST): + return [ + tuple(getattr(n, name, "-") for name in n._attributes) for n in pre_order(tree) + ] + + +def test_fix_missing_locations_matches_the_stdlib(): + ours, theirs = stripped(SOURCE), stripped(SOURCE) + assert fix_missing_locations(ours) is ours + ast.fix_missing_locations(theirs) + assert locations(ours) == locations(theirs) + compile(ours, "", "exec") + + +def built() -> ast.Module: + return ast.Module( + body=[ + ast.Expr( + value=ast.Call( + lineno=3, + func=ast.Name(id="f", ctx=ast.Load()), + args=[], + keywords=[], + ) + ) + ], + type_ignores=[], + ) + + +def test_fix_missing_locations_fills_a_built_tree(): + ours, theirs = built(), built() + fix_missing_locations(ours) + ast.fix_missing_locations(theirs) + assert locations(ours) == locations(theirs) + assert ours.body[0].value.func.lineno == 3 + + +def test_fix_missing_locations_on_a_deep_chain(): + tree = ast.parse("x" + ".y" * 20_000) + for node in ast.walk(tree): + if isinstance(node, ast.Attribute): + del node.end_lineno + fix_missing_locations(tree) + assert tree.body[0].value.end_lineno == 1 + + +def pairs(left: ast.AST, right: ast.AST): + return [(a, b) for a in pre_order(left) for b in pre_order(right)] + + +def test_subtree_hashes_agree_with_ast_dump(): + old = ast.parse(SOURCE) + new = copy.deepcopy(old) + new.body[2].body[0].value.elt.func.id = "changed" + old_hashes, new_hashes = subtree_hashes(old), subtree_hashes(new) + for a, b in pairs(old, new): + assert (old_hashes[id(a)] == new_hashes[id(b)]) == ( + ast.dump(a) == ast.dump(b) + ), ast.dump(a) + + +@pytest.mark.parametrize( + ("left", "right"), [("x = 1", "x = True"), ("x = 1", "x = 1.0"), ("x = 0", "x = False")] +) +def test_subtree_hashes_tell_equal_hashing_constants_apart(left, right): + a, b = ast.parse(left), ast.parse(right) + assert subtree_hashes(a)[id(a)] != subtree_hashes(b)[id(b)] + + +def test_subtree_hashes_ignore_dict_order(): + built = ast.Call(lineno=1, keywords=[], func=ast.Name(id="f", ctx=ast.Load()), args=[]) + parsed = ast.parse("f()").body[0].value + assert subtree_hashes(built)[id(built)] == subtree_hashes(parsed)[id(parsed)] + + +def test_refcount_neutral(): + tree = ast.parse(SOURCE) + nodes = list(ast.walk(tree)) + gc.collect() + before = [sys.getrefcount(n) for n in nodes] + for _ in range(50): + walk_of_types(tree, (ast.Name, ast.Call)) + subtree_hashes(tree) + fix_missing_locations(tree) + gc.collect() + assert [sys.getrefcount(n) for n in nodes] == before diff --git a/tests/test_frontier.py b/tests/test_frontier.py index 1eede54..49e2bdf 100644 --- a/tests/test_frontier.py +++ b/tests/test_frontier.py @@ -110,3 +110,10 @@ def test_refcount_neutral(): walk_frontier_edges(tree, (ast.Call, ast.Attribute, ast.Name)) gc.collect() assert [sys.getrefcount(n) for n in nodes] == before + + +def test_fields_are_read_by_name(): + call = ast.Call(lineno=1, func=ast.Name(id="f", ctx=ast.Load()), args=[]) + assert [(type(n).__name__, f, i) for n, _, f, i in walk_frontier_edges(call, (ast.Name,))] == [ + ("Name", "func", -1) + ] From d43f5f740ba1a10129de2c4c43d83ad97fe235fc Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 18:40:52 +0000 Subject: [PATCH 3/5] Add set_parents, subtree_hash, walk_frontier_events and a prune option - walk_of_types(node, kinds, prune) stops below nodes of the pruned types. - set_parents sets child.parent across a tree. - subtree_hash returns the root's structural hash without the per-node dict. - walk_frontier_events reports enter/leave around bracketed types, for a visitor that does work around its descent. - A field a constructor left out (ast.Constant(value=1) has no kind entry) is read through attribute lookup, as ast.iter_fields reads it, so it hashes like the parsed node. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hv7iDZqYjqWUscP5LcdwfB --- README.md | 11 ++- fast_walk.pyi | 28 ++++++- src/fields.rs | 182 ++++++++++++++++++++++++++++++++++++------- src/lib.rs | 3 + tests/test_fields.py | 133 ++++++++++++++++++++++++++++++- 5 files changed, 321 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index b2988f9..e98d8a8 100644 --- a/README.md +++ b/README.md @@ -56,14 +56,19 @@ needs to put a replacement back. ### Other helpers -- `walk_of_types(node, kinds)` — every node of the given exact types, in - pre-order, descending into matches. +- `walk_of_types(node, kinds, prune=None)` — every node of the given exact + types, in pre-order, descending into matches but not below a `prune` type. +- `walk_frontier_events(node, kinds, bracket)` — the frontier walk as + `(event, node)` pairs, with enter/leave events around `bracket` types, for a + visitor that does work around its generic descent. +- `set_parents(node)` — `child.parent = node` for the whole tree. - `fix_missing_locations(node)` — `ast.fix_missing_locations` without recursion. - `subtree_hashes(node)` — `{id(n): hash}` such that equal hashes mean equal `ast.dump` output (leaf types included, so `1` and `True` differ). + `subtree_hash(node)` is the root's entry alone. -These, and both frontier walks, read each child by its `_fields` name, so +These, and the frontier walks, read each child by its `_fields` name, so they also work on nodes a transformer built with keyword arguments in another order. `walk_dfs` and `walk_unordered` read fields by position, which is only guaranteed for parsed trees. diff --git a/fast_walk.pyi b/fast_walk.pyi index fc8bbe8..be1dda8 100644 --- a/fast_walk.pyi +++ b/fast_walk.pyi @@ -44,11 +44,35 @@ def walk_frontier_edges( """ def walk_of_types( - node: ast.AST, kinds: tuple[type[ast.AST], ...] + node: ast.AST, + kinds: tuple[type[ast.AST], ...], + prune: tuple[type[ast.AST], ...] | None = None, ) -> list[ast.AST]: """Return every node under `node` (`node` included) whose exact type is in `kinds`, in depth-first pre-order. Unlike :func:`walk_frontier`, - matches are descended into. + matches are descended into, except a node whose exact type is in + `prune`: it is returned if it matches and its children are skipped. + `node` itself is always descended. + """ + +def subtree_hash(node: ast.AST) -> int: + """``subtree_hashes(node)[id(node)]``, without building the dict.""" + +def set_parents(node: ast.AST) -> None: + """Set ``child.parent`` for every node under `node`, visiting parents in + depth-first pre-order. A node reachable from two parents keeps the last. + """ + +def walk_frontier_events( + node: ast.AST, + kinds: tuple[type[ast.AST], ...], + bracket: tuple[type[ast.AST], ...], +) -> list[tuple[int, ast.AST]]: + """The calls a ``NodeVisitor`` makes descending from `node`. + + ``(0, n)`` is a descendant whose exact type is in `kinds`, dispatched and + not descended. ``(1, n)`` and ``(2, n)`` surround the descent of a node + whose exact type is in `bracket`. `node` itself is always descended. """ def fix_missing_locations[T: ast.AST](node: T) -> T: diff --git a/src/fields.rs b/src/fields.rs index f0c3b20..14948e2 100644 --- a/src/fields.rs +++ b/src/fields.rs @@ -97,8 +97,10 @@ fn type_info(py: Python<'_>, kind: *mut PyTypeObject) -> *const TypeInfo { }) } -/// Call `visit(field_index, value)` for each field `node` has set, in -/// `_fields` order. Values are borrowed from the node's dict. +/// Call `visit(field_index, value)` for each field `node` has, in `_fields` +/// order, as `ast.iter_fields` reads them: from the node's dict, else by +/// attribute lookup, which finds the class default of an optional field a +/// constructor left out (`ast.Constant(value=1)` has no `kind` entry). unsafe fn for_each_field( py: Python<'_>, node: *mut PyObject, @@ -110,13 +112,24 @@ unsafe fn for_each_field( }; for (index, name) in info.fields.iter().enumerate() { let value = unsafe { ffi::PyDict_GetItemWithError(dict, name.as_ptr()) }; - if value.is_null() { - if unsafe { !ffi::PyErr_Occurred().is_null() } { - return Err(PyErr::fetch(py)); - } + if !value.is_null() { + visit(index, value)?; continue; } - visit(index, value)?; + if unsafe { !ffi::PyErr_Occurred().is_null() } { + return Err(PyErr::fetch(py)); + } + let inherited = unsafe { ffi::PyObject_GetAttr(node, name.as_ptr()) }; + if inherited.is_null() { + let error = PyErr::fetch(py); + if error.is_instance_of::(py) { + continue; + } + return Err(error); + } + let visited = visit(index, inherited); + unsafe { ffi::Py_DECREF(inherited) }; + visited?; } Ok(()) } @@ -178,35 +191,135 @@ impl<'py> Walk<'py> { } } -fn type_pointers(kinds: &Bound<'_, PyTuple>) -> PyResult> { - kinds - .iter() - .map(|kind| Ok(kind.cast_into::()?.as_type_ptr())) - .collect() +/// A set of exact types, sorted for a binary search. +pub(crate) struct TypeSet(Vec); + +impl TypeSet { + pub(crate) fn new(kinds: &Bound<'_, PyTuple>) -> PyResult { + let mut pointers = kinds + .iter() + .map(|kind| Ok(kind.cast_into::()?.as_type_ptr() as usize)) + .collect::>>()?; + pointers.sort_unstable(); + pointers.dedup(); + Ok(Self(pointers)) + } + + #[inline] + pub(crate) fn contains(&self, node: *mut PyObject) -> bool { + self.0 + .binary_search(&(unsafe { ffi::Py_TYPE(node) } as usize)) + .is_ok() + } } /// Return every node under `node` (`node` included) whose exact type is in /// `kinds`, in depth-first pre-order. Unlike `walk_frontier`, matches are -/// descended into. +/// descended into, except a node whose exact type is in `prune`: it is +/// returned if it matches, and its children are skipped. `node` itself is +/// always descended. #[pyfunction] +#[pyo3(signature = (node, kinds, prune=None))] pub fn walk_of_types<'py>( py: Python<'py>, node: Bound<'py, PyAny>, kinds: Bound<'py, PyTuple>, + prune: Option>, ) -> PyResult> { - let kinds = type_pointers(&kinds)?; + let kinds = TypeSet::new(&kinds)?; + let prune = prune.map(|prune| TypeSet::new(&prune)).transpose()?; let walk = Walk::new(py)?; + let root = node.as_ptr(); let mut result = Vec::new(); - let mut stack = vec![node.as_ptr()]; + let mut stack = vec![root]; while let Some(current) = stack.pop() { - if kinds.contains(&unsafe { ffi::Py_TYPE(current) }) { + if kinds.contains(current) { result.push(current); } + if current != root && prune.as_ref().is_some_and(|prune| prune.contains(current)) { + continue; + } walk.push_children(current, &mut stack)?; } vec_into_pylist(py, &result) } +/// Set `child.parent = parent` for every node under `node`, visiting parents +/// in depth-first pre-order and each parent's children in +/// `ast.iter_child_nodes` order. `node` itself is left alone. +#[pyfunction] +pub fn set_parents(py: Python<'_>, node: Bound<'_, PyAny>) -> PyResult<()> { + let walk = Walk::new(py)?; + let name = intern!(py, "parent"); + let mut children = Vec::new(); + for parent in walk.pre_order(node.as_ptr())? { + children.clear(); + walk.push_children(parent, &mut children)?; + for &child in children.iter().rev() { + if unsafe { ffi::PyObject_SetAttr(child, name.as_ptr(), parent) } == -1 { + return Err(PyErr::fetch(py)); + } + } + } + Ok(()) +} + +const VISIT: i64 = 0; +const ENTER: i64 = 1; +const LEAVE: i64 = 2; + +enum Step { + Child(*mut PyObject), + Leave(*mut PyObject), +} + +/// The calls a `NodeVisitor` makes descending from `node`, as a list of +/// `(event, node)` pairs: `(0, n)` for a descendant whose exact type is in +/// `kinds` (dispatched, not descended), and `(1, n)` / `(2, n)` around the +/// descent of a node whose exact type is in `bracket`. `node` itself is +/// always descended, and bracketed if its type is in `bracket`. +#[pyfunction] +pub fn walk_frontier_events<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, + kinds: Bound<'py, PyTuple>, + bracket: Bound<'py, PyTuple>, +) -> PyResult> { + let kinds = TypeSet::new(&kinds)?; + let bracket = TypeSet::new(&bracket)?; + let walk = Walk::new(py)?; + let mut events: Vec<(i64, *mut PyObject)> = Vec::new(); + let mut stack: Vec = Vec::new(); + let mut children = Vec::new(); + let mut open = |current: *mut PyObject, + stack: &mut Vec, + events: &mut Vec<(i64, *mut PyObject)>| + -> PyResult<()> { + if bracket.contains(current) { + events.push((ENTER, current)); + stack.push(Step::Leave(current)); + } + children.clear(); + walk.push_children(current, &mut children)?; + stack.extend(children.iter().map(|&child| Step::Child(child))); + Ok(()) + }; + open(node.as_ptr(), &mut stack, &mut events)?; + while let Some(step) = stack.pop() { + match step { + Step::Leave(current) => events.push((LEAVE, current)), + Step::Child(current) if kinds.contains(current) => events.push((VISIT, current)), + Step::Child(current) => open(current, &mut stack, &mut events)?, + } + } + let list = PyList::empty(py); + for (event, current) in events { + let pair = (event, unsafe { Bound::from_borrowed_ptr(py, current) }); + list.append(pair)?; + } + Ok(list.into_any()) +} + /// `ast.fix_missing_locations`, without recursing per tree level. Returns /// `node`. #[pyfunction] @@ -277,6 +390,27 @@ pub fn subtree_hashes<'py>( py: Python<'py>, node: Bound<'py, PyAny>, ) -> PyResult> { + let hashes = hash_subtrees(py, &node)?; + let result = pyo3::types::PyDict::new(py); + for (&pointer, &hash) in &hashes { + let key = unsafe { + Bound::from_owned_ptr_or_err( + py, + ffi::PyLong_FromVoidPtr(pointer as *mut std::ffi::c_void), + )? + }; + result.set_item(key, hash as i64)?; + } + Ok(result.into_any()) +} + +/// `subtree_hashes(node)[id(node)]`, without building the dict. +#[pyfunction] +pub fn subtree_hash(py: Python<'_>, node: Bound<'_, PyAny>) -> PyResult { + Ok(hash_subtrees(py, &node)?[&(node.as_ptr() as usize)] as i64) +} + +fn hash_subtrees(py: Python<'_>, node: &Bound<'_, PyAny>) -> PyResult> { let walk = Walk::new(py)?; let nodes = walk.pre_order(node.as_ptr())?; let mut hashes: PointerMap = PointerMap::default(); @@ -317,17 +451,7 @@ pub fn subtree_hashes<'py>( } hashes.insert(current as usize, hash); } - let result = pyo3::types::PyDict::new(py); - for (&pointer, &hash) in &hashes { - let key = unsafe { - Bound::from_owned_ptr_or_err( - py, - ffi::PyLong_FromVoidPtr(pointer as *mut std::ffi::c_void), - )? - }; - result.set_item(key, hash as i64)?; - } - Ok(result.into_any()) + Ok(hashes) } impl Walk<'_> { @@ -378,13 +502,13 @@ pub(crate) fn frontier_edges( node: &Bound<'_, PyAny>, kinds: &Bound<'_, PyTuple>, ) -> PyResult> { - let kinds = type_pointers(kinds)?; + let kinds = TypeSet::new(kinds)?; let walk = Walk::new(py)?; let mut result = Vec::new(); let mut stack = Vec::new(); walk.push_child_edges(node.as_ptr(), &mut stack)?; while let Some(edge) = stack.pop() { - if kinds.contains(&unsafe { ffi::Py_TYPE(edge.node) }) { + if kinds.contains(edge.node) { result.push(edge); } else { walk.push_child_edges(edge.node, &mut stack)?; diff --git a/src/lib.rs b/src/lib.rs index 619eb6a..e07d18f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -626,6 +626,9 @@ fn fast_walk(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(fields::walk_of_types, m)?)?; m.add_function(wrap_pyfunction!(fields::fix_missing_locations, m)?)?; m.add_function(wrap_pyfunction!(fields::subtree_hashes, m)?)?; + m.add_function(wrap_pyfunction!(fields::subtree_hash, m)?)?; + m.add_function(wrap_pyfunction!(fields::set_parents, m)?)?; + m.add_function(wrap_pyfunction!(fields::walk_frontier_events, m)?)?; m.add_function(wrap_pyfunction!(_walk_count, m)?)?; Ok(()) } diff --git a/tests/test_fields.py b/tests/test_fields.py index 0bb4a2b..023f329 100644 --- a/tests/test_fields.py +++ b/tests/test_fields.py @@ -1,4 +1,4 @@ -"""`walk_of_types`, `fix_missing_locations` and `subtree_hashes` read each +"""The functions here read each node's children by `_fields` name, so they must agree with the stdlib on parsed trees and on nodes a transformer built in another attribute order.""" @@ -12,7 +12,14 @@ import pytest -from fast_walk import fix_missing_locations, subtree_hashes, walk_of_types +from fast_walk import ( + fix_missing_locations, + set_parents, + subtree_hash, + subtree_hashes, + walk_frontier_events, + walk_of_types, +) SOURCE = textwrap.dedent(""" import os @@ -59,6 +66,98 @@ def test_walk_of_types_reads_fields_by_name(): assert [n.id for n in walk_of_types(call, (ast.Name,))] == ["f"] +SCOPES = (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda, ast.ClassDef) + + +def pruned_pre_order(node: ast.AST, prune: tuple[type, ...], root: bool = True): + yield node + if root or type(node) not in prune: + for child in ast.iter_child_nodes(node): + yield from pruned_pre_order(child, prune, root=False) + + +@pytest.mark.parametrize("kinds", [(ast.Name,), (ast.Return, *SCOPES), ()]) +def test_walk_of_types_stops_below_a_pruned_node(kinds): + tree = ast.parse(SOURCE) + for root in [tree, *walk_of_types(tree, SCOPES)]: + assert [id(n) for n in walk_of_types(root, kinds, SCOPES)] == [ + id(n) for n in pruned_pre_order(root, SCOPES) if type(n) in kinds + ] + + +def test_set_parents_matches_a_pre_order_loop(): + ours, theirs = ast.parse(SOURCE), ast.parse(SOURCE) + set_parents(ours) + for node in pre_order(theirs): + for child in ast.iter_child_nodes(node): + child.parent = node + assert not hasattr(ours, "parent") + position = {id(n): i for i, n in enumerate(pre_order(theirs))} + position |= {id(n): i for i, n in enumerate(pre_order(ours))} + assert [ + position[id(n.parent)] for n in pre_order(ours) if hasattr(n, "parent") + ] == [position[id(n.parent)] for n in pre_order(theirs) if hasattr(n, "parent")] + + +def test_set_parents_gives_a_shared_node_its_last_parent(): + shared = ast.Name(id="x", ctx=ast.Load()) + tree = ast.parse("f(a)\ng(b)") + tree.body[0].value.args[0] = shared + tree.body[1].value.args[0] = shared + set_parents(tree) + assert shared.parent is tree.body[1].value + + +def reference_events(node, kinds, bracket): + events = [] + + def descend(current): + if type(current) in bracket: + events.append((1, current)) + for child in ast.iter_child_nodes(current): + if type(child) in kinds: + events.append((0, child)) + else: + descend(child) + if type(current) in bracket: + events.append((2, current)) + + descend(node) + return events + + +EXPRESSIONS = tuple( + kind for kind in vars(ast).values() + if isinstance(kind, type) and issubclass(kind, ast.expr) and kind is not ast.expr +) + + +@pytest.mark.parametrize( + ("kinds", "bracket"), + [ + ((ast.Name, ast.Assign), EXPRESSIONS), + ((ast.Call,), (ast.Module, ast.Attribute)), + ((), EXPRESSIONS), + ((ast.FunctionDef,), ()), + ], +) +def test_walk_frontier_events_matches_a_recursive_descent(kinds, bracket): + tree = ast.parse(SOURCE) + for root in [tree, *walk_of_types(tree, (ast.Call, ast.Assign))]: + expected = reference_events(root, kinds, bracket) + assert [(e, id(n)) for e, n in walk_frontier_events(root, kinds, bracket)] == [ + (e, id(n)) for e, n in expected + ] + + +def test_walk_frontier_events_on_a_deep_chain(): + chain: ast.expr = ast.Name(id="a", ctx=ast.Load()) + for _ in range(50_000): + chain = ast.Attribute(value=chain, attr="b", ctx=ast.Load()) + events = walk_frontier_events(chain, (ast.Name,), EXPRESSIONS) + assert len(events) == 2 * 50_000 + 1 + + def stripped(source: str) -> ast.Module: tree = ast.parse(source) for node in ast.walk(tree): @@ -139,6 +238,31 @@ def test_subtree_hashes_tell_equal_hashing_constants_apart(left, right): assert subtree_hashes(a)[id(a)] != subtree_hashes(b)[id(b)] +@pytest.mark.parametrize( + ("built", "source"), + [ + (lambda: ast.Constant(value="x"), '"x"'), + (lambda: ast.Constant(value="x", kind=None), '"x"'), + (lambda: ast.Call(func=ast.Name(id="f"), args=[], keywords=[]), "f()"), + ], +) +def test_a_field_a_constructor_left_out_hashes_as_its_default(built, source): + node, parsed = built(), ast.parse(source).body[0].value + assert ast.dump(node) == ast.dump(parsed) + assert subtree_hash(node) == subtree_hash(parsed) + + +def test_a_field_a_constructor_left_out_is_not_a_child(): + node = ast.Constant(value="x") + assert walk_of_types(node, (ast.Constant,)) == [node] + assert walk_frontier_events(node, (), (ast.Constant,)) == [(1, node), (2, node)] + + +def test_subtree_hash_is_the_roots_entry(): + tree = ast.parse(SOURCE) + assert subtree_hash(tree) == subtree_hashes(tree)[id(tree)] + + def test_subtree_hashes_ignore_dict_order(): built = ast.Call(lineno=1, keywords=[], func=ast.Name(id="f", ctx=ast.Load()), args=[]) parsed = ast.parse("f()").body[0].value @@ -147,12 +271,17 @@ def test_subtree_hashes_ignore_dict_order(): def test_refcount_neutral(): tree = ast.parse(SOURCE) + set_parents(tree) nodes = list(ast.walk(tree)) gc.collect() before = [sys.getrefcount(n) for n in nodes] for _ in range(50): walk_of_types(tree, (ast.Name, ast.Call)) subtree_hashes(tree) + subtree_hash(tree) + set_parents(tree) fix_missing_locations(tree) + walk_of_types(tree, (ast.Name,), SCOPES) + walk_frontier_events(tree, (ast.Name,), EXPRESSIONS) gc.collect() assert [sys.getrefcount(n) for n in nodes] == before From f7ad164ac8658dcbc760ef25a5c9d723b183d060 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 20:12:45 +0000 Subject: [PATCH 4/5] Hold the nodes a field walk will still touch A field read can run Python code -- a property, a __getattr__, a leaf's __hash__ or __repr__ -- that drops the tree's own references, and the walks held borrowed pointers across it: a property computing a node returned a pointer to a freed object. Every pushed node, edge and event is now an owned reference, and a list is re-read item by item. - TypeInfo holds its type, so a freed type's address is not read with its _fields by a new type allocated there; the cache is filled without a borrow held across the Python calls that read it. - is_ast recognises a subclass at any depth, not only two levels down. - A node with no instance dict reads its fields by attribute. - Leaf values other than str, bytes, int, bool, None and Ellipsis hash by repr, as ast.dump compares them: NaNs together, -0.0 apart from 0.0, and an unhashable constant no longer raises. A child a property computes afresh is hashed on the spot rather than panicking. - set_parents reads each node's fields once. - walk_frontier and walk_frontier_events fill their lists directly. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hv7iDZqYjqWUscP5LcdwfB --- fast_walk.pyi | 5 +- src/fields.rs | 420 +++++++++++++++++++++++++++++-------------- src/lib.rs | 20 +-- tests/test_fields.py | 91 ++++++++++ 4 files changed, 393 insertions(+), 143 deletions(-) diff --git a/fast_walk.pyi b/fast_walk.pyi index be1dda8..ebeebe7 100644 --- a/fast_walk.pyi +++ b/fast_walk.pyi @@ -82,7 +82,10 @@ def subtree_hashes(node: ast.AST) -> dict[int, int]: """A structural hash per node under `node`, keyed by ``id(node)``. Subtrees with equal ``ast.dump`` output get equal hashes. A leaf value's - type is mixed in, so ``1``, ``1.0`` and ``True`` hash apart. + type is mixed in, so ``1``, ``1.0`` and ``True`` hash apart, and a value + that is not a str, bytes, int, bool, None or Ellipsis is hashed by its + ``repr``, so NaNs hash together, ``-0.0`` and ``0.0`` apart, and an + unhashable constant still hashes. """ def walk(node: ast.AST) -> list[ast.AST]: diff --git a/src/fields.rs b/src/fields.rs index 14948e2..aa4d449 100644 --- a/src/fields.rs +++ b/src/fields.rs @@ -5,6 +5,10 @@ //! has another order (`ast.Call(lineno=1, func=f)` stores `lineno` first), //! and these functions run on such trees, so each field is looked up by //! name instead. +//! +//! A field read can run Python code (a property, a `__getattr__`, a value's +//! `__hash__` or `__repr__`) that drops the tree's own references, so every +//! node a walk will still touch is held as a `Strong`. use std::cell::RefCell; use std::collections::HashMap; @@ -15,7 +19,69 @@ use pyo3::prelude::*; use pyo3::types::{PyList, PyTuple, PyType}; use pyo3::{PyTypeInfo, intern}; -use crate::{Edge, get_instance_dict_fast, issubclass_of_ast, resolve_base_types, vec_into_pylist}; +use crate::{Edge, get_instance_dict_fast, issubclass_of_ast, resolve_base_types}; + +/// An owned reference, dropped with the GIL held. +pub(crate) struct Strong(*mut PyObject); + +impl Strong { + /// # Safety + /// `object` is a live object and the GIL is held. + #[inline(always)] + pub(crate) unsafe fn borrowed(object: *mut PyObject) -> Self { + unsafe { ffi::Py_INCREF(object) }; + Self(object) + } + + /// # Safety + /// `object` is a new reference the caller gives up. + #[inline(always)] + unsafe fn owned(object: *mut PyObject) -> Self { + Self(object) + } + + #[inline(always)] + pub(crate) fn as_ptr(&self) -> *mut PyObject { + self.0 + } + + #[inline(always)] + pub(crate) fn into_ptr(self) -> *mut PyObject { + let object = self.0; + std::mem::forget(self); + object + } +} + +impl Clone for Strong { + #[inline(always)] + fn clone(&self) -> Self { + unsafe { Self::borrowed(self.0) } + } +} + +impl Drop for Strong { + #[inline(always)] + fn drop(&mut self) { + unsafe { ffi::Py_DECREF(self.0) }; + } +} + +/// A list holding `items`, whose references it takes over. +pub(crate) fn into_pylist<'py>( + py: Python<'py>, + items: impl ExactSizeIterator, +) -> PyResult> { + unsafe { + let list = + Bound::from_owned_ptr_or_err(py, ffi::PyList_New(items.len() as ffi::Py_ssize_t))?; + let slots = (*(list.as_ptr() as *mut ffi::PyListObject)).ob_item; + for (position, item) in items.enumerate() { + *slots.add(position) = item.into_ptr(); + } + Ok(list) + } +} #[derive(Default)] struct PointerHasher(u64); @@ -44,7 +110,10 @@ const END_LINENO: u8 = 4; const END_COL_OFFSET: u8 = 8; /// An AST type's `_fields` names, and which location `_attributes` it has. +/// It holds the type, so the address it is cached under cannot be reused by +/// another type. struct TypeInfo { + _kind: Py, fields: Vec>, locations: u8, } @@ -71,53 +140,62 @@ fn names_of(py: Python<'_>, kind: &Bound<'_, PyAny>, attribute: &str) -> Vec, kind: *mut PyTypeObject) -> *const TypeInfo { - TYPE_INFO.with(|cache| { - let mut cache = cache.borrow_mut(); - let entry = cache.entry(kind as usize).or_insert_with(|| { - let kind = unsafe { Bound::from_borrowed_ptr(py, kind.cast::()) }; - let attributes = names_of(py, &kind, "_attributes"); - let mut locations = 0; - for name in &attributes { - locations |= match name.bind(py).extract::<&str>().unwrap_or_default() { - "lineno" => LINENO, - "col_offset" => COL_OFFSET, - "end_lineno" => END_LINENO, - "end_col_offset" => END_COL_OFFSET, - _ => 0, - }; - } - Box::new(TypeInfo { - fields: names_of(py, &kind, "_fields"), - locations, - }) - }); - &**entry as *const TypeInfo - }) + let cached = TYPE_INFO.with(|cache| { + cache + .borrow() + .get(&(kind as usize)) + .map(|entry| &**entry as *const TypeInfo) + }); + if let Some(info) = cached { + return info; + } + let kind_object = unsafe { Bound::from_borrowed_ptr(py, kind.cast::()) }; + let mut locations = 0; + for name in names_of(py, &kind_object, "_attributes") { + locations |= match name.bind(py).extract::<&str>().unwrap_or_default() { + "lineno" => LINENO, + "col_offset" => COL_OFFSET, + "end_lineno" => END_LINENO, + "end_col_offset" => END_COL_OFFSET, + _ => 0, + }; + } + let info = Box::new(TypeInfo { + fields: names_of(py, &kind_object, "_fields"), + locations, + _kind: kind_object.unbind(), + }); + TYPE_INFO + .with(|cache| &**cache.borrow_mut().entry(kind as usize).or_insert(info) as *const TypeInfo) } /// Call `visit(field_index, value)` for each field `node` has, in `_fields` /// order, as `ast.iter_fields` reads them: from the node's dict, else by /// attribute lookup, which finds the class default of an optional field a /// constructor left out (`ast.Constant(value=1)` has no `kind` entry). +/// `node` must be kept alive by the caller. The value is borrowed: a `visit` +/// that runs Python code holds it first. unsafe fn for_each_field( py: Python<'_>, node: *mut PyObject, info: &TypeInfo, mut visit: impl FnMut(usize, *mut PyObject) -> PyResult<()>, ) -> PyResult<()> { - let Some(dict) = get_instance_dict_fast(node) else { - return Ok(()); - }; + let dict = get_instance_dict_fast(node).map(|dict| unsafe { Strong::borrowed(dict) }); for (index, name) in info.fields.iter().enumerate() { - let value = unsafe { ffi::PyDict_GetItemWithError(dict, name.as_ptr()) }; - if !value.is_null() { - visit(index, value)?; - continue; - } - if unsafe { !ffi::PyErr_Occurred().is_null() } { - return Err(PyErr::fetch(py)); + if let Some(dict) = &dict { + let value = unsafe { ffi::PyDict_GetItemWithError(dict.as_ptr(), name.as_ptr()) }; + if !value.is_null() { + visit(index, value)?; + continue; + } + if unsafe { !ffi::PyErr_Occurred().is_null() } { + return Err(PyErr::fetch(py)); + } } let inherited = unsafe { ffi::PyObject_GetAttr(node, name.as_ptr()) }; if inherited.is_null() { @@ -127,9 +205,23 @@ unsafe fn for_each_field( } return Err(error); } - let visited = visit(index, inherited); - unsafe { ffi::Py_DECREF(inherited) }; - visited?; + let inherited = unsafe { Strong::owned(inherited) }; + visit(index, inherited.as_ptr())?; + } + Ok(()) +} + +/// The items of `list`, borrowed, read one at a time so Python code run +/// between them cannot leave the loop reading a resized list. +#[inline(always)] +unsafe fn for_each_item( + list: *mut PyObject, + mut visit: impl FnMut(ffi::Py_ssize_t, *mut PyObject) -> PyResult<()>, +) -> PyResult<()> { + let mut index = 0; + while index < unsafe { ffi::PyList_GET_SIZE(list) } { + visit(index, unsafe { ffi::PyList_GET_ITEM(list, index) })?; + index += 1; } Ok(()) } @@ -149,28 +241,36 @@ impl<'py> Walk<'py> { }) } + /// Whether `value` is an `ast.AST` instance, at any depth of subclass. + #[inline(always)] fn is_ast(&self, value: *mut PyObject) -> bool { - issubclass_of_ast(unsafe { ffi::Py_TYPE(value) }, self.base) + let kind = unsafe { ffi::Py_TYPE(value) }; + if issubclass_of_ast(kind, self.base) { + return true; + } + let first = unsafe { (*kind).tp_base }; + if first.is_null() || first == &raw mut ffi::PyBaseObject_Type { + return false; + } + unsafe { ffi::PyType_IsSubtype(kind, self.base.0) != 0 } } /// Push the AST children of `node` onto `stack` so they pop in /// `ast.iter_child_nodes` order. - fn push_children(&self, node: *mut PyObject, stack: &mut Vec<*mut PyObject>) -> PyResult<()> { - let info = unsafe { &*type_info(self.py, ffi::Py_TYPE(node)) }; + fn push_children(&self, node: &Strong, stack: &mut Vec) -> PyResult<()> { + let info = unsafe { &*type_info(self.py, ffi::Py_TYPE(node.as_ptr())) }; let start = stack.len(); unsafe { - for_each_field(self.py, node, info, |_, value| { + for_each_field(self.py, node.as_ptr(), info, |_, value| { if ffi::Py_TYPE(value) == self.list_type { - let list = value as *mut ffi::PyListObject; - let length = (*(list as *mut ffi::PyVarObject)).ob_size; - for index in 0..length { - let item = *(*list).ob_item.offset(index); + for_each_item(value, |_, item| { if self.is_ast(item) { - stack.push(item); + stack.push(Strong::borrowed(item)); } - } + Ok(()) + })?; } else if self.is_ast(value) { - stack.push(value); + stack.push(Strong::borrowed(value)); } Ok(()) })?; @@ -180,12 +280,12 @@ impl<'py> Walk<'py> { } /// Every node under `root`, `root` included, in depth-first pre-order. - fn pre_order(&self, root: *mut PyObject) -> PyResult> { + fn pre_order(&self, root: &Bound<'_, PyAny>) -> PyResult> { let mut result = Vec::new(); - let mut stack = vec![root]; + let mut stack = vec![unsafe { Strong::borrowed(root.as_ptr()) }]; while let Some(node) = stack.pop() { + self.push_children(&node, &mut stack)?; result.push(node); - self.push_children(node, &mut stack)?; } Ok(result) } @@ -206,9 +306,9 @@ impl TypeSet { } #[inline] - pub(crate) fn contains(&self, node: *mut PyObject) -> bool { + pub(crate) fn contains(&self, node: &Strong) -> bool { self.0 - .binary_search(&(unsafe { ffi::Py_TYPE(node) } as usize)) + .binary_search(&(unsafe { ffi::Py_TYPE(node.as_ptr()) } as usize)) .is_ok() } } @@ -231,17 +331,18 @@ pub fn walk_of_types<'py>( let walk = Walk::new(py)?; let root = node.as_ptr(); let mut result = Vec::new(); - let mut stack = vec![root]; + let mut stack = vec![unsafe { Strong::borrowed(root) }]; while let Some(current) = stack.pop() { - if kinds.contains(current) { - result.push(current); + let pruned = current.as_ptr() != root + && prune.as_ref().is_some_and(|prune| prune.contains(¤t)); + if !pruned { + walk.push_children(¤t, &mut stack)?; } - if current != root && prune.as_ref().is_some_and(|prune| prune.contains(current)) { - continue; + if kinds.contains(¤t) { + result.push(current); } - walk.push_children(current, &mut stack)?; } - vec_into_pylist(py, &result) + into_pylist(py, result.into_iter()) } /// Set `child.parent = parent` for every node under `node`, visiting parents @@ -252,14 +353,17 @@ pub fn set_parents(py: Python<'_>, node: Bound<'_, PyAny>) -> PyResult<()> { let walk = Walk::new(py)?; let name = intern!(py, "parent"); let mut children = Vec::new(); - for parent in walk.pre_order(node.as_ptr())? { - children.clear(); - walk.push_children(parent, &mut children)?; - for &child in children.iter().rev() { - if unsafe { ffi::PyObject_SetAttr(child, name.as_ptr(), parent) } == -1 { + let mut stack = vec![unsafe { Strong::borrowed(node.as_ptr()) }]; + while let Some(parent) = stack.pop() { + walk.push_children(&parent, &mut children)?; + for child in children.iter().rev() { + if unsafe { ffi::PyObject_SetAttr(child.as_ptr(), name.as_ptr(), parent.as_ptr()) } + == -1 + { return Err(PyErr::fetch(py)); } } + stack.append(&mut children); } Ok(()) } @@ -269,8 +373,8 @@ const ENTER: i64 = 1; const LEAVE: i64 = 2; enum Step { - Child(*mut PyObject), - Leave(*mut PyObject), + Child(Strong), + Leave(Strong), } /// The calls a `NodeVisitor` makes descending from `node`, as a list of @@ -288,36 +392,50 @@ pub fn walk_frontier_events<'py>( let kinds = TypeSet::new(&kinds)?; let bracket = TypeSet::new(&bracket)?; let walk = Walk::new(py)?; - let mut events: Vec<(i64, *mut PyObject)> = Vec::new(); + let mut events: Vec<(i64, Strong)> = Vec::new(); let mut stack: Vec = Vec::new(); let mut children = Vec::new(); - let mut open = |current: *mut PyObject, - stack: &mut Vec, - events: &mut Vec<(i64, *mut PyObject)>| - -> PyResult<()> { - if bracket.contains(current) { - events.push((ENTER, current)); - stack.push(Step::Leave(current)); - } - children.clear(); - walk.push_children(current, &mut children)?; - stack.extend(children.iter().map(|&child| Step::Child(child))); - Ok(()) - }; - open(node.as_ptr(), &mut stack, &mut events)?; + let mut open = + |current: Strong, stack: &mut Vec, events: &mut Vec<(i64, Strong)>| -> PyResult<()> { + children.clear(); + walk.push_children(¤t, &mut children)?; + if bracket.contains(¤t) { + events.push((ENTER, current.clone())); + stack.push(Step::Leave(current)); + } + stack.extend(children.drain(..).map(Step::Child)); + Ok(()) + }; + open( + unsafe { Strong::borrowed(node.as_ptr()) }, + &mut stack, + &mut events, + )?; while let Some(step) = stack.pop() { match step { Step::Leave(current) => events.push((LEAVE, current)), - Step::Child(current) if kinds.contains(current) => events.push((VISIT, current)), + Step::Child(current) if kinds.contains(¤t) => events.push((VISIT, current)), Step::Child(current) => open(current, &mut stack, &mut events)?, } } - let list = PyList::empty(py); - for (event, current) in events { - let pair = (event, unsafe { Bound::from_borrowed_ptr(py, current) }); - list.append(pair)?; + let codes = + [VISIT, ENTER, LEAVE].map(|code| code.into_pyobject(py).unwrap().into_any().unbind()); + unsafe { + let list = + Bound::from_owned_ptr_or_err(py, ffi::PyList_New(events.len() as ffi::Py_ssize_t))?; + let slots = (*(list.as_ptr() as *mut ffi::PyListObject)).ob_item; + for (position, (event, current)) in events.into_iter().enumerate() { + let pair = Bound::from_owned_ptr_or_err(py, ffi::PyTuple_New(2))?; + ffi::PyTuple_SET_ITEM( + pair.as_ptr(), + 0, + codes[event as usize].clone_ref(py).into_ptr(), + ); + ffi::PyTuple_SET_ITEM(pair.as_ptr(), 1, current.into_ptr()); + *slots.add(position) = pair.into_ptr(); + } + Ok(list) } - Ok(list.into_any()) } /// `ast.fix_missing_locations`, without recursing per tree level. Returns @@ -338,21 +456,24 @@ pub fn fix_missing_locations<'py>( let zero = 0i64.into_pyobject(py)?.into_any().unbind(); let mut locations: Vec<[Py; 4]> = vec![[one.clone_ref(py), zero.clone_ref(py), one, zero]]; - let mut stack: Vec<(*mut PyObject, usize)> = vec![(node.as_ptr(), 0)]; + let mut stack: Vec<(Strong, usize)> = vec![(unsafe { Strong::borrowed(node.as_ptr()) }, 0)]; let mut children = Vec::new(); while let Some((current, inherited)) = stack.pop() { - let info = unsafe { &*type_info(py, ffi::Py_TYPE(current)) }; + let info = unsafe { &*type_info(py, ffi::Py_TYPE(current.as_ptr())) }; let mut own = inherited; if info.locations != 0 { - let dict = get_instance_dict_fast(current); + let dict = get_instance_dict_fast(current.as_ptr()) + .map(|dict| unsafe { Strong::borrowed(dict) }); let mut values: [Py; 4] = locations[inherited].each_ref().map(|v| v.clone_ref(py)); for (slot, (flag, name, none_is_missing)) in names.iter().enumerate() { if info.locations & flag == 0 { continue; } - let found = match dict { - Some(dict) => unsafe { ffi::PyDict_GetItemWithError(dict, name.as_ptr()) }, + let found = match &dict { + Some(dict) => unsafe { + ffi::PyDict_GetItemWithError(dict.as_ptr(), name.as_ptr()) + }, None => std::ptr::null_mut(), }; if found.is_null() && unsafe { !ffi::PyErr_Occurred().is_null() } { @@ -361,7 +482,7 @@ pub fn fix_missing_locations<'py>( let missing = found.is_null() || (*none_is_missing && unsafe { found == ffi::Py_None() }); if missing { - unsafe { Bound::from_borrowed_ptr(py, current) } + unsafe { Bound::from_borrowed_ptr(py, current.as_ptr()) } .setattr(name, values[slot].bind(py))?; } else { values[slot] = unsafe { Bound::from_borrowed_ptr(py, found) }.unbind(); @@ -371,8 +492,8 @@ pub fn fix_missing_locations<'py>( own = locations.len() - 1; } children.clear(); - walk.push_children(current, &mut children)?; - stack.extend(children.iter().map(|&child| (child, own))); + walk.push_children(¤t, &mut children)?; + stack.extend(children.drain(..).map(|child| (child, own))); } Ok(node) } @@ -384,7 +505,10 @@ fn mix(hash: u64, value: u64) -> u64 { /// A structural hash per node, keyed by `id(node)`. Two subtrees that would /// give the same `ast.dump` get the same hash; a leaf value's type is mixed -/// in, so `1`, `1.0` and `True` (equal hashes in Python) stay apart. +/// in, so `1`, `1.0` and `True` (equal hashes in Python) stay apart. A leaf +/// that is not a str, bytes, int, bool, None or Ellipsis is hashed by its +/// `repr`, as `ast.dump` prints it: NaNs compare equal there and `-0.0` and +/// `0.0` do not, and a list constant has a repr where it has no hash. #[pyfunction] pub fn subtree_hashes<'py>( py: Python<'py>, @@ -407,49 +531,85 @@ pub fn subtree_hashes<'py>( /// `subtree_hashes(node)[id(node)]`, without building the dict. #[pyfunction] pub fn subtree_hash(py: Python<'_>, node: Bound<'_, PyAny>) -> PyResult { - Ok(hash_subtrees(py, &node)?[&(node.as_ptr() as usize)] as i64) + let hashes = hash_subtrees(py, &node)?; + Ok(hashed(&hashes, node.as_ptr())? as i64) +} + +fn hashed(hashes: &PointerMap, node: *mut PyObject) -> PyResult { + hashes.get(&(node as usize)).copied().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err("the tree changed while it was being hashed") + }) +} + +/// Whether `value`'s hash is what `ast.dump` compares: its type's hash is +/// C code that agrees with its repr. +#[inline(always)] +fn hashes_as_dumped(value: *mut PyObject) -> bool { + unsafe { + let kind = ffi::Py_TYPE(value); + kind == &raw mut ffi::PyUnicode_Type + || kind == &raw mut ffi::PyLong_Type + || kind == &raw mut ffi::PyBool_Type + || kind == &raw mut ffi::PyBytes_Type + || value == ffi::Py_None() + || value == ffi::Py_Ellipsis() + } } fn hash_subtrees(py: Python<'_>, node: &Bound<'_, PyAny>) -> PyResult> { let walk = Walk::new(py)?; - let nodes = walk.pre_order(node.as_ptr())?; + let nodes = walk.pre_order(node)?; let mut hashes: PointerMap = PointerMap::default(); hashes.reserve(nodes.len()); let leaf = |value: *mut PyObject, hashes: &PointerMap| -> PyResult { if walk.is_ast(value) { - return Ok(hashes[&(value as usize)]); + if let Some(&hash) = hashes.get(&(value as usize)) { + return Ok(hash); + } + // A field that computes its node gives a new one on every read. + let child = unsafe { Bound::from_borrowed_ptr(py, value) }; + return hashed(&hash_subtrees(py, &child)?, value); } - let hashed = unsafe { ffi::PyObject_Hash(value) }; - if hashed == -1 && unsafe { !ffi::PyErr_Occurred().is_null() } { + let kind = unsafe { ffi::Py_TYPE(value) } as u64; + let hash = unsafe { + if hashes_as_dumped(value) { + ffi::PyObject_Hash(value) + } else { + let repr = Bound::from_owned_ptr_or_err(py, ffi::PyObject_Repr(value))?; + ffi::PyObject_Hash(repr.as_ptr()) + } + }; + if hash == -1 && unsafe { !ffi::PyErr_Occurred().is_null() } { return Err(PyErr::fetch(py)); } - Ok(mix(unsafe { ffi::Py_TYPE(value) } as u64, hashed as u64)) + Ok(mix(kind, hash as u64)) }; - for ¤t in nodes.iter().rev() { - if hashes.contains_key(&(current as usize)) { + for current in nodes.iter().rev() { + if hashes.contains_key(&(current.as_ptr() as usize)) { continue; } - let kind = unsafe { ffi::Py_TYPE(current) }; + let kind = unsafe { ffi::Py_TYPE(current.as_ptr()) }; let info = unsafe { &*type_info(py, kind) }; let mut hash = mix(0, kind as u64); unsafe { - for_each_field(py, current, info, |index, value| { - let value_hash = if ffi::Py_TYPE(value) == walk.list_type { - let list = value as *mut ffi::PyListObject; - let length = (*(list as *mut ffi::PyVarObject)).ob_size; - let mut items = mix(u64::MAX, length as u64); - for item in 0..length { - items = mix(items, leaf(*(*list).ob_item.offset(item), &hashes)?); - } + for_each_field(py, current.as_ptr(), info, |index, value| { + let value = Strong::borrowed(value); + let value_hash = if ffi::Py_TYPE(value.as_ptr()) == walk.list_type { + let mut items = mix(u64::MAX, ffi::PyList_GET_SIZE(value.as_ptr()) as u64); + for_each_item(value.as_ptr(), |_, item| { + let item = Strong::borrowed(item); + items = mix(items, leaf(item.as_ptr(), &hashes)?); + Ok(()) + })?; items } else { - leaf(value, &hashes)? + leaf(value.as_ptr(), &hashes)? }; hash = mix(mix(hash, index as u64), value_hash); Ok(()) })?; } - hashes.insert(current as usize, hash); + hashes.insert(current.as_ptr() as usize, hash); } Ok(hashes) } @@ -457,30 +617,28 @@ fn hash_subtrees(py: Python<'_>, node: &Bound<'_, PyAny>) -> PyResult { /// Push the AST children of `parent` onto `stack` so they pop in /// `ast.iter_child_nodes` order, each with the slot it was read from. - fn push_child_edges(&self, parent: *mut PyObject, stack: &mut Vec) -> PyResult<()> { - let info = unsafe { &*type_info(self.py, ffi::Py_TYPE(parent)) }; + fn push_child_edges(&self, parent: &Strong, stack: &mut Vec) -> PyResult<()> { + let info = unsafe { &*type_info(self.py, ffi::Py_TYPE(parent.as_ptr())) }; let start = stack.len(); unsafe { - for_each_field(self.py, parent, info, |index, value| { + for_each_field(self.py, parent.as_ptr(), info, |index, value| { let key = info.fields[index].as_ptr(); if ffi::Py_TYPE(value) == self.list_type { - let list = value as *mut ffi::PyListObject; - let length = (*(list as *mut ffi::PyVarObject)).ob_size; - for item in 0..length { - let child = *(*list).ob_item.offset(item); + for_each_item(value, |item, child| { if self.is_ast(child) { stack.push(Edge { - node: child, - parent, + node: Strong::borrowed(child), + parent: parent.clone(), key, index: item, }); } - } + Ok(()) + })?; } else if self.is_ast(value) { stack.push(Edge { - node: value, - parent, + node: Strong::borrowed(value), + parent: parent.clone(), key, index: -1, }); @@ -506,12 +664,12 @@ pub(crate) fn frontier_edges( let walk = Walk::new(py)?; let mut result = Vec::new(); let mut stack = Vec::new(); - walk.push_child_edges(node.as_ptr(), &mut stack)?; + walk.push_child_edges(&unsafe { Strong::borrowed(node.as_ptr()) }, &mut stack)?; while let Some(edge) = stack.pop() { - if kinds.contains(edge.node) { + if kinds.contains(&edge.node) { result.push(edge); } else { - walk.push_child_edges(edge.node, &mut stack)?; + walk.push_child_edges(&edge.node, &mut stack)?; } } Ok(result) diff --git a/src/lib.rs b/src/lib.rs index e07d18f..731968e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,7 +20,7 @@ use pyo3::exceptions::PyDeprecationWarning; use pyo3::ffi::{self, PyListObject, PyObject, PyTypeObject}; use pyo3::types::{PyList, PyModule, PyTuple, PyType}; -use fields::frontier_edges; +use fields::{Strong, frontier_edges, into_pylist}; use pyo3::{PyTypeInfo, prelude::*}; /// Open-addressed, direct-mapped lookup from `*mut PyTypeObject` to an @@ -519,10 +519,9 @@ fn walk_unordered<'py>(py: Python<'py>, node: Bound<'py, PyAny>) -> PyResult` when `index < 0`, else `parent.[index]`. -#[derive(Clone, Copy)] struct Edge { - node: *mut PyObject, - parent: *mut PyObject, + node: Strong, + parent: Strong, key: *mut PyObject, index: ffi::Py_ssize_t, } @@ -537,8 +536,7 @@ fn walk_frontier<'py>( kinds: Bound<'py, PyTuple>, ) -> PyResult> { let edges = frontier_edges(py, &node, &kinds)?; - let nodes: Vec<*mut PyObject> = edges.iter().map(|edge| edge.node).collect(); - vec_into_pylist(py, &nodes) + into_pylist(py, edges.into_iter().map(|edge| edge.node)) } /// Like `walk_frontier`, but each match comes as `(node, parent, field, @@ -558,7 +556,7 @@ fn walk_frontier_edges<'py>( } let list = Bound::from_owned_ptr(py, list_ptr); let ob_item = (*(list_ptr as *mut ffi::PyListObject)).ob_item; - for (position, edge) in edges.iter().enumerate() { + for (position, edge) in edges.into_iter().enumerate() { let index = ffi::PyLong_FromSsize_t(edge.index); if index.is_null() { return Err(PyErr::fetch(py)); @@ -568,10 +566,10 @@ fn walk_frontier_edges<'py>( ffi::Py_DECREF(index); return Err(PyErr::fetch(py)); } - for (slot, value) in [edge.node, edge.parent, edge.key].into_iter().enumerate() { - ffi::Py_INCREF(value); - ffi::PyTuple_SET_ITEM(item, slot as ffi::Py_ssize_t, value); - } + ffi::Py_INCREF(edge.key); + ffi::PyTuple_SET_ITEM(item, 0, edge.node.into_ptr()); + ffi::PyTuple_SET_ITEM(item, 1, edge.parent.into_ptr()); + ffi::PyTuple_SET_ITEM(item, 2, edge.key); ffi::PyTuple_SET_ITEM(item, 3, index); *ob_item.add(position) = item; } diff --git a/tests/test_fields.py b/tests/test_fields.py index 023f329..e732204 100644 --- a/tests/test_fields.py +++ b/tests/test_fields.py @@ -285,3 +285,94 @@ def test_refcount_neutral(): walk_frontier_events(tree, (ast.Name,), EXPRESSIONS) gc.collect() assert [sys.getrefcount(n) for n in nodes] == before + + +def name(id_: str = "x") -> ast.Name: + return ast.Name(id=id_, ctx=ast.Load()) + + +class Computed(ast.expr): + _fields = ("child",) + + @property + def child(self): + return name("computed") + + +def test_a_field_a_property_computes_survives_the_walk(): + node = Computed() + found = walk_of_types(node, (ast.Name,)) + assert [n.id for n in found] == ["computed"] + assert [n.id for n in walk_frontier_events(node, (ast.Name,), ())[0][1:]] == ["computed"] + assert len({subtree_hash(Computed()) for _ in range(3)}) == 1 + + +class Dropping(ast.expr): + _fields = ("first", "second") + + @property + def second(self): + self.__dict__.pop("first", None) + + +def test_a_field_read_may_drop_the_only_reference_to_a_child(): + node = Dropping() + node.first = name("first") + assert [n.id for n in walk_of_types(node, (ast.Name,))] == ["first"] + node = Dropping() + node.first = name("first") + set_parents(node) + + +def test_a_type_whose_address_is_reused_is_read_by_its_own_fields(): + for index in range(300): + kind = type(f"Kind{index}", (ast.expr,), {"_fields": (f"field{index}",)}) + node = kind() + setattr(node, f"field{index}", name(str(index))) + assert [n.id for n in walk_of_types(node, (ast.Name,))] == [str(index)] + del kind, node + gc.collect() + + +class MyAssign(ast.Assign): + pass + + +class Deeper(MyAssign): + pass + + +def test_a_subclass_of_any_depth_is_a_node(): + statement = Deeper(targets=[name("a")], value=name("b")) + tree = ast.Module(body=[statement], type_ignores=[]) + assert walk_of_types(tree, (ast.Name,)) == list(statement.targets) + [statement.value] + assert walk_of_types(tree, (Deeper,)) == [statement] + set_parents(tree) + assert statement.parent is tree + + +@pytest.mark.parametrize( + ("left", "right"), + [ + (lambda: float("nan"), lambda: float("nan")), + (lambda: (1, float("nan")), lambda: (1, float("nan"))), + (lambda: complex(float("nan"), 0), lambda: complex(float("nan"), 0)), + (lambda: 0.0, lambda: -0.0), + (lambda: 0j, lambda: -0j), + (lambda: [1], lambda: [1]), + (lambda: [1], lambda: [2]), + ], +) +def test_constant_values_hash_as_ast_dump_compares_them(left, right): + a, b = ast.Constant(value=left()), ast.Constant(value=right()) + assert (subtree_hash(a) == subtree_hash(b)) == (ast.dump(a) == ast.dump(b)) + + +class Defaulted(ast.AST): + _fields = ("child",) + child = name("default") + + +def test_a_node_with_no_instance_dict_reads_its_fields_by_attribute(): + node = Defaulted.__new__(Defaulted) + assert walk_of_types(node, (ast.Name,)) == list(ast.iter_child_nodes(node)) From a46b584ceae0046deeedf40503ed35b70e894bfe Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 20:23:06 +0000 Subject: [PATCH 5/5] Release 0.3.0 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hv7iDZqYjqWUscP5LcdwfB --- Cargo.lock | 2 +- Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 429eb88..2dd2439 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4,7 +4,7 @@ version = 4 [[package]] name = "fast-walk" -version = "0.2.2" +version = "0.3.0" dependencies = [ "pyo3", ] diff --git a/Cargo.toml b/Cargo.toml index 6b16a40..245c66a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "fast-walk" -version = "0.2.2" +version = "0.3.0" edition = "2024" license = "MIT" readme = "README.md"