Генерация типаж-объектов на лету (или безумие с Rust)

В этой статье мы немного потешимся с языком программирования Rust, а в частности, с типаж-объектами.

Когда я знакомился с Rust, одна из деталей реализации типаж-объектов мне показалась интересной. А именно то, что виртуальная таблица функций находится не в самих данных, а в "толстом" указателе на них. Каждый указатель на типаж-объект ) содержит указатель на сами данные, а так же ссылку на виртуальную таблицу, где будут находиться адреса функций, реализующий данный типаж-объект для заданной структуры (но так как это деталь реализации, то поведение может поменяться.

Начнём с простого примера, демонстрирующего "толстые" указатели. Следующий код выведет на 64-разрядной архитектуре 8 и 16:

fn main () {
    let v: &String = &"hello".into();
    let disp: &std::fmt::Display = v;
    println!("Указатель на строку: {}", std::mem::size_of_val(&v));
    println!("Толстый указатель на типаж-объект: {}", std::mem::size_of_val(&disp));
}

Чем это интересно? Когда я занимался энтерпрайз-Java, одна из задач, которая довольно-таки регулярно возникала была адаптация существующих объектов под заданные интерфейсы. То есть объект вот уже есть, выдан в виде ссылки, а надо его адаптировать к заданному интерфейсу. И менять входной объект нельзя, он есть такой, какой есть.

Приходилось делать как-то так:

Person adapt(Json value) {
    // ...какая-нибудь логика, например, проверим, что "value" действительно
    // соответствует контракту Person
    return new PersonJsonAdapter(value);
}

У такого подхода вылазили разные проблемы. Например, если один и тот же объект "адаптируется" два раза, то получим два разных Person (с точки зрения сравнения ссылок). Да и сам факт, что приходится создавать новые объекты каждый раз как-то некрасиво.

Когда я увидел типаж-объекты в Rust, у меня возникла мысль, что в Rust это можно сделать гораздо более элегантно! Можно же взять и приписать данным другую виртуальную таблицу и получить новый типаж-объект! И не надо выделять память на каждый экземпляр. При этом, вся логика "заимствования" остаётся на месте — наша функция адаптирования будет иметь вид что-то вроде fn adapt<'a>(value: &'a Json) -> &'a Person (то есть мы как бы заимствуем из исходных данных).

Даже более того, можно "заставить" один и тот же тип (например, String) реализовать наш типаж-объект несколько раз, с разным поведением. Зачем? Да мало ли чего может понадобиться в энтерпрайзе?!

Давайте попробуем это реализовать.

Постановка задачи

Поставим задачу таким образом: сделать функцию annotate, которая "припишет" обычному типу String следующий типаж-объект:

trait Object {
  fn type_name(&self) -> &str;
  fn as_string(&self) -> &String;
}

И сама функция annotate:

/// Адаптировать строку к типажу-объекту `Object`, который будет декларировать,
/// что его "тип" -- тот, который задан в `type_name`.
fn annotate<'a>(input: &'a String, type_name: &str) -> &'a dyn Object {
    // ...
}

Напишем сразу тест. Во-первых, убедимся, что "приписанный" тип совпадает с ожидаемым. Во вторых, убедимся, что мы можем достать исходную строку и это будет всё та же строка (с точки зрения указателей):

#[test]
fn test() {
    let input: String = "hello".into();
    let annotated1 = annotate(&input, "Widget");
    let annotated2 = annotate(&input, "Gadget");

    // Типаж-объект ведёт себя так, как мы ожидаем 
    assert_eq!("Widget", annotated1.type_name());
    assert_eq!("Gadget", annotated2.type_name());

    let unwrapped1 = annotated1.as_string();
    let unwrapped2 = annotated2.as_string();

    // Это физически всё та же строка -- сравниваем указатели
    assert_eq!(unwrapped1 as *const String, &input as *const String);
    assert_eq!(unwrapped2 as *const String, &input as *const String);
}

Подход №1: а после нас хоть потоп!

Для начала попробуем сделать совсем наивную реализацию. Просто завернём наши данные в "обёртку", которая будет дополнительно содержать type_name:

struct Wrapper<'a> {
    value: &'a String,
    type_name: String,
}

impl<'a> Object for Wrapper<'a> {
    fn type_name(&self) -> &str {
        &self.type_name
    }

    fn as_string(&self) -> &String {
        self.value
    }
}

