Параметризация

В прошлой главе значение, посчитанное при сборке, было одно на всю программу: comptime SIZE = 10 — и всё. Теперь научимся делать такое значение входом: функция или структура получает его снаружи, и компилятор собирает под каждое своё исполнение.

Два слова, которые нельзя путать

В Mojo эти термины закреплены жёстко, и вся документация ими пользуется:

Термин	Где записан	Когда известен
параметр	в квадратных скобках [...]	на этапе сборки
аргумент	в круглых скобках (...)	во время выполнения

В Python «параметр» и «аргумент» — почти синонимы, и их путают все. Здесь путать нельзя: это два разных механизма.

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

def multiplier[factor: Int](x: Int) -> Int:
    return x * factor


def main():
    print(multiplier[7](6))

Результат

42

factor — параметр, он известен компилятору. x — аргумент, он появится при запуске. Компилятор собрал функцию, в которой умножение на семь зашито намертво.

Параметром не может быть обычная переменная

Самая частая ошибка новичка — попытаться подставить в скобки то, что вычислится позже:

var k = 3
print(multiplier[k](6))   # ошибка

error: invalid call to 'multiplier': cannot use a dynamic value in a parameter list

Лечится либо comptime k = 3, либо переносом factor в обычные аргументы, если значение и правда становится известно только при работе.

Параметр типа и его вывод

Параметром бывает не только число, но и тип. Так пишутся функции, работающие с чем угодно подходящим:

def biggest[T: Comparable & ImplicitlyCopyable](a: T, b: T) -> T:
    return a if a > b else b


def main():
    print(biggest(3, 9))
    print(biggest(String("аз"), String("буки")))

Результат

9
буки

Обратите внимание на вызовы: biggest[Int](3, 9) писать не пришлось. Компилятор вывел T из типов аргументов. Это правило общее: если параметр однозначно определяется аргументами, указывать его не нужно.

Запись Comparable & ImplicitlyCopyable — это список требований к типу. Мы уже встречали такое в главе Структуры: без Comparable компилятор не разрешит сравнение a > b, без ImplicitlyCopyable — возврат значения:

error: value of type 'T' cannot be implicitly copied, it does not conform to 'ImplicitlyCopyable'

Подробный разговор о трейтах — в главе Трейты и обобщённый код.

Иногда требований нужно больше, чем кажется

Функция, которая просто печатает список, выглядит безобидно:

def show[T: Writable](items: List[T]):
    for item in items:
        print(item)

Но компилятор потребует ещё одно свидетельство — уже не про вашу функцию, а про сам List:

error: invalid bindings in signature: lacking evidence to prove correctness
note: constraint declared here needs evidence for 'conforms_to(T, Movable)'
note: add a trailing 'where' clause that requires 'conforms_to(T, Movable)'

Он же подсказывает и решение:

def show[T: Writable](items: List[T]) where conforms_to(T, Movable):

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

Значения по умолчанию и имена

Параметры ведут себя как обычные аргументы Python: у них бывают значения по умолчанию, и передавать их можно по имени.

def repeat[count: Int = 3](text: String):
    comptime for _ in range(count):
        print(text)


def speak[msg: StaticString = "баа"]():
    print(msg)


def main():
    repeat("три раза")
    repeat[1]("один раз")
    speak()
    speak[msg="мее"]()

Результат

три раза
три раза
три раза
один раз
баа
мее

Параметризованные структуры

Тот же механизм работает и для типов. Классический случай — контейнер, размер которого зашит в тип:

struct Buffer[size: Int](Copyable):
    var data: Array[Int, Self.size]

    def __init__(out self):
        self.data = Array[Int, Self.size](fill=0)

Buffer[4] и Buffer[8] — разные типы. Присвоить одно другому нельзя, и это ровно то, чего мы хотели: длина проверяется при сборке, а не при обращении к элементу.

Внутри структуры нужен `Self.`

Ловушка, на которой спотыкаются все. Внутри тела структуры к своему же параметру нельзя обратиться просто по имени:

var data: Array[Int, size]     # ошибка

error: unqualified access to struct parameter 'size'; use 'Self.size' instead

Пишите Self.size — и в объявлении поля, и в методах. Компилятор, к его чести, называет решение прямо в тексте ошибки.

А вот снаружи структуры префикс не нужен: у готового значения параметр читается как обычное поле — print(buffer.size).

Вывод параметра структуры

Функция может принять «любой Vec, какой бы длины он ни был», объявив длину параметром и позволив компилятору её вывести:

def length_squared[n: Int](v: Vec[n]) -> Float64:
    return v.dot(v)

Вызывается это без квадратных скобок — просто length_squared(v). Компилятор достал n из типа аргумента. Рабочий Vec целиком разобран в конце главы.

