Ускорение форматирования Gleam с помощью арен в Rust

Джакомо Кавальери (Giacomo Cavalieri) из основной команды Gleam (функционального языка, чей компилятор написан на Rust) закрыл трёхлетний тикет, переписав pretty printer форматировщика на выделение памяти аренами. В самом тикете задачу назвали "довольно долгой ручной работой", и так оно и вышло: pull request на +2963/-1032 строк, собранный за пару дней поиска с автозаменой. Результат - форматировщик работает примерно на 24% быстрее и потребляет на ~10% меньше пиковой памяти.

Что изменилось

Pretty printer в Gleam строит рекурсивное дерево Document, которое описывает, как должен отображаться код и где можно переносить строки. Варианты, хранящие дочерние документы (Nest и другие), держали их в Box, из-за чего каждый вложенный узел требовал отдельной аллокации в куче:

pub enum Document<'string> {
    String(&'string str),
    Break { broken: &'string str, unbroken: &'string str },
    Group(Vec<Self>),
    Nest(Box<Self>),
    // ...
}

Изменение заменяет владеющий Box на заимствованную ссылку на арену, что потребовало добавить второй lifetime в enum:

Nest(&'doc Self)   // was: Nest(Box<Self>)

Узлы аллоцируются через крейт typed_arena. arena.alloc(value) возвращает ссылку, которой borrow checker не позволит пережить саму арену, поэтому освобождать память вручную не требуется: когда арена уничтожается (drop) в конце форматирования, вместе с ней освобождается всё дерево Document. Дополнительный плюс - кэширование: сотни мелких повторяющихся документов (ключевые слова вроде String("fn"), запятая - разделитель списков) аллоцируются один раз и переиспользуются по ссылке повсюду вместо повторной упаковки в Box при каждом появлении.

Цена изменений оказалась эргономической, а не концептуальной. Каждая функция, которая раньше вызывала Box::new, теперь должна принимать арену дополнительным аргументом и пробрасывать её lifetime - именно поэтому diff получился таким большим и механическим.

Цифры

Форматирование реального проекта squirrel ускорилось с 13 мс до 9.8 мс внутри pretty printer'а (примерно на 24% быстрее). Поскольку само форматирование - лишь часть gleam format (чтение и парсинг исходников занимают большую часть времени), вся команда целиком ускорилась примерно на 13%. Пиковое потребление памяти снизилось с 8.4 МБ до 7.6 МБ. Урок банален, но о нём легко забыть: сокращение аллокаций в куче делает код одновременно быстрее и легче.

В основе алгоритма pretty-printing лежит статья Кристиана Линдига (Christian Lindig) "Strictly Pretty" - подход в стиле Вадлера (Wadler).

Интересное из обсуждения

Тред на Lobsters превратился в полезное сравнение двух подходов к аренам. Кавальери использовал арену, раздающую ссылки; комментатор отметил, что арена на дескрипторах (handle'ах) - непрозрачных 16- или 32-битных индексах вместо 64-битных ссылок, как в compact_arena, - упаковала бы множество мелких узлов плотнее и могла бы увеличить выигрыш по памяти. Контраргумент от автора известного обзора арен заключается в том, что по ссылкам гораздо удобнее делать pattern matching: Rust не умеет матчить сквозь Box в рекурсивном enum'е, но прозрачно матчит по ссылкам, а единая арена даёт всему дереву один согласованный lifetime. Сам rustc использует подход с дескрипторами: он интернирует типы и преобразует AST в размещённый в арене HIR с дескрипторами HirId, а не ссылками, именно потому, что интернирование делает дедупликацию и кэширование дешёвыми. Это всё тот же компромисс между индексами и указателями, на который опираются и arena-allocation, и safe-gc.