Пока ничего особенного. Всё как в Java. Но у нас же нет сборщика мусора, где мы хранить-то будет эту обёртку? Нам же надо ссылку вернуть, да так, чтобы она осталась действительной после вызова функции annotate. Ничего, страшного засунем в коробку (Box), чтобы обёртка (Wrapper) была выделена на куче. А потом вернём на неё ссылку. А чтобы обёртка осталась жить после вызова функции annotate, мы эту коробку "утечём":

fn annotate<'a>(input: &'a String, type_name: &str) -> &'a dyn Object {
    let b = Box::new(Wrapper {
        value: input,
        type_name: type_name.into(),
    });
    Box::leak(b)
}

… и тест проходит!

Но это какое-то сомнительное решение. Мало того, что мы всё так же выделяем память при каждом "аннотировании", так ещё и память утекает (Box::leak возвращает ссылку на данные, сохранённые на куче, но при этом "забывает" саму коробку, то есть автоматического освобождения не произойдет).

Подход №2: арена!

Для начала попробуем куда-нибудь сохранить эти обёртки так, чтобы они всё-таки высвобождались в какой-то момент. Но при этом сохранив сигнатуру annotate как она есть. То есть вернуть ссылку с подсчётом ссылок (например, Rc<Wrapper>) — не подходит.

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

Как-то так. Используется библиотека typed-arena для хранения обёрток, но можно было обойтись и типом Vec<Box<Wrapper>>, главное, гарантировать, что Wrapper никуда не перемещается (в ночном Rust можно для этого было использорвать pin API):

struct TypeSystem {
    wrappers: typed_arena::Arena<Wrapper>,
}

impl TypeSystem {
    pub fn new() -> Self {
        Self {
            wrappers: typed_arena::Arena::new(),
        }
    }

    /// Результат заимствует из параметра `input`, и при этом должен жить меньше,
    /// чем система типов (иначе возможна ситуация, когда все обёртки высвободятся,
    /// а при этом на них ещё будут ссылки)!
    pub fn annotate<'a: 'b, 'b>(
        &'a self,
        input: &'b String,
        type_name: &str
    ) -> &'b dyn Object {
        self.wrappers.alloc(Wrapper {
            value: input,
            type_name: type_name.into(),
        })
    }
}

Но куда же делся параметр, отвечающий за время жизни ссылки у типа Wrapper? От него пришлось избавиться, так как мы не можем приписать какое-то фиксированное время жизни в типе typed_arena::Arena<Wrapper<'?>>. Каждая обёртка имеет уникальный параметр, зависящий от input!

Вместо этого, мы присыпем немного небезопасного Rust, чтобы избавиться от параметра-времени жизни:

struct Wrapper {
    value: *const String,
    type_name: String,
}

impl Object for Wrapper {
    fn type_name(&self) -> &str {
        &self.type_name
    }

    /// Эта конверсия -- безопасна, так как мы гарантируем (через сигнатуру
    /// `annotate`), что ссылка на обёртку (как часть ссылки на типаж-объект
    /// `&Object`) живет меньше, чем ссылка на сами данные (`String`).
    fn as_string(&self) -> &String {
        unsafe { &*self.value }
    }
}

И тесты снова проходят, тем самым давая нам уверенность в правильности решения. В дополнение к ощущению лёгкой неловкости из-за unsafe (как и должно быть, с небезопасным Rust лучше не шутить!).

Но всё же, а как же обещанный вариант, не требующий дополнительных выделений памяти на обёртки?

Подход №3: да разверзнутся врата ада

Идея. Для каждого уникального "типа" ("Widget", "Gadget"), мы создадим виртуальную таблицу. Руками, во время выполнения программы. И припишем её к переданной нам ссылке на сами данные (которые у нас, как мы помним, просто String).

Для начала небольшое описание, того что же нам нужно получить. Итак, ссылка на типаж объект, как она устроена? По сути, это просто два указателя, один на сами данные, а другой — на виртуальную таблицу. Так и запишем:

#[repr(C)]
struct TraitObject {
    pub data: *const (),
    pub vtable: *const (),
}

