← На главную

Reflex ускорили обход AST в 220 раз, переписав ast.walk на Rust с PyO3

16.06.2026 16:25 · hackernews

Разработчики Reflex, AI-конструктора для веб-приложений, столкнулись с типичной проблемой: их нейросеть генерирует тонны Python-кода, но часто с мелкими ошибками — неправильный порядок аргументов, return в асинхронных генераторах, устаревший синтаксис. Компилятор reflex compile ловит эти баги, но только по одному за раз. Если ошибок много, время отклика растёт катастрофически. Решили использовать линтер, но готовый не подходил — нужны были свои, специфичные для Reflex правила. Пришлось писать собственный.

Первая версия была простой, но сразу упёрлась в производительность: AI генерирует слишком много кода. Самый тормозной участок — функция ast.walk из стандартной библиотеки Python. Автор замерил: обход difflib (около 7000 узлов) занимал ~2 миллисекунды, то есть ~285 наносекунд на узел — тысячи тактов CPU для простого обхода дерева. Он начал копать.

Первым делом убрал yield в генераторах — переписал ast.walk на обычный список с накоплением. Прирост — жалкие 5%. Дальше — заглянул внутрь iter_child_nodes, которая тоже оказалась генератором. Инлайнинг дал уже 25%. Осталось 75% — пошёл глубже. Оказалось, что iter_fields — снова генератор, да ещё возвращает кортеж (имя поля, значение), хотя имя не нужно. Заменил на getattr(node, field, None) — и ещё 25%, итого 50% ускорения.

На этом возможности чистого Python закончились. Сделать итеративным вместо рекурсивного — почти не помогло. Тогда автор взял Rust и подключил PyO3. Простая транслитерация дала 78% ускорения. Но он пошёл дальше: вместо getattr начал читать __dict__ напрямую (словарь объекта), обходя вызовы Python API. Вместо isinstance — проверку по памяти: подклассов ast.AST всего 132, вычислил их адреса и сложил в хеш-множество. Ускорение — 93% (в 14 раз). Остался один вызов PyDict_Next — его тоже переписал на Rust. Затем добавил кеш: предвычислил для каждого типа, является ли он подклассом AST и сколько у него полей _fields. Таблица размером ~2 КБ помещается в L1-кэш процессора. Последний трюк: поля __dict__ в AST-узлах идут в фиксированном порядке — сначала _fields, потом _attributes и lineno/col_offset. Зная длину _fields, сканируем только нужные записи и игнорируем бэк-ссылки типа .parent.

Итог: ускорение примерно в 220 раз (99,5%). В репозитории github.com/reflex-dev/fast-walk/ добавили ещё пакетную предвыборку данных и выжали чуть больше.

Читать оригинал →