Проект: ускоряем Python-скрипт

Самый частый сценарий, ради которого питонисты смотрят на Mojo: есть работающая программа на Python, и она медленная. Переписывать всё не хочется и не нужно. Нужно найти то место, где на самом деле уходит время, и перенести в Mojo только его.

В этом проекте мы пройдём этот путь целиком: профилировщик, перенос, замеры, несколько шагов ускорения — и честное сравнение с готовой библиотекой, про которую стоило вспомнить с самого начала.

Исходная программа

Скрипт ищет возможные опечатки: слова, которые встретились в тексте один раз и отличаются от частого слова ровно одной буквой. Мерой сходства служит расстояние Левенштейна — сколько вставок, удалений и замен букв нужно, чтобы из одного слова получить другое. «Компилятр» от «компилятор» — на расстоянии 1.

Сердце программы — классическая динамика:

def levenshtein(a, b):
    prev = list(range(len(b) + 1))
    for i, ca in enumerate(a, 1):
        cur = [i]
        for j, cb in enumerate(b, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb)))
        prev = cur
    return prev[-1]

А вокруг — двойной цикл: каждое редкое слово сравнивается с каждым частым, пары слишком разной длины отбрасываются сразу:

for word in rare:
    for candidate in common:
        if abs(len(word) - len(candidate)) > MAX_DIST:
            continue
        if levenshtein(word, candidate) <= MAX_DIST:
            found.append((word, candidate))

Запустим на текстах всех глав этого курса (735 КБ на момент написания главы):

$ python typos.py course.txt
редких слов: 3325, частых: 1559, пар-подозрений: 567
  about → abort
  bound → found
  builtins → builtin
  closure → closures
  ...
поиск занял 15.93 с

В русском тексте находки — в основном формы одного слова: «аргументу → аргумента», «данным → данных». Как проверка орфографии скрипт наивен, но как пример узкого места — то, что нужно: 16 секунд на пол-мегабайта текста.

Замеры в главе — на двух машинах: облачный сервер с Intel Xeon 2,1 ГГц (Emerald Rapids) и AMD Ryzen 7 9700X. Python 3.14.7 на обеих.

	Xeon	Ryzen
весь скрипт, чистый Python	16,1 с	7,4 с

Шаг 0. Где уходит время

Не угадывайте — измерьте. В Python для этого есть встроенный профилировщик:

python -m cProfile -s tottime typos.py course.txt

Главное из его отчёта:

         228669178 function calls (228669162 primitive calls) in 83.909 seconds

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
  1737013   49.283    0.000   78.430    0.000 typos.py:17(levenshtein)
104816431   15.104    0.000   15.104    0.000 {built-in method builtins.min}
104817017   13.842    0.000   13.842    0.000 {method 'append' of 'list' objects}
        1    3.518    3.518   83.789   83.789 typos.py:29(find_typos)
        …
        1    0.059    0.059    0.062    0.062 {method 'read' of '_io.TextIOWrapper' objects}
        1    0.042    0.042    0.042    0.042 {method 'findall' of 're.Pattern' objects}

Под профилировщиком программа идёт в пять раз медленнее — 84 секунды вместо 16, — но пропорции он показывает верно. 93 % времени уходит в levenshtein и то, что она вызывает. Её вызывают 1,7 миллиона раз, а внутри — по 105 миллионов вызовов min и append. Чтение файла, подсчёт слов и сортировка без профилировщика занимают 0,03 с. Переносить будем только поиск пар.

Прежде чем переносить — подумайте об алгоритме

Нам нужно знать только одно: расстояние не больше 1 или нет. Значит, как только все числа в очередной строке таблицы превысили порог, считать дальше бессмысленно — они только растут. Эта проверка в две строки ускоряет чистый Python в 2,8 раза (5,7 с на Xeon и 2,6 с на Ryzen). Хороший алгоритм выигрывает у любого языка. Мы ещё вернёмся к этой идее на шаге 3.

Шаг 1. Та же функция на Mojo

Модуль расширения Python на Mojo мы уже писали в главе «Mojo из Python». Здесь то же самое: функция принимает и возвращает PythonObject, а PyInit_fuzzy регистрирует её в модуле.

def to_codepoints(obj: PythonObject) raises -> List[UInt32]:
    """Строка Python → коды символов: сравниваем буквы, а не байты UTF-8."""
    var out = List[UInt32]()
    for c in String(py=obj).codepoints():
        out.append(c.to_u32())
    return out^