(#[repr(C)] нам нужен, чтобы гарантировать правильное расположение в памяти).

Вроде всё просто, мы сгенерируем новую таблицу для заданных параметров и "соберём" ссылку на типаж-объект! Но из чего же состоит эта таблица?

Правильный ответ на этот вопрос будет "это деталь реализации". Но мы сделаем так; создадим файл rust-toolchain в корне нашего проекта и запишем туда: nightly-2018-12-01. Ведь зафиксировання сборка может считаться стабильной, правда?

Теперь, когда мы зафиксировали версию Rust (на самом деле, ночная сборка нам понадобится для одной из библиотек чуть ниже).

После некоторого поиска в интернете выясняем, что формат таблицы простой: сначала идёт ссылка на деструктор, затем два поля связанных с выделением памяти (размер типа и выравнивание), а потом идут функции, одна за другой (порядок — на усмотрение компилятора, но у нас всего две функции, поэтому вероятность угадать довольно велика, 50%).

Так и запишем:

#[repr(C)]
#[derive(Clone, Copy)]
struct VirtualTableHeader {
    destructor_fn: fn(*mut ()),
    size: usize,
    align: usize,
}

#[repr(C)]
struct ObjectVirtualTable {
    header: VirtualTableHeader,
    type_name_fn: fn(*const String) -> *const str,
    as_string_fn: fn(*const String) -> *const String,
}

Аналогично, #[repr(C)] нужен, чтобы гарантировать правильное расположение в памяти. Я разделил на две структуры, чуть позже это нам пригодится.

Попробуем теперь написать нашу систему типов, которая будет предоставлять функцию annotate. Нам нужно будет кешировать сгенерированные таблицы, поэтому заведём кеш:

struct TypeInfo {
    vtable: ObjectVirtualTable,
}

#[derive(Default)]
struct TypeSystem {
    infos: RefCell<HashMap<String, TypeInfo>>,
}

Мы используем внутреннее состояние RefCell чтобы наша функция TypeSystem::annotate могла получать &self как разделяемую ссылку. Это важно, так как мы "заимствуем" у TypeSystem, чтобы гарантировать, что сгенерированные нами виртуальные таблицы жили дольше, чем ссылка на типаж-объект, который мы возвращаем из annotate.

Так как мы хотим, чтобы можно было аннотировать много экземляров, мы не можем заимствовать &mut self, как изменяемую ссылку.

И набросаем вот такой код:

impl TypeSystem {
    pub fn annotate<'a: 'b, 'b>(
        &'a self,
        input: &'b String,
        type_name: &str
    ) -> &'b dyn Object {
        let type_name = type_name.to_string();
        let mut infos = self.infos.borrow_mut();
        let imp = infos.entry(type_name).or_insert_with(|| unsafe {
            // Откуда мы её возьмём, эту таблицу?
            let vtable = unimplemented!();
            TypeInfo { vtable }
        });

        let object_obj = TraitObject {
            data: input as *const String as *const (),
            vtable: &imp.vtable as *const ObjectVirtualTable as *const (),
        };

        // Сконвертируем сконструированную структуру в ссылку на типаж-объект
        unsafe { std::mem::transmute::<TraitObject, &dyn Object>(object_obj) }
    }
}

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

trait Whatever {}
impl<T> Whatever for T {}

Он нам пригодится, чтобы получить эту самую "любую другую виртуальную таблицу". А потом, скопируем у него эти три записи:

let whatever = input as &dyn Whatever;
let whatever_obj = std::mem::transmute::<&dyn Whatever, TraitObject>(whatever);
let whatever_vtable_header = whatever_obj.vtable as *const VirtualTableHeader;
let vtable = ObjectVirtualTable {
    // Скопируем записи!
    header: *whatever_vtable_header,
    type_name_fn: unimplemented!(),
    as_string_fn: unimplemented!(),
};

TypeInfo { vtable }

В принципе, размер и выравнивание мы могли бы получить через std::mem::size_of::<String>() и std::mem::align_of::<String>(). Но вот откуда еще можно "украсть" деструктор, я не знаю.

Хорошо, но где же мы возьмём адреса этих функций, type_name_fn и as_string_fn? Можно заметить, что as_string_fn в общем-то и не нужна, указатель на данные-то всегда идёт первой записью в представлении типаж-объекта. То есть это функция всегда одна и та же:

impl Object for String {
    // ...

    fn as_string(&self) -> String {
        self
    }
}

Но вот со второй функцией уже не так просто! Она же зависит от нашего имени "типа", type_name.

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

Для простоты предположим, что нас интересует только Mac OS и Linux (после всех этих веселых трансформаций, совместимость нас уже не особо волнует, правильно?). А, да, исключительно x86-64, конечно.

Вторую функцию, as_string, реализовать легко. Нам обещают, что первый параметр будет в регистре RDI. А вернуть значение в RAX. То есть код функции будет что-то вроде:

