让 ast.walk 速度提升 220 倍

# 让 `ast.walk` 提速 220 倍 来源：https://reflex.dev/blog/why-ast-walk-when-you-can-ast-sprint 在我们的 AI reflex 应用构建器（https://build.reflex.dev/）中，我们会生成大量 Python 代码。有时，这种代码生成会在相当琐碎的地方失败：关键字参数后的位置参数、异步生成器中带有返回值的 return、使用之前版本框架的过时语法约定，等等。 运行 `reflex compile` 最终会找到所有这些错误，但每次只能找到一个。这意味着如果 AI 犯了多个错误，我们就会为本来可以相对简单修复的问题增加大量延迟。 因此，我们决定使用 linter 是解决这个问题的最佳方法。而且由于我们需要添加 reflex 特定的规则，我们不能使用现有的 linter，必须自己构建。 我们最初的 linter 看起来很简单： ```python # 示例代码（未在原文中给出，但通常类似） for node in ast.walk(tree): if isinstance(node, ...): fix(node) ``` 不幸的是，这种方法很快就遇到了性能问题，因为我们正在处理大量生成的代码。 值得注意的是，上述代码中最慢的部分不是 `isinstance` 检查，而是 `ast.walk`。当然，有很多优化方法并不需要让 `ast.walk` 更快，我们首先实现了那些“低垂的果实”的改动。然而，我们很快意识到，如果不挑战让 walk 本身更快，就很难让 linter 显著提速。 然而，遍历一棵抽象语法树（就这段代码而言，它只是一棵普通的树）不一定会慢。在我的设备上，遍历 `difflib` 模块大约需要 2 毫秒，处理约 7000 个节点。单独来看，这并不算糟糕，但很快就会累积起来。粗略计算得出每个节点大约 285 纳秒，相当于上千个 CPU 周期——对于如此简单的遍历来说，这个开销太大了。那么，`ast.walk` 到底在做什么？ 首先要注意的是使用了 `yield`。生成器和 `yield` 语法是 Python 的一个强大特性，但它们会带来高昂的代价：在热路径中反复挂起和恢复循环的执行，而我们实际上会消耗整个列表。当然，它节省了内存，但这不是我们目前的问题，尤其是当列表会被清理掉时。如果我们简单地存储一个列表并不断追加，就可以最小化这种开销： ```python # 使用列表代替生成器 def fast_walk(node): stack = [node] result = [] while stack: n = stack.pop() result.append(n) for child in iter_child_nodes(n): stack.append(child) return result ``` 但在运行之后，我发现只获得了 5% 的改进。这远低于我的预期，让我好奇：`iter_child_nodes` 做了什么？ 啊，又一个生成器。如果我们内联它，应该能看到一些不错的性能提升： ```python # 内联 iter_child_nodes def fast_walk(node): stack = [node] result = [] while stack: n = stack.pop() result.append(n) for field in n._fields: child = getattr(n, field, None) if isinstance(child, ast.AST): stack.append(child) elif isinstance(child, list): stack.extend(child) return result ``` 我们确实得到了期望的提升，大约 25%。这就引出了一个问题：剩下的 75% 在哪里？顺便问一下，`iter_fields` 是什么？Python，你好？ ```python # iter_fields 实际上是一个生成器 def iter_fields(node): for field in node._fields: yield field, getattr(node, field) ``` 又是一个生成器。我们还生成了一个元组，但我们最终并没有使用它，因为我们只关心值，不关心名称。 `getattr(node, field, None)` 应该更快，因为异常处理在委托给 CPython 时可以更快。 让我们把这两点结合起来： ```python # 直接使用 getattr 并避免 tuple def fast_walk(node): stack = [node] result = [] while stack: n = stack.pop() result.append(n) for field in n._fields: value = getattr(n, field, None) if isinstance(value, ast.AST): stack.append(value) elif isinstance(value, list): stack.extend(value) return result ``` 这让我们达到了大约 50% 的累积改进。这时我们已经去掉了从 `ast` 模块中调用的最后一个函数，所以剩下的只能怪我们自己了。 快了两倍并不差，但以两倍速度的行走能叫冲刺吗？我不这么认为。让我们更进一步。 我们可以在同一个调用中读取 `_fields` 并检查子类。任何在格式良好的 `ast` 树中出现在 `._fields` 里的其他对象都不可能存在，所以应该是安全的。 ```python # 直接检查 _fields 而不是 isinstance def fast_walk(node): stack = [node] result = [] while stack: n = stack.pop() result.append(n) if hasattr(n, '_fields'): for field in n._fields: value = getattr(n, field, None) if isinstance(value, ast.AST): stack.append(value) elif isinstance(value, list): stack.extend(value) # 如果节点没有 _fields（如 Module），可能不需要处理 return result ``` 但这只是增量改进，可能累积达到 55%。在这一点上，我已经接近用 Python 能做到的极限了。将其从递归改为迭代几乎没什么作用。 但 Python 还有最后一招：绑定（bindings）。它允许我们用原生机器代码（或能编译成它的东西）编写这个逻辑。虽然我可以在这里使用 C，但我选择了 Rust，因为这是我习惯的语言。让我们做一个简单的直译： ```rust // 使用 PyO3 的 Rust 绑定 use pyo3::prelude::*; fn fast_walk(node: &PyAny) -> PyResult<Vec<&PyAny>> { let mut stack = vec![node]; let mut result = vec![]; while let Some(n) = stack.pop() { result.push(n); if let Ok(fields) = n.getattr("_fields") { for field in fields.iter()? { let field_name: String = field.extract()?; if let Ok(value) = n.getattr(field_name.as_str()) { // 检查 value 是否是 ast.AST 实例或列表 if value.is_instance()? { /* 检查类型 */ } // ... } } } } Ok(result) } ``` 我使用了 `cast_unchecked`，因为如果使用常规类型转换，PyO3 会进行大量的类型检查。有时这种类型检查很有用，但在这里不是。除此之外，上面的代码并没有太大不同，`getattr_opt` 就是 `getattr(..., None)`。 这达到了大约 78% 的累积改进。不错！ 这也让我们能够做一些更有趣的事情。由于我们做了大量的 `getattr`，它最终会编译成读取字典，我们可以直接迭代字典本身。在 Python 中，这个字典叫做 `__dict__`。我们可以直接在内存偏移处读取它： ```rust // 通过直接读取 __dict__ 来加速 use pyo3::ffi; let dict = n.getattr("__dict__")?; let keys = dict.call_method0("keys")?; for key in keys.iter()? { let value = dict.get_item(key)?; // ... } ``` 我们还可以改进子类检查。只有 132 个类继承自 `ast.AST`，所以我们可以存储所有这些类的内存地址到一个集合中，然后检查成员身份，而不需要真正的 `isinstance` 调用。（我最初尝试使用 fastset（https://crates.io/crates/fastset），但它专为小的、紧凑的整数设计，在 64 位指针值上会 panic，所以使用普通的哈希集合。） 然后将两者结合： ```rust // 预计算 AST 子类的内存地址集合 let ast_subclasses: HashSet<usize> = /* 从 Python 收集 */; fn is_ast(value: &PyAny) -> bool { ast_subclasses.contains(&value.get_type().as_ptr() as usize) } ``` 这让我们达到了大约 93%。总共快了约 14 倍。 我们剩下的唯一 CPython 调用是在 `BorrowedDictIter` 内部，它调用了 `PyDict_Next`。这个函数并不慢，但它做了很多我们的 Rust 大脑无法理解的检查和引用计数。所以让我们重写它！ ```rust // 手动迭代字典，避免 PyDict_Next use pyo3::ffi::PyDict_Next; unsafe { let mut pos = 0; let mut key: *mut ffi::PyObject = std::ptr::null_mut(); let mut value: *mut ffi::PyObject = std::ptr::null_mut(); while PyDict_Next(dict_ptr, &mut pos, &mut key, &mut value) != 0 { // 直接使用 value let child = PyObject::from_borrowed_ptr(value); // ... } } ``` 另一个有用的想法是，所有 `ast.AST` 子类都在一个相对较短的继承链中。最长的链是两步（`BinOp -> expr -> AST`）。 这让我们达到了大约 99%，也就是两个数量级！我们越来越接近极限了。 我们之前注意到 AST 子类数量很少。如果我们为每个子类缓存两条信息：它是否是 `AST` 子类，如果是，它的 `_fields` 中元素的数量，会怎么样？ 我们可以将两者预先计算到一个小的直接映射表中，以类型指针为键： ```rust // 使用大小为 2048 的直接映射缓存 const CACHE_SIZE: usize = 2048; static CACHE: [u8; CACHE_SIZE] = [0; CACHE_SIZE]; // 实际上应该用 Atomic 或 UnsafeCell fn lookup(type_ptr: usize) -> u8 { let index = (type_ptr.wrapping_shr(4)) % CACHE_SIZE; CACHE[index] } ``` 结果要么是 0（如果不是 AST 子类），否则 `result - 1` 就是 `_fields` 的元素数量。这个映射的大小约为 2KB，可以放入 L1 缓存。 最后一个观察是，`AST` 子类的 `__dict__` 非常可预测，前几个值是 `_fields`，然后是 `_attributes`。所以我们只需要扫描前 `len(_fields)` 个条目（我们已经预先计算好了），就可以避免浪费时间检查 `lineno`/`col_offset`。我们还能忽略任何用户附加的 `.parent` 反向引用，某些 lint 规则可能会这样做。 所有这些结合起来，我们实现了大约 99.5% 的改进，大约提升了 220 倍。:) 这篇博客就写到这里，不过 https://github.com/reflex-dev/fast-walk 上有进一步的优化（使用批量预取，还能再挤出一点性能）。