def levenshtein(py_a: PythonObject, py_b: PythonObject) raises -> PythonObject:
    var a = to_codepoints(py_a)
    var b = to_codepoints(py_b)
    var rows = List[Int](length=2 * (len(b) + 1), fill=0)
    return PythonObject(distance(a, b, rows))

to_codepoints важна: строки в Mojo хранятся в UTF-8, и русская буква занимает два байта. «ё» — это D1 91, «е» — D0 B5: они различаются обоими байтами, и побайтово «ёж» и «еж» отстояли бы на 2, а не на 1. Поэтому слова сперва превращаются в списки кодов символов, как строки Python.

`String(py=obj)`, а не `String(obj)`

Обе записи превращают строку Python в строку Mojo, но по-разному реагируют на строку, которую нельзя записать в UTF-8 — например, с одиночным суррогатом "\ud800", какие остаются после неудачного декодирования. String(py=obj) бросает обычное исключение, и Python его поймает. String(obj) роняет весь процесс Python:

ABORT: …/python_object.mojo:1388:18: failed to write PythonObject to writer:
'utf-8' codec can't encode character '\ud800' in position 0: surrogates not allowed

Сама distance — та же динамика, что на Python, только две строки таблицы лежат в одном буфере rows, а «поменять строки местами» значит поменять два смещения. Полностью код — в конце главы.

В Python появляется импорт модуля, а в цикле вместо своей levenshtein вызывается модульная:

import mojo.importer  # учит Python импортировать .mojo-файлы
import fuzzy

... fuzzy.levenshtein(word, candidate) ...

поиск пар	Xeon	Ryzen	ускорение
чистый Python	15,93 с	7,48 с	—
fuzzy.levenshtein, цикл на Python	2,21 с	1,11 с	×7

В семь раз — неплохо, но функция на Mojo должна быть быстрее Python в десятки раз. Где остальное?

Во что обходится вызов

Каждый вызов fuzzy.levenshtein из Python — это сам переход в Mojo, перевод двух строк Python в коды символов, выделение буфера и только потом расчёт. Разложим время по частям: оставим тот же цикл на Python с его 1,74 миллиона вызовов и будем вызывать функции, которые делают всё меньше (Xeon):

что делает функция на Mojo	время цикла
цикл на Python без вызова	0,29 с
ничего, возвращает число	0,33 с
переводит обе строки в коды символов	1,30 с
… и выделяет буфер rows	1,39 с
… и считает расстояние (fuzzy.levenshtein)	2,42 с

Сам переход из Python в Mojo почти ничего не стоит — около 0,02 мкс. Дорого другое: перевод строк — около 0,56 мкс на вызов, почти секунда на весь цикл, — и сам цикл на Python, 0,29 с.

Значит, нужно перенести в Mojo весь двойной цикл и переводить слова один раз. Функция all_pairs получает оба списка, переводит их в коды символов на входе и возвращает готовый список пар:

def all_pairs(py_rare: PythonObject, py_common: PythonObject, py_max: PythonObject) raises -> PythonObject:
    var max_dist = Int(py=py_max)
    var rare = to_codepoint_lists(py_rare)
    var common = to_codepoint_lists(py_common)
    var rows = List[Int](length=2 * (longest(common) + 1), fill=0)
    var out = Python.list()
    for i in range(len(rare)):
        for j in range(len(common)):
            if abs(len(rare[i]) - len(common[j])) > max_dist:
                continue
            if distance(rare[i], common[j], rows) <= max_dist:
                out.append(Python.tuple(py_rare[i], py_common[j]))
    return out

поиск пар	Xeon	Ryzen	ускорение
fuzzy.levenshtein, цикл на Python	2,21 с	1,11 с	×7
промежуточный вариант: similar() на каждое слово	2,18 с	1,24 с	×6–7
all_pairs() — весь цикл на Mojo	1,01 с	0,44 с	×16–17

Промежуточный вариант поучителен: similar(word, common) вызывается всего 3325 раз, по разу на редкое слово, но работает не быстрее первого — потому что переводит все 1559 частых слов в коды символов заново для каждого редкого. Дорог не вызов, а перевод данных.

Правило, которое стоит запомнить

