Когда Zig круче Rust — массивы перечислений, позволяющие сэкономить память

Перечисления (или размеченные объединения), отличающиеся вариативностью и, следовательно, размером, провоцируют в Rust серьёзную фрагментацию памяти. Дело в том, что нам приходится выделять достаточно данных, чтобы их хватило на самый крупный вариант.

Рассмотрим следующее перечисление:

pub enum Foo { 
  A(u8), 
  B(u16),
  C(u32),
  D(u64),
}

Поскольку приходится предусматривать пространство, чтобы разметить и выровнять данные, этот тип имеет длину 16 байт.

Это очень мешает, когда приходится собрать очень много таких структур в Vec или HashMap. Справиться с заполнением можно, применив в той или иной форме преобразование структур в массивы (SoA); при этом для хранения метки выделяется отдельный участок. Но сократить фрагментацию вариантов уже гораздо сложнее.

Можно вручную выкатывать специализированные структуры данных для конкретного перечисления, минимизируя таким образом фрагментацию. Но, если попытаетесь унифицировать такую операцию, чтобы она была применима к любому произвольному перечислению с сохранением максимальной эффективности памяти, то увидите, что в Rust это практически невозможно. Единственный вариант в данном случае — воспользоваться процедурными макросами, которые плохо компонуются (работая с ними, нельзя применять ни #[derive] со сторонним кодом, ни псевдонимы типов). Кроме того, они не обладают информацией о типах (если только не прибегать к обходным манёврам с использованием generic_const_expr , но в таком случае граф вызовов засоряется пространными конструкциями where, а ещё не удаётся работать с обобщёнными параметрами типов). Zig, в свою очередь, позволяет преобразовывать структуры данных самым невероятным образом, и всё это делается лаконично и универсально.

Ниже перейдем к деталям реализации, но сначала я хотел бы объяснить, каков вообще прок от подобного снижения фрагментации памяти.

Контекст

На мой взгляд, разработка эффективных массивов перечислений во многом продиктована опытом работы с компиляторами. Есть проблема, то и дело всплывающая при проектировании абстрактного синтаксического дерева (AST): как бы выделить под это дерево поменьше памяти. Но, когда мы компилируем такие деревья, именно в этот период производительность сильно падает. Дело в том, что пропускная способность памяти и, соответственно, задержки — традиционное узкое место во фронтендах компиляторов. На форумах по языку Rust давно обсуждается видео о компиляторе Carbon от Чендлера Кэррата. В этом материале описано, что распарсенное абстрактное синтаксическое дерево clang может потреблять в 50 раз больше памяти, чем оригинальный исходный код, и такое в порядке вещей!

Хорошо, так как же это всё связано с перечислением? Что ж, чаще всего для представления узлов синтаксического дерева подбирается какая-нибудь рекурсивная (или рекурсивно-подобная) структура данных. Давайте определим узел для выражений в Rust, причём косвенность обеспечим при помощи индексирования новых типов:

enum Expr {
    Unit,
    Number,
    Binary(Operation, ExprId, ExprId),
    Ident(Symbol),
    Eval(ExprId, ExprSlice),
    BlockExpression(ExprId, StatementSlice)
}

Примечание: для сравнения рассмотрим, как можно написать узел AST на языке OCaml:

type expr = 
  | Unit
  | Number
  | Binary of op * expr * expr
  | Ident of symbol
  | Eval of expr * stmt list

Серьёзное отличие от Rust состоит в том, что здесь мы без какой-либо явной косвенности можем выражать по-настоящему рекурсивные типы данных. Дело в том, что здесь нам не приходится вести учёт памяти — за нас это делают среда времени исполнения и сборщик мусора.

Теперь наша задача сводится к тому, что хочется повысить эффективность упаковки этих перечислений. Простой Vec(Expr) будет потреблять память в объёме sizeof(Enum) на каждый элемент, что соответствует размеру наибольшего варианта + метка + отступ. К счастью, есть способы с этим справиться.

Сокращение фрагментации

Рассмотрим простой пример с 3-вариантным перечислением, в которое входят элементы размером 8, 16 и 32 бита. Если хранить их в обычном Vec, то получится такая картина:

Рис. 2: Здесь под каждый элемент резервируется достаточно места, чтобы на каждом участке помещался 32-разрядный вариант, и при этом выполнялись требования к выравниванию.

Самый распространённый способ повысить эффективность упаковки — просто следить, чтобы варианты перечислений были как можно меньше; это делается при помощи размеченных индексов (*).

(*): Если вас интересуют подобные примеры из Rust, взгляните на крейт tagged_index, применяемый в компиляторе, а также ознакомьтесь с этим свежим постом об оптимизации строк. Такие оптимизации постоянно попадаются в высокопроизводительном коде, используемом, например, в языковых средах выполнения, сборщиках мусора, компиляторах, игровых движках или ядрах ОС.

К сожалению, всё это не решает проблему фрагментации окончательно. Зайдём с другой стороны — попробуем поработать напрямую с типом контейнера! Можно было бы воспользоваться подходом «структуры из массивов», чтобы хранить дискриминант и значение в двух разных регионах. На самом деле, именно так и делается в компиляторе Zig.

Рис. 3: Метки и значения объединений хранятся в двух разных регионах, поэтому отступы более можно не учитывать. Правда, здесь в коллекции объединений по-прежнему наблюдается фрагментация вариантов.

Поскольку в Zig применяется поэтапная компиляция, здесь у нас могут быть как раз такие контейнерные типы, которые в обобщённом виде обеспечивают такое SoA-преобразование для любого типа. В Rust приходится удовлетвориться процедурными макросами, например, soa_derive  , а у них есть определённые недостатки. Например, нельзя ставить #[derive] на сторонние типы, не меняя при этом их источника).

Сокращаем фрагментацию вариантов

При таком SoA-преобразовании удаётся в значительной степени избавиться от лишних заполнений, возникающих при работе с метками перечислений, но этот вариант не оптимален. Чтобы как следует избавиться от фрагментации значений, можно создавать по одному вектору на каждый вариант.

Рис. 4: По сравнению с предложенной выше SoA-компоновкой, здесь мы получаем частичный, а не полный порядок. Поэтому, как только произойдёт вставка, мы возвращаемся к размеченному индексу, где содержится как метка перечисления, так и номер в конкретном массиве вариантов.  

Не думаю, что эта коллекция как-то называется, поэтому назову её массив вариантных массивов (или AoVA). Такую структуру можно реализовать и в Rust, и в Zig. В первом для реализации нам понадобятся процедурные макросы, а во втором — выполнение кода во время компиляции (comptime).

Классы эквивалентности размера

На этом можно было бы остановиться, но давайте рассмотрим перечисления, в которых очень много, и эти варианты можно сгруппировать в небольшое множество кластеров, где все типы имеют одинаковый размер:

enum Foo {
    A(u8, u8),
    B(u16),
    C(u16),
    D([u8; 2]),
    E(u32),
    F(u16, u16),
    G(u32),
    H(u32),
    I([u8; 4]),
    J(u32, u32),
    K(u32, (u16, u16)),
    L(u64),
    M(u64),
    N(u32, u16, u16),
    O([u8; 8])
}

Как видите, при подходе «один вектор на вариант» получится в сумме 15 векторов. Вероятно, что количество (повторных) выделений и системных вызовов весьма возрастёт, и для его амортизации потребуется очень много памяти по сравнению с обычным Vec. Кроме того, векторы могут быть произвольным образом распределены в памяти, а из-за этого возрастает вероятность конфликтов в кэше. Коллекция AoVA как таковая также потребляет очень много памяти, и поэтому сильно раздувает любую структуру, в которую встраивается.

Теперь, если сгруппировать все варианты по размеру, то получится три кластера: по 2, 4 и 8 байт. Такие кластеры могут выделяться вместе в рамках одного и того же вектора — тем самым суммарное количество векторов, содержащихся у нас в контейнере, сокращается на 80%. Так, вполне реально сохранить варианты Foo в трёх кластерах:

struct FooVec {
    c_2: Vec<[u8; 2]>, // A - D
    c_4: Vec<[u8; 4]>, // E - I
    c_8: Vec<[u8; 8]>, // J - O
}

Можно было бы сказать, что это плотная версия нашего паттерна AoVA. Но, если мы попытаемся компактно разместить различные варианты в одной и той же выделенной области памяти, то больше не сможем перебирать вектор типобезопасным образом. Добраться до элементов в таком контейнере можно было бы лишь при помощи меченого указателя, который был создан в момент вставки. Если при том паттерне доступа, которым вы пользуетесь, слепая итерация не требуется (так может обстоять дело с уплощёнными древовидными структурами, выстраиваемыми на основе индекса), то такой компромисс, возможно, оправдан.

Я реализовал прототип такой структуры данных на Zig. Здесь наиболее важны встроенные элементы компилятора, обеспечивающие рефлексию при работе с типами полей, байтовыми и битовыми размерами; также они позволяют проверять дискриминант.

Листинг: в сущности, этот код выполняет обычную рефлексию во время компиляции, что позволяет нам вычислять кластеры и отображения полей на кластер. Мы выполняем псевдодинамическое выделение памяти, пользуясь вектором, выделенным в стеке. На основе информации о кластере конструируется структура данных AoVA. Весь исходный код находится здесь.

// определить, что за тип перед нами (напр., структура, объединение, т.д.)
switch (@typeInfo(inner)) {
    .Union => |u| {
        // сохранить отображение с поля объединения на индекс кластера 
        var field_map = [_]u8{0} ** u.fields.len;

        // перебрать поля объединений
        for (u.fields, 0..) |field, idx| {
            // вычислить размер
            const space = @max(field.alignment, @sizeOf(field.type));

            // вставить в хеш-таблицу 
            if (!svec.contains_slow(space)) {
                svec.push(space) catch @compileError(ERR_01);
            }

            field_map[idx] = svec.len - 1;
        }

        // вернуть кластеры
        return .{ .field_map = field_map, .sizes = svec };
    },
    else => @compileError("only unions allowed"),
}

Если вы хотите организовать типобезопасный перебор, то можно добиться этого, раскошелившись на заполнение, и поставить метку обратно:

Рис. 5: В сущности, мы сегментировали перечисление на уровне данных, а интерпретацию на уровне типов оставляем нетронутой

Если отступы занимают слишком много места, то преобразование SoA можно проделать с каждым из вариантных массивов.

Рис. 6: Здесь показан похожий вариант сегментирования, но без отступов. Недостаток в данном случае в том, что мы удваиваем количество векторов.

Вот, как видите, здесь приходится соглашаться на много компромиссов — и все они зависят от того, как именно конкретное перечисление расположено в памяти.

Притом, что создавать такие структуры данных в Zig не составляет труда, в Rust, напротив, воспроизвести эти примеры при помощи процедурных макросов практически невозможно. Дело в том, что процедурные макросы не имеют доступа к информации о типах, например, каковы у них размер и выравнивание. Да, процедурный макрос можно заставить генерировать const fn, вычисляющий кластеры для конкретного перечисления, но при помощи этой функции нельзя указать длину массива для обобщённого типа.

Другое ограничение, присущее дженерикам Rust, таково: из того, как реализован типичный контейнер, не проистекает, является ли данный тип перечислением или структурой. В Zig же мы легко можем сделать нечто подобное:

// это псевдокод
struct EfficientContainer<T> {
    if(T.isEnum()) {
        x: EfficientEnumArray<T>,
    } else {
        x: EfficientStructArray<T>,
    }
}

Также можно конкретизировать, какова разновидность нашей реализации AoVA, основываясь на перечислении. Может быть, польза от совместного размещения различных вариантов начинает становиться целесообразной только на том этапе, когда количество векторов удаётся снизить более чем на 90%.

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

Бонус: как определить битовую ширину индекса во время компиляции

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

Когда этот размеченный индекс включается в последующую структуру данных, например ещё в одно перечисление, информация, естественно, перетекает туда, а те разряды, которые нам не требуются, можно использовать для дискриминанта!

Итак, Zig позволяет добиться эффективности при компоновке памяти. Конкретно указывая, какие биты тебе нужны, можно пользоваться ими в различных частях кода. А при неявном расширении, вызванном приведением целых чисел, работа с API разной битовой ширины остаётся эргономичной. В каком-то отношении эта практика весьма напоминает мне уточнение типов и работу с диапазонами целых чисел (ranged integers), поэтому данная тема сильно перекликается с моим постом о самостоятельном подборе битовой ширины для целых чисел.

Заключение

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