dynasm!(ops
    ; mov rax, rdi
    ; ret
);

А вот первая функция немного хитрее. Во-первых, нам надо вернуть &str, а это толстый указатель. Его первая часть — указатель на строку, а вторая часть — длина строкового среза. К счастью, соглашение выше позволяет возвращать и 128-разрядные результаты, используя регистр EDX для второй части.

Осталось где-то добыть ссылку на строковый срез, который содержит нашу строку type_name. Полагаться на type_name мы не хотим (хотя через аннотации времени жизни можно гарантировать, что type_name будет жить дольше, чем возвращённое значение).

Но у нас есть копия этой строки, которую мы помещаем в хеш-таблицу. Скрестив пальцы, мы сделаем предположение что расположение строкового слайса который на вернёт String::as_str не поменяется от перемещения самой строки String (а перемещаться String будет в процессе смены размера HashMap, где эта строка хранится ключём). Не знаю, гарантирует ли стандартная библиотека такое поведение, но мы ж так, поиграть просто?

Получаем необходимые компоненты:

let type_name_ptr = type_name.as_str().as_ptr();
let type_name_len = type_name.as_str().len();

И пишем такую функцию:

dynasm!(ops
    ; mov rax, QWORD type_name_ptr as i64
    ; mov rdx, QWORD type_name_len as i64
    ; ret
);

И, наконец, финальный код annotate:

pub fn annotate<'a: 'b, 'b>(&'a self, input: &'b String, type_name: &str) -> &'b Object {
    let type_name = type_name.to_string();

    // Запоминаем расположение и длину строкового слайса
    let type_name_ptr = type_name.as_str().as_ptr();
    let type_name_len = type_name.as_str().len();
    let mut infos = self.infos.borrow_mut();
    let imp = infos.entry(type_name).or_insert_with(|| unsafe {
        let mut ops = dynasmrt::x64::Assembler::new().unwrap();

        // Создаём код для функции `type_name`
        let type_name_offset = ops.offset();
        dynasm!(ops
            ; mov rax, QWORD type_name_ptr as i64
            ; mov rdx, QWORD type_name_len as i64
            ; ret
        );

        // Создаём код для функции `as_string`
        let as_string_offset = ops.offset();
        dynasm!(ops
            ; mov rax, rdi
            ; ret
        );
        let buffer = ops.finalize().unwrap();

        // Копируем части из аналогичной таблицы
        let whatever = input as &dyn Whatever;
        let whatever_obj =
            std::mem::transmute::<&dyn Whatever, TraitObject>(whatever);
        let whatever_vtable_header =
            whatever_obj.vtable as *const VirtualTableHeader;
        let vtable = ObjectVirtualTable {
            header: *whatever_vtable_header,
            type_name_fn: std::mem::transmute(buffer.ptr(type_name_offset)),
            as_string_fn: std::mem::transmute(buffer.ptr(as_string_offset)),
        };

        TypeInfo { vtable, buffer }
    });

    assert_eq!(imp.vtable.header.size, std::mem::size_of::<String>());
    assert_eq!(imp.vtable.header.align, std::mem::align_of::<String>());

    let object_obj = TraitObject {
        data: input as *const String as *const (),
        vtable: &imp.vtable as *const ObjectVirtualTable as *const (),
    };
    unsafe { std::mem::transmute::<TraitObject, &dyn Object>(object_obj) }
}

Для целей dynasm нужно ещё добавить поле buffer в нашу структуру TypeInfo. Это поле управляет памятью, которая хранит код наших сгенерированных функций:

#[allow(unused)]
buffer: dynasmrt::ExecutableBuffer,

И все тесты проходят!

Готово, мастер!

Вот так легко и непринуждённо можно генерировать свои реализации типаж-объектов в Rust коде!

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

Есть, правда, (ещё) одна особенность, на которую я тут полагаюсь. А именно то, что безопасно высвобождать память, занятую виртуально таблицей после того, как нет ссылок на типаж-объект, её использующий. С одной стороны, это логично, использовать виртуальную таблицу можно только через ссылки типаж-объектов. С другой стороны, таблицы, предоставляемые Rust имеют время жизни 'static. Вполне можно предположить какой-нибудь код, который отделит таблицу от ссылки для каких-то своих целей (мало ли, например, для каких-то своих грязных делишек).

Исходный код можно найти тут.

Автор: Иван Дубров

Источник