В скобках можно считать

Параметр — не только имя для константы. В квадратных скобках работает арифметика, и компилятор вычисляет тип результата:

def concat[m: Int](self, other: Vec[m]) -> Vec[Self.n + m]:

Склеили вектор длины 3 с вектором длины 2 — получили Vec[5], и это посчитано при сборке. Никакой проверки длины во время работы не будет: её просто не может не сойтись.

Вот ради чего всё и затевается. Vec[3] и Vec[2] — разные типы, и перепутать их нельзя:

error: invalid call to 'dot': value passed to 'other' cannot be converted from 'Vec[Int(2)]' to 'Vec[Int(3)]'

Что это значит

Скалярное произведение объявлено как dot(self, other: Self), то есть требует вектор той же длины. Вы передали другой.

Как исправить

Ошибка не в вызове, а в замысле: у векторов разной длины скалярного произведения не бывает. Если длины должны совпасть — приведите их к одной. Если операция и правда разнодлинная, объявите второй параметр: def dot[m: Int](self, other: Vec[m]).

В Python такая ошибка вылезла бы при запуске — в лучшем случае IndexError, в худшем молча посчиталось бы не то.

Как читать чужие сигнатуры

Дальше — четыре записи, которые вы почти наверняка встретите в чужом коде и в стандартной библиотеке раньше, чем захотите написать сами. Держите их как справку.

Одно имя для типа, две записи. Обе делают одно и то же — дают короткое имя частично заданному типу, — но вторая гибче: она умеет считать.

comptime StringKeyDict = Dict[String, _]      # прочерк: параметр оставлен на потом
comptime TwoOfAKind[dt: DType] = SIMD[dt, 2]  # параметр свой, тип вычисляется

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

Выводимые параметры: //. Всё, что объявлено до двух косых черт, компилятор обязан вывести сам — указать это при вызове нельзя:

