Making ast.walk 220x Faster让 ast.walk 快速 220 倍

In our AI reflex-app builder we generate massive amounts of Python code. Sometimes, this code generation fails in rather trivial manners; positional parameters after keyword ones, returns with values in async generators, using outdated syntax conventions from previous versions of our framework, etc.在我们的 AI reflex‑app 构建器中，我们会生成大量的 Python 代码。有时，这些代码生成会因一些非常琐碎的错误而失败；关键字参数后面出现位置参数、在 async 生成器中返回值、使用我们框架旧版本的语法约定等。

Running reflex compile will eventually find all of those bugs, but it only finds one issue at a time. That means if the AI made multiple mistakes, we are increasing the latency massively for what could be relatively simple fixes.运行 reflex compile 最终会发现所有这些错误，但一次只能找到一个问题。这意味着如果 AI 出了多个错误，我们就会大幅增加延迟，而这些本可以相对简单地修复。

As such, we decided using a linter would be the best approach to fix this. And since we need to add reflex-specific rules, we couldn't use an existing one, and had to build our own.因此，我们决定使用 linter 来解决这个问题。由于需要添加 reflex‑specific 规则，现有的 linter 无法满足需求，只能自行构建。

Our initial linter looked simple:我们最初的 linter 看起来很简单：

def is_background_event(node): ...


class NoBackgroundEvent:
    def lint(self, root: ast.Module):
        violations = []
        for node in ast.walk(root):
            if isinstance(node, ast.FunctionDef) and any(
                is_background_event(decorator) for decorator in node.decorator_list
            ):
                violations.append(node)
        return violations

Unfortunately, this approach ran into performance problems quickly, as we are processing a lot of generated code.不幸的是，这种做法很快就遇到了性能问题，因为我们要处理大量生成的代码。

Notably, the slowest part in the above code is not the isinstance checks, it is ast.walk. Of course, there are many ways of optimizing things in ways that aren't making ast.walk faster, and we implemented those "lower-hanging-fruit" changes first. However, we quickly realized it was hard to make the linter significantly faster without taking on the challenge of making walk itself faster.值得注意的是，上面代码中最慢的部分不是 isinstance 检查，而是 ast.walk。当然，有很多优化方式并没有让 ast.walk 本身更快，我们先实现了那些“低垂的果实”。然而，我们很快意识到，如果不挑战让 walk 本身更快，就很难显著提升 linter 的速度。

However, walking an abstract tree (which for the purposes of this code, is just a regular tree) doesn't have to be slow. Walking the difflib module took ~2ms on my device, for ~7,000 nodes. On its own, this isn't horrible, but it quickly adds up. A rough calculation gives us ~285 nanoseconds per node, on the order of a thousand CPU cycles - far more than such a simple traversal should need. So what on earth is ast.walk doing?不过，遍历抽象树（对这段代码而言，就是普通树）不一定要慢。对 difflib 模块的遍历在我的设备上大约用了 2 ms，遍历约 7,000 个节点。单独来看这并不算糟糕，但累计起来就会很快增加。粗略计算得到每个节点约 285 ns，约千个 CPU 周期——远高于如此简单遍历应有的开销。那么 ast.walk 到底在干什么？

def walk(node):
    """
    Recursively yield all descendant nodes in the tree starting at *node*
    (including *node* itself), in no specified order.  This is useful if you
    only want to modify nodes in place and don't care about the context.
    """
    from collections import deque

    todo = deque([node])
    while todo:
        node = todo.popleft()
        todo.extend(iter_child_nodes(node))
        yield node

First thing to note here is the use of yield. Generators and yield syntax are a powerful feature of Python, but they come at a sharp cost: suspending the execution of the loop and unsuspending it repeatedly in a hot path where we are consuming the full list anyways. Sure, it saves memory, but that's not our problem at the moment, especially if that list would get cleaned up. If we simply store a list, and keep appending to it, we can minimize this:首先要注意的是 yield 的使用。生成器和 yield 语法是 Python 的强大特性，但它们代价高昂：在热点路径中不断挂起和恢复循环，而我们实际上已经把整个列表全部消费掉了。确实可以节省内存，但这不是我们当前关注的问题，尤其是如果列表会被清理的话。如果我们直接存储一个列表并不断追加，就可以把这部分开销降到最低：

