Rust — сохраняем безразмерные типы в статической памяти

23648fc101d779dab6536075fc1741cb.jpg?v=1

Не так давно в качестве хобби я решил погрузиться в изучение embedded разработки на Rust и через какое-то время мне захотелось сделать себе логгер, который бы просто писал логи через UART, но при этом не знал какая конкретно реализация используется. Вот тут я быстро осознал, именно в этом конкретном случае я не могу полагаться на статический полиморфизм и мономорфизацию, ведь компилятор не знает сколько нужно памяти выделять под конкретную реализацию. Фактически это означает, что нам нужно как-то уметь сохранять в памяти типы, размер которых неизвестен на этапе компиляции. Такой способностью обладает тип Box, но он использует для этого динамическое выделение памяти из кучи. В итоге возникла идея написать свой аналог данного типа, но хранящий объект в предоставленном пользователем буфере, а не в глобальной куче.

А почему бы просто не взять какой-нибудь linked_list_allocator от Фила, дать ему пару килобайт памяти и воспользоваться обычным Box типом, или даже взять какой-нибудь простейший bump аллокатор, ведь мы хотим использовать его лишь для того, чтобы создать несколько глобальных объектов, но есть множество сценариев, когда куча не используется принципиально? Это и дополнительная зависимость от целого alloc крейта и дополнительные риски, что использование кучи выйдет за рамки строго детерминированных сценариев, что будет приводить к трудноуловимым ошибкам.

С другой стороны, мы можем просто принимать &'static dyn Trait и таким образом переложить заботу о том, как получить такую ссылку, на конечного пользователя, но чтобы обеспечить потом доступ к этой ссылке, нам необходимо использовать примитивы синхронизации или же воспользоваться unsafe кодом, с другой стороны, конечный пользователь тоже должен воспользоваться ими, чтобы создать такую ссылку. В конечном итоге у нас получается или двойная работа или unsafe в публичном API, что довольно плохо. Да и в целом, Box обладает гораздо более широкой областью применения, например, его можно использовать для организации очереди задач в очередном futures executor.

Что же такое безразмерные типы?

Хочу немного напомнить о том, что такое типы динамического размера (DST) и чем они отличаются от обычных. Большинство типов имеют фиксированный размер, известный на этапе компиляции, такие типы компилятор помечает трейтом Sized, но есть типы, размер которых в памяти неизвестен на этапе компиляции, поэтому компилятор не может размещать их в статической памяти или на стэке по значению, такие типы еще называют безразмерными.

Тут очень важно понимать, что для каждого конкретного типа компилятор точно знает внутреннее устройство его безразмерного аналога, проблема в том, что один и тот же безразмерный dyn Display может быть получен из самых разнообразных конкретных типов, чем и обеспечивается динамический полиморфизм. И именно поэтому можно приводить к безразмерным типам лишь ссылки и указатели, уж размер указателя компилятору всегда известен.

Ссылки и указатели в Rust это не всегда просто адреса в памяти, в случае с DST типами, помимо адреса хранится еще и объект с метаданными указателя, но гораздо проще это все осознать, если просто взглянуть на код стандартной библиотеки.

#[repr(C)]
pub(crate) union PtrRepr {
    pub(crate) const_ptr: *const T,
    pub(crate) mut_ptr: *mut T,
    pub(crate) components: PtrComponents,
}

#[repr(C)]
pub(crate) struct PtrComponents {
    pub(crate) data_address: *const (),
    pub(crate) metadata: ::Metadata,
}

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

pub struct DynMetadata {
    vtable_ptr: &'static VTable,
    phantom: crate::marker::PhantomData,
}
/// The common prefix of all vtables. It is followed by function pointers for trait methods.
///
/// Private implementation detail of DynMetadata::size_of etc.
#[repr(C)]
struct VTable {
    drop_in_place: fn(*mut ()),
    size_of: usize,
    align_of: usize,
}

Таким образом, в текущей реализации, размер &dyn Display на x86_64 составляет 16 байт, а когда мы пишем такой вот код:

let a: u64 = 42;
let dyn_a: &dyn Display = &a;

Компилятор генерирует объект VTable и сохраняет его где-то в статической памяти, а обычную ссылку заменяет на широкую, содержащую кроме адреса еще и указатель на таблицу виртуальных функций. Ссылка на таблицу виртуальных функций статическая и не зависит от места расположения значения, таким образом, для того, чтобы создать желаемый Box из искомого значения a, нам необходимо извлечь метаданные из ссылки на dyn_a и все это вместе скопировать в заранее приготовленный для этого буфер. Чтобы все это сделать, нам необходимо использовать nightly features:  unsize и ptr_metadata.