def first[T: Copyable, //](items: List[T]) -> T:
    return items[0].copy()


def main():
    print(first([10, 20, 30]))
    print(first([String("аз"), String("буки")]))

Результат

10
аз

Смысл // — отделить «служебные» параметры, которые нужны сигнатуре, но не должны попадаться на глаза вызывающему.

Автопараметризация: многоточие. Если конкретные параметры вам безразличны, их можно не объявлять вовсе:

def take_simd(vec: SIMD[...]):
    print(vec.dtype, len(vec))

Компилятор сам заведёт скрытые параметры и выведет их из аргумента. Вызов take_simd(SIMD[DType.float32, 4](1.5)) напечатает float32 4.

Вариативные параметры. Список, длина которого известна при сборке:

def dims[*sizes: Int]():
    comptime for s in sizes:
        print(s)


def main():
    dims[2, 3, 5]()

Результат

2
3
5

Вариативные параметры — только простые типы

Элементы должны быть однородными и помещаться в регистр. Int годится, String — нет:

def dims[*sizes: String]():   # не соберётся

error: function instantiation failed
note: call expansion failed with parameter value(s): (...)

Чем за это платят

Обычно пишут коротко: «специализация раздувает бинарник». Мы проверили — и оказалось, что это верно далеко не всегда.

Постановка одна на оба замера. Функция считает хеш в цикле; её вызывают шестьдесят раз с разной длиной цикла. В параметрической версии длина — параметр, и цикл разворачивается при сборке; в обычной — аргумент, и цикл остаётся циклом. Различаются замеры только тем, что цикл делает внутри.

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

Вариант	Размер .text
обычный аргумент	4 911 байт
параметр	4 911 байт

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

Случай второй: тело читает массив. Ровно тот же цикл, но числа берутся из списка, который заполняется при запуске.

Вариант	Размер .text
обычный аргумент	7 231 байт
параметр	143 103 байта

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

Вывод не «параметры дорогие» и не «параметры бесплатные», а такой:

Специализация обменивает код на константы. Пока есть что сворачивать — обмен выгоден или хотя бы безвреден. Как только развёрнутый цикл начинает работать с данными, которых компилятор не видит, платить приходится размером, и платить много.

Как измерить у себя

Размер файла на диске поглощает мелкие различия: в первом замере .text совпал полностью, а файлы отличались на 8 байт — это шум выравнивания. Большую разницу файл, разумеется, покажет (во втором замере 26 992 против 172 488 байт), но мерить надёжнее секцию с кодом:

readelf -S ./program | grep -A1 '\.text'   # весь код
nm -S ./program | grep ' T main'            # одна функция

Во втором замере main в параметрической сборке занял всего 83 байта: весь объём разошёлся по шестидесяти специализациям.

Две другие статьи расхода никуда не деваются: время сборки (компилятор выполняет ваш код) и читаемость ошибок (диагностика приходит из сгенерированного кода, а не из написанного).

Когда параметр оправдан

Значение действительно известно заранее и от него зависит структура кода: размер буфера, тип элемента, ширина SIMD-вектора, включённость отладочных проверок.

Во всех остальных случаях берите обычный аргумент. Параметр там, где хватило бы аргумента, не ускоряет ничего, а мешает многому.

Всё вместе

Вектор фиксированной длины: длина живёт в типе, значения — в памяти.

vectors.mojo

# Вектор фиксированной длины. Длина — параметр: она известна компилятору,
# поэтому он проверяет её сам и не даёт сложить несовместимое.


struct Vec[n: Int](Copyable, Movable):
    var data: Array[Float64, Self.n]

    def __init__(out self, fill: Float64 = 0.0):
        self.data = Array[Float64, Self.n](fill=fill)

    def copy(self) -> Self:
        """Нужен из-за Copyable: Array не копируется неявно."""
        var result = Self()
        comptime for i in range(Self.n):
            result.data[i] = self.data[i]
        return result^

    def __getitem__(self, index: Int) -> Float64:
        return self.data[index]

    def __setitem__(mut self, index: Int, value: Float64):
        self.data[index] = value

    def dot(self, other: Self) -> Float64:
        """Скалярное произведение. Цикл разворачивается при сборке."""
        var total = 0.0
        comptime for i in range(Self.n):
            total += self.data[i] * other.data[i]
        return total

    def concat[m: Int](self, other: Vec[m]) -> Vec[Self.n + m]:
        """Длина результата посчитана компилятором: Self.n + m."""
        var result = Vec[Self.n + m]()
        comptime for i in range(Self.n):
            result.data[i] = self.data[i]
        comptime for j in range(m):
            result.data[Self.n + j] = other.data[j]
        return result^


def length_squared[n: Int](v: Vec[n]) -> Float64:
    """Параметр n выводится из аргумента — писать его при вызове не нужно."""
    return v.dot(v)


def main():
    var a = Vec[3]()
    a[0] = 1.0
    a[1] = 2.0
    a[2] = 2.0

    print("длина известна типу:", a.n)
    print("квадрат длины:", length_squared(a))

    var ones = Vec[3](fill=1.0)
    print("скалярное произведение:", a.dot(ones))

    var copied = a.copy()
    copied[0] = 100.0
    print("копия независима:", a[0], "и", copied[0])

    # Длина результата вычислена при сборке: 3 + 2 = 5.
    var short = Vec[2](fill=5.0)
    var joined = a.concat(short)
    print("склейка длиной", joined.n, ":", joined[0], joined[3], joined[4])

    # А вот ради чего всё затевалось. Раскомментируйте — и программа
    # не соберётся, потому что Vec[3] и Vec[2] это разные типы:
    #
    #     print(a.dot(short))
    #
    # error: invalid call to 'dot': value passed to 'other' cannot be
    # converted from 'Vec[Int(2)]' to 'Vec[Int(3)]'
    print("несовпадение длин поймано бы при сборке, а не при запуске")

Результат

длина известна типу: 3
квадрат длины: 9.0
скалярное произведение: 5.0
копия независима: 1.0 и 100.0
склейка длиной 5 : 1.0 5.0 5.0
несовпадение длин поймано бы при сборке, а не при запуске

Здесь собрано всё из главы: параметр структуры с Self.n, вывод параметра в length_squared, параметрическое выражение Self.n + m в concat и comptime for, разворачивающий короткие циклы прямо в арифметику. Последние строки — самое важное: закомментированный вызов, который не даст программе собраться.

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

Длину буфера программа читает из аргумента командной строки. Параметром её объявить или аргументом?

Только аргументом. Параметр обязан быть известен компилятору, а содержимое командной строки появляется при запуске — на попытку подставить его в квадратные скобки компилятор ответит cannot use a dynamic value in a parameter list. Значит, и контейнер нужен с длиной времени выполнения — List, а не Array.

`a` имеет тип `Vec[3]`, `short` — `Vec[2]`. Что скажет компилятор на `a.dot(short)`?

Откажется собирать: value passed to 'other' cannot be converted from 'Vec[Int(2)]' to 'Vec[Int(3)]'. Длина входит в тип, поэтому несовпадение — обычная ошибка типов, пойманная при сборке. В Python то же самое дало бы IndexError при запуске или, хуже, молча посчитало бы не то.

Почему `biggest(3, 9)` работает без `biggest[Int](3, 9)`?

Компилятор выводит параметр типа из типов аргументов. Правило общее: если параметр однозначно определяется аргументами, писать его при вызове не нужно.

Верно ли, что параметризация всегда увеличивает размер программы?

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

Что дальше

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