def walk_helper(node: ast.AST, nodes: list[ast.AST]):
    nodes.append(node)
    for value in iter_child_nodes(node):
        walk_helper(value, nodes)


def walk(node: AST) -> list[AST]:
    nodes: list[AST] = []
    walk_helper(node, nodes)
    return nodes

But after running it, I realized it only gets us, 5% improvement. Much less than I expected, and it makes me wonder: what is iter_child_nodes doing?但运行后我发现这只能提升约 5%。远低于预期，这让我好奇：iter_child_nodes 到底在做什么？

def iter_child_nodes(node):
    """
    Yield all direct child nodes of *node*, that is, all fields that are nodes
    and all items of fields that are lists of nodes.
    """
    for name, field in iter_fields(node):
        if isinstance(field, AST):
            yield field
        elif isinstance(field, list):
            for item in field:
                if isinstance(item, AST):
                    yield item

Ah, another generator. If we inline it, we should see some decent performance improvements:啊，又是一个生成器。如果我们把它内联，应该能看到相当可观的性能提升：

def walk_helper(node: ast.AST, nodes: list[ast.AST]):
    nodes.append(node)
    for name, field in iter_fields(node):
        if isinstance(field, AST):
            walk_helper(field, nodes)
        elif isinstance(field, list):
            for item in field:
                if isinstance(item, AST):
                    walk_helper(item, nodes)

We do get the improvement we are looking for, around 25%. That begs the question: where is the remaining 75%? And while we're on the topic, what is iter_fields? Hello Python?我们确实得到了大约 25% 的提升。这就引出了问题：剩下的 75% 到底在哪里？顺便说一下，iter_fields 是什么？你好，Python？

def iter_fields(node):
    """
    Yield a tuple of ``(fieldname, value)`` for each field in ``node._fields``
    that is present on *node*.
    """
    for field in node._fields:
        try:
            yield field, getattr(node, field)
        except AttributeError:
            pass

Here we go again, another generator. We are also yielding a tuple, which we don't end up using, as we only care about the value, not the name.又来了，又是一个生成器。我们还在产生一个元组，但实际上我们只关心值，不需要名字。

getattr(node, field, None) should be faster, as the exception handling can be faster when delegated to CPython.getattr(node, field, None) 应该更快，因为异常处理交给 CPython 处理会更高效。

Let's combine both of those:把这两点结合起来：

def walk_helper(node: ast.AST, nodes: list[ast.AST]):
    nodes.append(node)
    for name in node._fields:
        field = getattr(node, name, None)
        if isinstance(field, AST):
            walk_helper(field, nodes)
        elif isinstance(field, list):
            for item in field:
                if isinstance(item, AST):
                    walk_helper(item, nodes)

That gets us to around 50% cumulative improvement. This removes the last function we're calling from the ast modules, so it's only us to blame here.这样累计提升约 50%。我们已经把 ast 模块中最后一个被调用的函数去掉了，唯一的责任就在我们这边。

Twice as fast isn't bad, but is walking at 2x speed called sprinting? I don't think so. Let's push this further.提升两倍已经不错，但把遍历速度提升 2 倍称为 sprint 吗？我不这么认为。让我们再进一步。

We can read _fields and check the subclassing in the same call. No other object that exists within ._fields can appear in a well-formed ast tree, so we should be safe.我们可以一次性读取 _fields 并检查子类关系。只要在一个合法的 ast 树中，._fields 中出现的对象一定是合法的，所以这样做是安全的。

def walk_helper(node: AST, fields: list[str], nodes: list[AST]):
    nodes.append(node)
    for field in fields:
        value = getattr(node, field, None)
        if type(value) is list:
            for item in value:
                if (subfields := getattr(item, "_fields", None)) is not None:
                    walk_helper(item, subfields, nodes)
        elif (subfields := getattr(value, "_fields", None)) is not None:
            walk_helper(value, subfields, nodes)

That only adds another incremental improvement though, maybe 55% cumulative. At this point I'm reaching the ends of what I can do with Python. Making this iterative instead of recursive barely moves the needle.这只能再带来一点点增益，累计大约 55%。此时我已经触及 Python 能做到的极限。把递归改为迭代几乎没有效果。

But Python has one last trick up its sleeve: bindings. It allows us to write this logic in native machine code (or something that would compile to it). While I could use C here, I opted to use rust just because that's what I'm used to. Let's do a simple transliteration:但 Python 还有最后一招：绑定（bindings）。它允许我们用本机机器码（或能编译成机器码的代码）来实现这段逻辑。虽然可以用 C 实现，我还是选用了 Rust，因为我更熟悉它。下面是一个简单的翻译示例：