Для получения &dyn T из &Value используется специальный маркерный трейт Unsize, который выражает отношение между Sized типом и его безразмерным альтер-эго. То есть,  T это Unsize в том случае, если T реализует Trait.

А чтобы работать с метаданными указателя используется функция core::ptr::metadata и типаж Pointee, который связывает тип указателя и тип его метаданных, в случае с безразмерными типами метаданные имеют тип DynMetadata, где T это искомый безразмерный тип.

#[inline]
fn meta_offset_layout(value: &Value) -> (DynMetadata, Layout, usize)
where
    T: ?Sized + Pointee>,
    Value: Unsize + ?Sized,
{
    // Get dynamic metadata for the given value.
    let meta = ptr::metadata(value as &T);
    // Compute memory layout to store the value + its metadata.
    let meta_layout = Layout::for_value(&meta);
    let value_layout = Layout::for_value(value);
    let (layout, offset) = meta_layout.extend(value_layout).unwrap();
    (meta, layout, offset)
}

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

Обратите внимание, что мы берем Layout от ссылки на метаданные, а не DynMetadata::layout, последний описывает размещение VTable, но нас интересует размещение самого DynMetadata, будьте внимательны!

Пишем свой Box

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

Конструктор, который копирует данные и метаданные в предоставленный буфер, используя довольно удобное API указателя.

impl Box
where
    T: ?Sized + Pointee>,
    M: AsRef<[u8]> + AsMut<[u8]>,
{
    pub fn new_in_buf(mut mem: M, value: Value) -> Self
    where
        Value: Unsize,
    {
        let (meta, layout, offset) = meta_offset_layout(&value);
        // Check that the provided buffer has sufficient capacity to store the given value.
        assert!(layout.size() > 0);
        assert!(layout.size() <= mem.as_ref().len());

        unsafe {
            let ptr = NonNull::new(mem.as_mut().as_mut_ptr()).unwrap();
            // Store dynamic metadata at the beginning of the given memory buffer.
            ptr.cast::>().as_ptr().write(meta);
            // Store the value in the remainder of the memory buffer.
            ptr.cast::()
                .as_ptr()
                .add(offset)
                .cast::()
                .write(value);

            Self {
                mem,
                phantom: PhantomData,
            }
        }
    }
}

А вот и код, который собирает байты назад в &dyn Trait:

    #[inline]
    fn meta(&self) -> DynMetadata {
        unsafe { *self.mem.as_ref().as_ptr().cast() }
    }

    #[inline]
    fn layout_meta(&self) -> (Layout, usize, DynMetadata) {
        let meta = self.meta();
        let (layout, offset) = Layout::for_value(&meta).extend(meta.layout()).unwrap();
        (layout, offset, meta)
    }

    #[inline]
    fn value_ptr(&self) -> *const T {
        let (_, offset, meta) = self.layout_meta();
        unsafe {
            let ptr = self.mem.as_ref().as_ptr().add(offset).cast::<()>();
            ptr::from_raw_parts(ptr, meta)
        }
    }

    #[inline]
    fn value_mut_ptr(&mut self) -> *mut T {
        let (_, offset, meta) = self.layout_meta();
        unsafe {
            let ptr = self.mem.as_mut().as_mut_ptr().add(offset).cast::<()>();
            ptr::from_raw_parts_mut(ptr, meta)
        }
    }

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

impl Deref for Box
where
    T: ?Sized + Pointee>,
    M: AsRef<[u8]> + AsMut<[u8]>,
{
    type Target = T;

    #[inline]
    fn deref(&self) -> &T {
        self.as_ref()
    }
}

impl DerefMut for Box
where
    T: ?Sized + Pointee>,
    M: AsRef<[u8]> + AsMut<[u8]>,
{
    #[inline]
    fn deref_mut(&mut self) -> &mut T {
        self.as_mut()
    }
}
running 8 tests
test tests::test_box_dyn_fn ... ok
test tests::test_box_nested_dyn_fn ... ok
test tests::test_box_in_provided_memory ... ok
test tests::test_box_trait_object ... ok
test tests::test_box_move ... ok
test tests::test_drop ... ok
test tests::test_layout_of_dyn ... ok
test tests::test_box_insufficient_memory ... ok

Miri

Казалось бы, все замечательно, можно использовать библиотеку в боевом коде… Но, постойте, мы же написали unsafe код, как мы вообще можем быть уверены в том, что нигде не нарушили никакие инварианты? К счастью, существует такой проект, как Miri, который интерпретирует промежуточное представление MIR, генерируемое компилятором rustc, используя специальную виртуальную машину. Таким образом, можно находить очень многие ошибки в unsafe коде, подробнее об этом можно почитать в этой статье. Давайте попробуем запустить наши тесты используя Miri.