Переносите в Mojo не функцию, а цикл, который её вызывает, и переводите данные в представление Mojo один раз, на входе. Вызов с переводом двух коротких строк стоит около полумикросекунды, и миллион таких вызовов — это полсекунды, которую никакая скорость Mojo не вернёт.

Шаг 2. Без проверок границ

Это мы уже видели в проекте «Матричная библиотека»: индексация List проверяет выход за границы, и во внутреннем цикле это дорого. Та же функция через unsafe_ptr() и unsafe_load / unsafe_store (distance_fast в полном коде):

поиск пар	Xeon	Ryzen	ускорение
all_pairs()	1,01 с	0,44 с	×16–17
без проверок границ	0,35 с	0,19 с	×40–45

Шаг 3. Ранний выход

Теперь вспомним совет из шага 0: как только вся строка таблицы превысила порог, ответ уже известен.

for i in range(1, len(a) + 1):
    # ... заполняем строку cur, заодно считаем её минимум row_min
    if cutoff and row_min > max_dist:
        return max_dist + 1

поиск пар	Xeon	Ryzen	ускорение
без проверок границ	0,35 с	0,19 с	×40–45
и с ранним выходом	0,14 с	0,07 с	×101–117

Мы посчитали: в среднем счёт обрывается на 2,2-й строке таблицы из 7,6 — большинство пар непохожи с первых букв.

Шаг 4. Другой алгоритм: битовый метод Майерса

Классическая динамика считает таблицу len(a) × len(b) клетка за клеткой. В 1999 году Джин Майерс показал, как считать целый столбец таблицы за раз — битовыми операциями над 64-битным числом; мы используем вариант этого метода, который Хейкки Хюрё приспособил для расстояния между двумя словами. Идея в том, что соседние клетки столбца отличаются не больше чем на 1. Поэтому столбец можно хранить не числами, а двумя масками: в каких строках значение на 1 больше, чем строкой выше (vp), и в каких на 1 меньше (vn). Сложение и сдвиги масок пересчитывают весь столбец сразу.

def myers(pattern_len: Int, peq: List[UInt64], text: List[UInt32]) -> Int:
    if pattern_len == 0:
        return len(text)
    var top = UInt64(1) << UInt64(pattern_len - 1)
    var vp = ~UInt64(0)
    var vn = UInt64(0)
    var dist = pattern_len
    var pp = peq.unsafe_ptr()
    var pt = text.unsafe_ptr()
    for j in range(len(text)):
        var eq = pp.unsafe_load(Int(pt.unsafe_load(j)))
        var x = eq | vn
        var d0 = (((x & vp) + vp) ^ vp) | x
        var hp = vn | ~(d0 | vp)
        var hn = vp & d0
        if hp & top:
            dist += 1
        if hn & top:
            dist -= 1
        x = (hp << 1) | 1
        vn = x & d0
        vp = (hn << 1) | ~(x | d0)
    return dist

Здесь peq[c] — маска позиций, где в первом слове стоит символ c. Её строим один раз на каждое редкое слово. Чтобы искать маску по индексу, а не по словарю, всем буквам заранее выдаются небольшие номера. Одна буква второго слова обрабатывается десятком битовых операций. Маска — 64 бита, поэтому в нашей реализации первое слово не длиннее 64 букв; более длинные слова считаются обычной динамикой из шага 3.

Это тот самый случай, где Mojo особенно хорош: UInt64 — настоящее машинное число, и каждая операция здесь — одна инструкция процессора. В Python каждая операция — это работа с объектом int. Мы проверили: тот же метод на чистом Python (с масками, построенными один раз на слово, как в Mojo) ищет пары за 4,7 с на Xeon. Это в 3,4 раза быстрее исходной динамики, но лишь в 1,2 раза быстрее динамики с ранним выходом, тогда как в Mojo тот же переход дал больше чем вдвое.

поиск пар	Xeon	Ryzen	ускорение
с ранним выходом	0,14 с	0,07 с	×101–117
битовый метод Майерса	0,059 с	0,032 с	×234–270

А что, если поискать готовое?

Для расстояния Левенштейна давно есть библиотека rapidfuzz на C++. Устанавливается одной командой:

pip install rapidfuzz

поиск пар	Xeon	Ryzen	ускорение
Levenshtein.distance, цикл на Python	0,50 с	0,24 с	×31–32
process.cdist, 1 поток	0,074 с	0,034 с	×215–220
наш all_pairs_myers, 1 поток	0,059 с	0,032 с	×234–270