fn walk_node<'py>(
    node: Bound<'py, PyAny>,
    field_names: Bound<'py, PyTuple>,
    result_list: &mut Vec<Bound<'py, PyAny>>,
) -> PyResult<()> {
    result_list.push(node.clone());
    for field in field_names {
        if let Ok(Some(child)) = node.getattr_opt(unsafe { field.cast_unchecked() }) {
            if child.is_exact_instance_of::<PyList>() {
                for item in unsafe { child.cast_unchecked::<PyList>() } {
                    if let Ok(Some(subfields)) = item.getattr_opt(intern!(item.py(), "_fields")) {
                        walk_node(item, unsafe { subfields.cast_into_unchecked() }, result_list)?;
                    }
                }
            } else if let Ok(Some(subfields)) = child.getattr_opt(intern!(child.py(), "_fields")) {
                walk_node(child, unsafe { subfields.cast_into_unchecked() }, result_list)?;
            }
        }
    }
    Ok(())
}

I'm using cast_unchecked as PyO3 would do heavy type checking if I used regular casting. Sometimes that type checking is useful, but not here. Other than that, the above code is not so different, getattr_opt is getattr(..., ..., None).我使用 cast_unchecked，因为如果使用普通的强制转换，PyO3 会进行大量类型检查。有时这种检查是有用的，但这里不需要。除此之外，上面的代码并没有太大区别，getattr_opt 就是 getattr(..., ..., None)。

That gets around 78% cumulative improvement. Nice!这样累计提升约 78%。不错！

This also allows us to do something more interesting. See, since we are doing a lot of getattr, which compiles down to reading the dictionary, we can simply iterate over the dictionary itself. In Python, that dictionary is called __dict__. We can simply read it at a memory offset:这还让我们可以做更有趣的事。因为我们大量使用 getattr，而 getattr 最终会读取字典，我们可以直接遍历该字典本身。在 Python 中，这个字典叫 __dict__。我们可以直接在内存偏移处读取它：

fn get_instance_dict_fast(obj: *mut PyObject) -> Option<*mut PyObject> {
    unsafe {
        let dict_offset = (*obj).ob_type.as_ref()?.tp_dictoffset;   // 1
        if dict_offset != 0 {
            let dict_ptr_addr =
                (obj as *mut u8).offset(dict_offset as isize) as *mut *mut ffi::PyObject;  // 2
            let dict = *dict_ptr_addr;                              // 3
            if !dict.is_null() {
                return Some(dict);
            }
        }
        None
    }
}

We can also improve our subclass checking. There's only 132 classes that subclass ast.AST, so instead of a real isinstance call we can store the memory addresses of all of those classes in a set and simply check membership. (I first reached for fastset here, but it's built for small, dense integers and panics on 64-bit pointer values, so a plain hash set it is.)我们还能改进子类检查。ast.AST 只有 132 个子类，所以我们可以把这些类的内存地址存入一个集合，然后直接检查成员关系。（我最初想用 fastset，但它只适用于小的、密集的整数，并且在 64 位指针值上会 panic，于是改用普通的哈希集合。）

fn isinstance_of_ast(obj: *mut PyObject, all_ast_classes: &Set) -> bool {
    let el = unsafe { ffi::Py_TYPE(obj) };          // the object's exact type
    all_ast_classes.contains(&(el as usize))         // is that type pointer in the set?
}

Then we combine both of those:然后把两者结合起来：

fn walk_node(node: *mut PyObject, all_ast_classes: &Set,
             py_list_type: *mut PyTypeObject, result_list: &mut Vec<*mut PyObject>) -> PyResult<()> {
    result_list.push(node);
    let Some(dict) = get_instance_dict_fast(node) else { return Ok(()) };
    for item_ptr in BorrowedDictIter::new(dict) {            // PyDict_Next under the hood
        if isinstance_of_ast(item_ptr, all_ast_classes) {    // set.contains(Py_TYPE(obj))
            walk_node(item_ptr, all_ast_classes, py_list_type, result_list)?;
        } else if isinstance_of_list(item_ptr, py_list_type) {
            for i in 0..get_length_of_list(item_ptr) {
                let item_ptr = get_item_of_list(item_ptr, i);
                if isinstance_of_ast(item_ptr, all_ast_classes) {
                    walk_node(item_ptr, all_ast_classes, py_list_type, result_list)?;
                }
            }
        }
    }
    Ok(())
}