cargo miri test
   Compiling static-box v0.0.1 (/home/aleksey/Projects/opensource/static-box)
    Finished test [unoptimized + debuginfo] target(s) in 0.40s
     Running unittests (target/x86_64-unknown-linux-gnu/debug/deps/static_box-e2c02215f3157959)
running 8 tests
test tests::test_box_dyn_fn ... error: Undefined Behavior: accessing memory with alignment 1, but alignment 8 is required
   --> /home/aleksey/.rustup/toolchains/nightly-2021-04-25-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/ptr/mod.rs:886:9
    |
886 |         copy_nonoverlapping(&src as *const T, dst, 1);
    |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ accessing memory with alignment 1, but alignment 8 is required
    |
    = help: this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior
    = help: see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information

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

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

        // Construct a box to move the specified memory into the necessary location.
        // SAFETY: This code relies on the fact that this method will be inlined.
        let mut new_box = Self {
            align_offset: 0,
            mem,
            phantom: PhantomData,
        };

        let raw_ptr = new_box.mem.as_mut().as_mut_ptr();
        // Compute the offset that needs to be applied to the pointer in order to make
        // it aligned correctly.
        new_box.align_offset = raw_ptr.align_offset(layout.align());

Вот собственно и все, после этого Miri больше не показывает ошибок выравнивания.

cargo miri test
   Compiling static-box v0.1.0 (/home/aleksey/Projects/opensource/static-box)
    Finished test [unoptimized + debuginfo] target(s) in 0.30s
     Running unittests (target/x86_64-unknown-linux-gnu/debug/deps/static_box-ce23f69c165cf930)

running 11 tests
test tests::test_box_dyn_fn ... ok
test tests::test_box_in_provided_memory ... ok
test tests::test_box_in_static_mem ... ok
test tests::test_box_in_unaligned_memory ... ok
test tests::test_box_insufficient_memory ... ok
test tests::test_box_move ... ok
test tests::test_box_nested_dyn_fn ... ok
test tests::test_box_trait_object ... ok
test tests::test_drop ... ok
test tests::test_layout_of_dyn_split_at_mut ... ok
test tests::test_layout_of_dyn_vec ... ok
test result: ok. 11 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out
   Doc-tests static-box
running 2 tests
test src/lib.rs - (line 24) ... ok
test src/lib.rs - (line 48) ... ok
test result: ok. 2 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.15s

Хочу еще сказать несколько слов относительно типа Layout, в нем содержится два поля size, которое содержит размер памяти в байтах, необходимый для размещения объекта, и align — это число (причем всегда степень двойки), которому должен быть кратен указатель на объект данного типа. И таким образом, чтобы починить выравнивание, мы просто вычисляем сколько нам нужно прибавить к адресу начала буфера, чтобы получить адрес кратный align. Дополнительно довольно доступно написано про выравнивание у у Фила.

Заключение

Ура, теперь мы можем писать вот такой вот код!

use static_box::Box;

struct Uart1Rx {
    // Implementation details...
}

impl SerialWrite for Uart1Rx {
    fn write(&mut self, _byte: u8) {
        // Implementation details
    }
}
let rx = Uart1Rx { /* ... */ };

SOME_GLOBAL_WRITER.init_once(move || Box::::new(rx));

// A bit of code later.

SOME_GLOBAL_WRITER.lock().unwrap().write_str("Hello world!");

Итак, мы при помощи unsafe и некоторого количества nightly фич смогли написать тип, позволяющий размещать полиморфные объекты на стеке или в статической памяти без использования кучи, что может быть полезным во многих случаях. Хотя, конечно, каждый раз при получении ссылки на объект приходится дополнительно вычислять адрес метаданных и значения, но мы не можем просто так взять и сохранить эти адреса как поля структуры, в этом случае она станет самоссылающиеся, что довольно неприятно в Rust контексте, это не работает с семантикой перемещения. В целом, если воспользоваться pin API, и сделать нашBox неперемещаемым, то можно будет позволить себе эту оптимизацию, а заодно и обеспечить возможность работать с любыми Future типами.

Хочу еще сказать напоследок, что не стоит бояться писать низкоуровневый unsafe код, но стоит 10 раз подумать над его корректностью и обязательно использовать Miri в CI тестах, он отлавливает довольно много ошибок, а разработка низкоуровневого кода требует очень большой внимательности к деталям всевозможным граничным случаям. В конечном счете, именно знания того, как в реальности реализована та или иная языковая абстракция, позволяет перестать воспринимать её как черную магию. Часто все намного проще и очевиднее, чем кажется, стоит просто копнуть чуть поглубже.

А еще важно иногда выходить за рамки stable Rust, чтобы быть в курсе, куда же язык дальше развивается и тем самым расширять свой кругозор.

Ссылка на крейт

© Habrahabr.ru