Внутри rapidfuzz — те же идеи, только доведённые дальше. Мы заглянули в исходники: для одной пары с маленьким порогом (score_cutoff < 4) она отрезает общее начало и конец слов и перебирает несколько вариантов правок, а cdist считает вариантом Хюрё сразу много слов одной SIMD-инструкцией. С workers=-1 cdist занимает оба ядра и на Xeon укладывается в 0,030 с — быстрее нашей однопоточной версии.

Наша версия на одном потоке вровень с rapidfuzz на Ryzen и быстрее в 1,25 раза на Xeon. Но rapidfuzz получен одной строкой pip install. Обратите внимание и на первую строку таблицы: даже библиотека на C++, вызванная из цикла на Python, в семь раз медленнее своего же cdist: больше половины этого времени — сам цикл на Python.

Честный вывод

Если задачу решает готовая библиотека — начните с неё. Mojo нужен там, где готового нет: своя формула сходства, особые правила для букв, нестандартная обработка данных. Тогда весь путь из этой главы — профилировщик, перенос цикла целиком, снятие проверок, лучший алгоритм — даёт то же, что дал здесь.

Итог

Весь скрипт целиком, от запуска до выхода:

	Xeon	Ryzen
чистый Python	16,1 с	7,4 с
Python + fuzzy.all_pairs_myers	0,19 с	0,10 с
первый запуск: компиляция модуля	6,6 с	3,3 с

Ускорение программы целиком — в 75–85 раз, а не в 250. Поиск пар теперь занимает 0,03–0,06 с, чтение и подсчёт слов — ещё 0,03 с, а остальное — запуск интерпретатора Python и импорт модулей. Это закон Амдала в действии: ускорив одну часть, вы делаете заметными все остальные.

При первом запуске mojo.importer компилирует fuzzy.mojo и кладёт результат в каталог __mojocache__ рядом с модулем; дальше импорт мгновенный, пока исходник не изменится. Чтобы отдать программу другим людям без Mojo, модуль нужно собрать заранее и приложить библиотеки рантайма — это разобрано в главе «Упаковка и распространение».

Весь код

Исходная программа на Python:

typos.py

"""Ищет вероятные опечатки: редкие слова, очень похожие на частые.

Запуск: python typos.py текст.txt
"""

import re
import sys
import time
from collections import Counter

RARE = 1  # слово встретилось не больше RARE раз — подозреваемое
COMMON = 5  # слово встретилось хотя бы COMMON раз — образец
MAX_DIST = 1  # сколько правок допускаем
MIN_LEN = 5  # короткие слова слишком похожи друг на друга


def levenshtein(a, b):
    """Расстояние Левенштейна: сколько вставок, удалений и замен букв
    нужно, чтобы из a получить b."""
    prev = list(range(len(b) + 1))
    for i, ca in enumerate(a, 1):
        cur = [i]
        for j, cb in enumerate(b, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb)))
        prev = cur
    return prev[-1]


def find_typos(counts):
    common = sorted(w for w, c in counts.items() if c >= COMMON and len(w) >= MIN_LEN)
    rare = sorted(w for w, c in counts.items() if c <= RARE and len(w) >= MIN_LEN)
    found = []
    for word in rare:
        for candidate in common:
            if abs(len(word) - len(candidate)) > MAX_DIST:
                continue
            if levenshtein(word, candidate) <= MAX_DIST:
                found.append((word, candidate))
    return rare, common, found


def main():
    text = open(sys.argv[1], encoding="utf-8").read().lower().replace("\\n", " ")
    counts = Counter(re.findall(r"[а-яёa-z]+", text))
    start = time.perf_counter()
    rare, common, found = find_typos(counts)
    elapsed = time.perf_counter() - start
    print(f"редких слов: {len(rare)}, частых: {len(common)}, пар-подозрений: {len(found)}")
    for word, candidate in found[:10]:
        print(f"  {word} → {candidate}")
    print(f"поиск занял {elapsed:.2f} с")


if __name__ == "__main__":
    main()

Модуль на Mojo — со всеми шагами из главы, чтобы их можно было сравнить:

fuzzy.mojo