That gets us to ~93%. In total, that's ~14 times faster.这样累计提升约 93%。总体来说快了约 14 倍。

The only CPython call we have left is inside of BorrowedDictIter which calls PyDict_Next. That function isn't that slow, but it does a lot of checking and reference-counting that our Rust-brain cannot comprehend. So let's rewrite it!我们唯一剩下的 CPython 调用在 BorrowedDictIter 中，它会调用 PyDict_Next。这个函数本身并不慢，但会进行大量检查和引用计数，Rust 代码无法直接处理。所以我们把它重写！

impl Iterator for DictValuesIter {
    type Item = *mut PyObject;
    fn next(&mut self) -> Option<Self::Item> {
        while self.current < self.end {
            let entry = &unsafe { *self.entries.add(self.current) };
            self.current += 1;
            if !entry.me_value.is_null() { return Some(entry.me_value); }
        }
        None
    }
}

Another useful idea is that all ast.AST subclasses are in a relatively short chain. The longest is two hops (BinOp -> expr -> AST).另一个有用的想法是，所有 ast.AST 子类形成的链条相对较短。最长的只有两层（BinOp → expr → AST）。

unsafe fn is_subtype(subtype: *mut PyTypeObject, base: *mut PyTypeObject) -> bool {
    if subtype == base { return true; }
    let mut current = subtype;
    for _ in 0..2 {
        current = (*current).tp_base;
        if current.is_null() { return false; }
        if current == base { return true; }
    }
    false
}

That gets us to around 99%. That's two orders of magnitude! We're getting close to the limit now.这样累计提升约 99%。提升了两个数量级！我们已经接近极限。

One observation we previously had is that there's a small amount of AST subclasses. What if we cached, for each subclass, two pieces of information: whether this is an AST subclass, and if so, the number of elements in its _fields.我们之前观察到 AST 子类数量很少。如果我们为每个子类缓存两条信息：它是否是 AST 子类，以及如果是的话，它的 _fields 有多少元素，会怎样？

We can precompute both into a tiny direct-mapped table keyed by the type pointer:我们可以把这两项预先计算好，放进一个以类型指针为键的微型直接映射表中：

fn lookup(&self, ptr: *mut PyTypeObject) -> u8 {
    let key = ptr as u64;
    let mut idx = ((key >> 4) as usize) & FIELD_TABLE_MASK;
    loop {
        let k = unsafe { *self.keys.get_unchecked(idx) };
        if k == key { return unsafe { *self.values.get_unchecked(idx) }; }  // n_fields + 1
        if k == 0 { return 0; }                                             // not an AST type
        idx = (idx + 1) & FIELD_TABLE_MASK;
    }
}

The result is either 0 if it isn't an AST subclass, and otherwise result - 1 is the number of elements _fields. The size of this mapping is ~2KB, so it can fit inside of L1 cache.结果要么是 0（表示不是 AST 子类），要么是 result‑1，表示 _fields 的元素数量。该映射大小约 2 KB，能够放进 L1 缓存。

The last observation is that the __dict__ of AST subclasses are very predictable, the first values are the _fields then _attributes. So if we only need to scan the first len(_fields) entries (which we have precomputed), we save ourselves from wasting time checking lineno/col_offset. We also gain the ability of ignoring any user-attached .parent back-references, which some lint rules might do.最后的观察是，AST 子类的 __dict__ 非常可预测，前面的值是 _fields，随后是 _attributes。因此如果我们只扫描前 len(_fields) 项（我们已经预先计算好），就可以避免检查 lineno/col_offset 的开销。我们还能忽略任何用户附加的 .parent 反向引用，这在某些 lint 规则中会出现。

All of those combined, we get ~99.5%, which is roughly, a ~220x improvement. B)把所有这些技巧综合起来，累计提升约 99.5%，相当于约 220 倍的加速。B)

I will stop the blog here, but https://github.com/reflex-dev/fast-walk/ takes it further (with batched prefetching, it squeezes out a bit more still).我就此止笔，但 https://github.com/reflex-dev/fast-walk/ 进一步推进（通过批量预取，还能再挤出一点性能）。