"""Модуль для Python: поиск похожих слов по расстоянию Левенштейна.

Все шаги ускорения из главы — отдельными функциями, чтобы их можно было
сравнить. В настоящей программе нужна только последняя, all_pairs_myers.
"""

from std.os import abort
from std.python import Python, PythonObject
from std.python.bindings import PythonModuleBuilder


@export
def PyInit_fuzzy() abi("C") -> PythonObject:
    try:
        var m = PythonModuleBuilder("fuzzy")
        m.def_function[levenshtein](
            "levenshtein", docstring="Расстояние между двумя словами"
        )
        m.def_function[similar](
            "similar", docstring="Слова из списка, похожие на данное"
        )
        m.def_function[all_pairs](
            "all_pairs", docstring="Все пары похожих слов"
        )
        m.def_function[all_pairs_unchecked](
            "all_pairs_unchecked", docstring="То же без проверок границ"
        )
        m.def_function[all_pairs_cutoff](
            "all_pairs_cutoff", docstring="То же с ранним выходом"
        )
        m.def_function[all_pairs_myers](
            "all_pairs_myers", docstring="То же битовым алгоритмом Майерса"
        )
        return m.finalize()
    except e:
        abort(String("не удалось создать модуль fuzzy: ", e))


# --- Общее: слова как списки кодов символов ---------------------------------


def to_codepoints(obj: PythonObject) raises -> List[UInt32]:
    """Строка Python → коды символов: сравниваем буквы, а не байты UTF-8.

    String(py=...), а не String(...): если строку Python нельзя перевести
    в UTF-8 (одиночный суррогат вроде U+D800), первое бросает обычное
    исключение, а второе роняет весь процесс."""
    var out = List[UInt32]()
    for c in String(py=obj).codepoints():
        out.append(c.to_u32())
    return out^


def to_codepoint_lists(words: PythonObject) raises -> List[List[UInt32]]:
    var out = List[List[UInt32]]()
    for w in words:
        out.append(to_codepoints(w))
    return out^


def longest(words: List[List[UInt32]]) -> Int:
    var n = 0
    for w in words:
        n = max(n, len(w))
    return n


# --- Шаг 1. Та же динамика, что в Python ------------------------------------


def distance(a: List[UInt32], b: List[UInt32], mut rows: List[Int]) -> Int:
    """Две строки таблицы лежат в одном буфере rows; prev и cur — смещения,
    и «поменять строки» значит поменять смещения."""
    var n = len(b)
    var prev = 0
    var cur = n + 1
    for j in range(n + 1):
        rows[prev + j] = j
    for i in range(1, len(a) + 1):
        rows[cur] = i
        for j in range(1, n + 1):
            var cost = 0 if a[i - 1] == b[j - 1] else 1
            rows[cur + j] = min(
                rows[prev + j] + 1,
                rows[cur + j - 1] + 1,
                rows[prev + j - 1] + cost,
            )
        swap(prev, cur)
    return rows[prev + n]


def levenshtein(py_a: PythonObject, py_b: PythonObject) raises -> PythonObject:
    """Одна пара за вызов: Python зовёт её миллион раз."""
    var a = to_codepoints(py_a)
    var b = to_codepoints(py_b)
    var rows = List[Int](length=2 * (len(b) + 1), fill=0)
    return PythonObject(distance(a, b, rows))


def similar(
    py_word: PythonObject, py_candidates: PythonObject, py_max: PythonObject
) raises -> PythonObject:
    """Одно слово против всех кандидатов за вызов."""
    var word = to_codepoints(py_word)
    var max_dist = Int(py=py_max)
    var out = Python.list()
    for py_c in py_candidates:
        var c = to_codepoints(py_c)
        if abs(len(word) - len(c)) > max_dist:
            continue
        var rows = List[Int](length=2 * (len(c) + 1), fill=0)
        if distance(word, c, rows) <= max_dist:
            out.append(py_c)
    return out


def all_pairs(
    py_rare: PythonObject, py_common: PythonObject, py_max: PythonObject
) raises -> PythonObject:
    """Весь двойной цикл за один вызов; слова переводятся в коды один раз."""
    var max_dist = Int(py=py_max)
    var rare = to_codepoint_lists(py_rare)
    var common = to_codepoint_lists(py_common)
    var rows = List[Int](length=2 * (longest(common) + 1), fill=0)
    var out = Python.list()
    for i in range(len(rare)):
        for j in range(len(common)):
            if abs(len(rare[i]) - len(common[j])) > max_dist:
                continue
            if distance(rare[i], common[j], rows) <= max_dist:
                out.append(Python.tuple(py_rare[i], py_common[j]))
    return out


# --- Шаги 2 и 3. Без проверок границ и с ранним выходом ---------------------


def distance_fast(
    a: List[UInt32],
    b: List[UInt32],
    mut rows: List[Int],
    max_dist: Int,
    cutoff: Bool,
) -> Int:
    """Та же динамика через указатели. Если cutoff — бросаем счёт, как только
    вся строка таблицы больше max_dist: дальше значения только растут."""
    var n = len(b)
    var pa = a.unsafe_ptr()
    var pb = b.unsafe_ptr()
    var r = rows.unsafe_ptr()
    var prev = 0
    var cur = n + 1
    for j in range(n + 1):
        r.unsafe_store(prev + j, j)
    for i in range(1, len(a) + 1):
        r.unsafe_store(cur, i)
        var ca = pa.unsafe_load(i - 1)
        var row_min = i
        for j in range(1, n + 1):
            var cost = 0 if ca == pb.unsafe_load(j - 1) else 1
            var v = min(
                r.unsafe_load(prev + j) + 1,
                r.unsafe_load(cur + j - 1) + 1,
                r.unsafe_load(prev + j - 1) + cost,
            )
            r.unsafe_store(cur + j, v)
            row_min = min(row_min, v)
        if cutoff and row_min > max_dist:
            return max_dist + 1
        swap(prev, cur)
    return r.unsafe_load(prev + n)


def pairs_fast(
    py_rare: PythonObject,
    py_common: PythonObject,
    py_max: PythonObject,
    cutoff: Bool,
) raises -> PythonObject:
    var max_dist = Int(py=py_max)
    var rare = to_codepoint_lists(py_rare)
    var common = to_codepoint_lists(py_common)
    var rows = List[Int](length=2 * (longest(common) + 1), fill=0)
    var out = Python.list()
    for i in range(len(rare)):
        for j in range(len(common)):
            if abs(len(rare[i]) - len(common[j])) > max_dist:
                continue
            if (
                distance_fast(rare[i], common[j], rows, max_dist, cutoff)
                <= max_dist
            ):
                out.append(Python.tuple(py_rare[i], py_common[j]))
    return out


def all_pairs_unchecked(
    py_rare: PythonObject, py_common: PythonObject, py_max: PythonObject
) raises -> PythonObject:
    return pairs_fast(py_rare, py_common, py_max, cutoff=False)


def all_pairs_cutoff(
    py_rare: PythonObject, py_common: PythonObject, py_max: PythonObject
) raises -> PythonObject:
    return pairs_fast(py_rare, py_common, py_max, cutoff=True)


# --- Шаг 4. Другой алгоритм: битовый метод Майерса --------------------------


def myers(pattern_len: Int, peq: List[UInt64], text: List[UInt32]) -> Int:
    """Расстояние Левенштейна битовым методом Майерса (в варианте Хюрё).

    Столбец таблицы динамики хранится не числами, а двумя масками по 64 бита:
    в каких строках значение на 1 больше, чем строкой выше (vp), и в каких
    на 1 меньше (vn). Один символ text обрабатывается десятком битовых
    операций — сразу для всех символов образца. peq[c] — маска позиций,
    где в образце стоит символ c. Нужно pattern_len <= 64."""
    if pattern_len == 0:
        return len(text)
    var top = UInt64(1) << UInt64(pattern_len - 1)
    var vp = ~UInt64(0)
    var vn = UInt64(0)
    var dist = pattern_len
    var pp = peq.unsafe_ptr()
    var pt = text.unsafe_ptr()
    for j in range(len(text)):
        var eq = pp.unsafe_load(Int(pt.unsafe_load(j)))
        var x = eq | vn
        var d0 = (((x & vp) + vp) ^ vp) | x
        var hp = vn | ~(d0 | vp)
        var hn = vp & d0
        if hp & top:
            dist += 1
        if hn & top:
            dist -= 1
        x = (hp << 1) | 1
        vn = x & d0
        vp = (hn << 1) | ~(x | d0)
    return dist


def encode(
    obj: PythonObject, mut letters: Dict[UInt32, UInt32]
) raises -> List[UInt32]:
    """Каждой букве — свой небольшой номер, чтобы маски искать по индексу."""
    var out = List[UInt32]()
    for c in String(py=obj).codepoints():
        var code = c.to_u32()
        var idx = letters.get(code, UInt32(len(letters)))
        if Int(idx) == len(letters):
            letters[code] = idx
        out.append(idx)
    return out^


def all_pairs_myers(
    py_rare: PythonObject, py_common: PythonObject, py_max: PythonObject
) raises -> PythonObject:
    var max_dist = Int(py=py_max)
    var letters = Dict[UInt32, UInt32]()
    var rare = List[List[UInt32]]()
    for w in py_rare:
        rare.append(encode(w, letters))
    var common = List[List[UInt32]]()
    for w in py_common:
        common.append(encode(w, letters))

    var peq = List[UInt64](length=len(letters), fill=0)
    var rows = List[Int](length=2 * (longest(common) + 1), fill=0)
    var out = Python.list()
    for i in range(len(rare)):
        ref word = rare[i]
        # Маска — 64 бита, поэтому слова длиннее считаем обычной динамикой
        var bits = len(word) <= 64
        if bits:
            for k in range(len(word)):
                peq[Int(word[k])] |= UInt64(1) << UInt64(k)
        for j in range(len(common)):
            if abs(len(word) - len(common[j])) > max_dist:
                continue
            var d: Int
            if bits:
                d = myers(len(word), peq, common[j])
            else:
                d = distance_fast(word, common[j], rows, max_dist, cutoff=True)
            if d <= max_dist:
                out.append(Python.tuple(py_rare[i], py_common[j]))
        if bits:
            for k in range(len(word)):
                peq[Int(word[k])] = 0
    return out

Итоговая программа — та же, но поиск пар в Mojo:

typos_mojo.py

"""Та же программа, но поиск пар — в модуле на Mojo (fuzzy.mojo рядом).

Запуск из окружения, где установлен mojo: python typos_mojo.py текст.txt
При первом запуске модуль компилируется — это займёт несколько секунд.
"""

import re
import sys
import time
from collections import Counter

import mojo.importer  # noqa: F401  — учит Python импортировать .mojo
import fuzzy

from typos import COMMON, MAX_DIST, MIN_LEN, RARE


def find_typos(counts):
    common = sorted(w for w, c in counts.items() if c >= COMMON and len(w) >= MIN_LEN)
    rare = sorted(w for w, c in counts.items() if c <= RARE and len(w) >= MIN_LEN)
    return rare, common, fuzzy.all_pairs_myers(rare, common, MAX_DIST)


def main():
    text = open(sys.argv[1], encoding="utf-8").read().lower().replace("\\n", " ")
    counts = Counter(re.findall(r"[а-яёa-z]+", text))
    start = time.perf_counter()
    rare, common, found = find_typos(counts)
    elapsed = time.perf_counter() - start
    print(f"редких слов: {len(rare)}, частых: {len(common)}, пар-подозрений: {len(found)}")
    for word, candidate in found[:10]:
        print(f"  {word} → {candidate}")
    print(f"поиск занял {elapsed:.2f} с")


if __name__ == "__main__":
    main()

И сравнение всех вариантов, которым получены таблицы:

compare.py

"""Все шаги из главы на одном тексте: время и совпадение с чистым Python.

Запуск из окружения, где установлен mojo:
    python compare.py текст.txt
Если установлен rapidfuzz (pip install rapidfuzz numpy), он тоже
попадёт в сравнение.
"""

import re
import sys
import time
from collections import Counter

import mojo.importer  # noqa: F401
import fuzzy
import typos

text = open(sys.argv[1], encoding="utf-8").read().lower().replace("\\n", " ")
counts = Counter(re.findall(r"[а-яёa-z]+", text))
rare, common, expected = typos.find_typos(counts)
M = typos.MAX_DIST


def levenshtein_cutoff(a, b, limit):
    """Python-версия с ранним выходом — для честности сравнения."""
    prev = list(range(len(b) + 1))
    for i, ca in enumerate(a, 1):
        cur = [i]
        for j, cb in enumerate(b, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb)))
        if min(cur) > limit:
            return limit + 1
        prev = cur
    return prev[-1]


def pairs(distance):
    """Двойной цикл на Python с данной функцией расстояния."""
    found = []
    for word in rare:
        for cand in common:
            if abs(len(word) - len(cand)) <= M and distance(word, cand) <= M:
                found.append((word, cand))
    return found


variants = [
    ("чистый Python", lambda: typos.find_typos(counts)[2], 1),
    ("Python с ранним выходом", lambda: pairs(lambda a, b: levenshtein_cutoff(a, b, M)), 1),
    ("1. levenshtein() на Mojo, цикл на Python", lambda: pairs(fuzzy.levenshtein), 3),
    ("   similar() на Mojo для каждого слова", lambda: [(w, c) for w in rare for c in fuzzy.similar(w, common, M)], 3),
    ("   all_pairs() — весь цикл на Mojo", lambda: fuzzy.all_pairs(rare, common, M), 5),
    ("2. без проверок границ", lambda: fuzzy.all_pairs_unchecked(rare, common, M), 5),
    ("3. с ранним выходом", lambda: fuzzy.all_pairs_cutoff(rare, common, M), 5),
    ("4. битовый алгоритм Майерса", lambda: fuzzy.all_pairs_myers(rare, common, M), 5),
]

try:
    import numpy as np
    from rapidfuzz import process
    from rapidfuzz.distance import Levenshtein

    def rapidfuzz_cdist():
        d = process.cdist(rare, common, scorer=Levenshtein.distance, score_cutoff=M, workers=1)
        return [(rare[i], common[j]) for i, j in zip(*np.nonzero(d <= M))]

    variants += [
        ("rapidfuzz: цикл на Python", lambda: pairs(lambda a, b: Levenshtein.distance(a, b, score_cutoff=M)), 3),
        ("rapidfuzz.cdist, 1 поток", rapidfuzz_cdist, 5),
    ]
except ImportError:
    print("(rapidfuzz не установлен — его в сравнении не будет)")

base = None
for name, run, repeat in variants:
    best = float("inf")
    for _ in range(repeat):
        start = time.perf_counter()
        result = run()
        best = min(best, time.perf_counter() - start)
    base = base or best
    same = sorted(map(tuple, result)) == sorted(expected)
    print(f"{name:42} {best:8.3f} с  ×{base / best:6.1f}  {'совпадает' if same else 'НЕ СОВПАДАЕТ'}")

Запуск — из окружения, где установлен Mojo (в uv-проекте из главы об установке):

uv run python typos.py course.txt
uv run python typos_mojo.py course.txt
uv add rapidfuzz numpy           # по желанию, для сравнения
uv run python compare.py course.txt

Текст для опытов — любой большой файл. Мы брали все главы курса, склеенные в один, на момент написания этой главы. Курс с тех пор вырос, так что у вас числа слов и пар будут другими, а соотношения — примерно теми же.

Что дальше

Последний проект — первое знакомство с GPU: те же идеи параллелизма, но на видеокарте.

🎯 Проверь себя

Как найти, что переносить в Mojo?

Профилировщиком, например python -m cProfile -s tottime script.py. Он покажет, на какие функции уходит время и сколько раз их вызывали. Профилировщик замедляет программу, но пропорции показывает верно.

Функцию перенесли в Mojo, а поиск ускорился всего в 7 раз. Почему?

Функцию по-прежнему вызывают из цикла на Python больше миллиона раз. Сам переход в Mojo дешёвый, но каждый раз заново переводятся строки — около 0,56 мкс на вызов, — да и сам цикл на Python стоит времени. Перенесите в Mojo весь цикл и переводите данные один раз — у нас это дало ×16–17 вместо ×7.

Почему нельзя сравнивать слова побайтово, раз строки в Mojo — UTF-8?

Русская буква занимает два байта, и замена одной буквы меняет один или два байта: «ё» и «е» различаются обоими. Расстояние получилось бы другим. Слова нужно сначала превратить в последовательности символов — codepoints().

Весь поиск ускорился в 250 раз, а программа целиком — только в 75–85. Почему?

Остальная часть программы осталась прежней: запуск интерпретатора, импорт модулей, чтение файла и подсчёт слов. Когда основная работа стала быстрой, они стали заметны. Это закон Амдала.

Есть готовая библиотека на C++, которая делает то же самое. Нужен ли Mojo?

Скорее всего, нет: rapidfuzz почти не уступил нашей лучшей версии, а на двух потоках обогнал её — и ставится одной командой. Mojo стоит брать там, где готового решения нет или оно не подходит.