Высокоинтеллектуальный #meme

#prog #rust #моё

Опытные Rust-программисты знают, что в Rust можно сделать мапу, которая индексируется не рантайм-ключами, а типами. Сделать это довольно просто поверх обычной мапы: достаточно хранить значения в виде Box<dyn Any>, а в качестве ключей использовать TypeId (хотя и есть немного неочевидный нюанс). В геттере мы конструируем TypeId от переданного типа, индексируем им мапу, и если значение есть, то приводим к нужному типа через downcast_ref/downcast_mut.

#[derive(Default)]
struct TypeMap(HashMap<TypeId, Box<dyn Any>>);

impl TypeMap {
    fn insert<T: 'static>(&mut self, val: T) {
        self.0.insert(TypeId::of::<T>(), Box::new(val));
    }

    fn get<T: 'static>(&self) -> Option<&T> {
        self.0.get(&TypeId::of::<T>())?.downcast_ref::<T>()
    }
}

fn main() {
    let mut map = TypeMap::default();
    map.insert(42_i32);
    assert_eq!(map.get::<i32>(), Some(&42));
}

С дополнительными обёртками можно делать мапы с отображениями из T в тип, параметризованный T:

T -> Option<T>
T -> [T; 2]
T -> Vec<T>

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

Пусть для определённости мы хотим хранить коллбеки вида FnMut(&mut T). Как я уже говорил, в Rust нет общего типа для замыканий одной сигнатуры — есть лишь общий трейт. Для того, чтобы сохранить замыкание в TypeMap, нам потребуется каким-то образом стереть тип дважды: сначала от конкретного типа замыкания до dyn FnMut(&mut T), а потом от него до dyn Any.

Наиболее простой способ сделать это — забоксить замыкание дважды:

// определение `TypeMap` выше

#[derive(Default)]
struct CallbackMap(TypeMap);

impl CallbackMap {
    fn insert<T: 'static>(&mut self, f: impl FnMut(&mut T) + 'static) {
        let boxed = Box::new(f) as Box<dyn FnMut(&mut T) + 'static>;
        self.0.insert(boxed);
    }
    
    fn get_by_arg<T: 'static>(&mut self) -> Option<&mut (dyn FnMut(&mut T) + 'static)> {
        self.0.get_mut::<Box<dyn FnMut(&mut T) + 'static>>().map(|f| &mut **f)
    }
}

fn main() {
    let mut map = CallbackMap::default();
    let multiplier = 7;
    let f = move |x: &mut i32| *x *= multiplier;
    map.insert(f);
    let mut x = 6;
    map.get_by_arg().unwrap()(&mut x);
    assert_eq!(x, 42);
}

Решение выше полностью безопасно, но неудовлетворительно ввиду двойной аллокации и двойной индирекции (даже тройной, считая саму мапу). Можно ли сделать лучше? Можно... Но для этого нам придётся задействовать фичи nightly, а именно — ptr_metadata. Почему? Потому что — забегая вперёд — нам потребуется стирать типы у жирных толстых указателей на трейт-объекты, и без этого API это невозможно сделать, не закладываясь на детали реализации текущей версии.

Что можно сказать наверняка — так это то, что один раз боксить всё-таки придётся. Вне зависимости от того, что мы сконструируем, это что-то должно будет хранить !Sized значения. И это что-то само по себе должно быть Sized, чтобы его можно было положить в мапу. Таким образом, нам потребуется что-то, что будет хранить указатели. Мы, очевидно, не можем хранить просто Box<dyn FnMut(&mut T)>, поскольку это разные типы для разных T, а для хранения в мапе тип нужен одинаковый.

Немного подумаем о том, что из себя представляет указатель на трейт-объект. Так как это — толстый указатель, он состоит из указателя на данные и... DynMetadata. DynMetadata — это особый тип, который хранит в себе указатель на таблицу. Если DynMetadata инстанцирован неким dyn-типом, то таблица содержит в себе размер и выравнивание реального значения, его деструктор и методы соответствующего трейта. Методы трейт-объекта являются тонкой обёрткой, которые просто загружают нужный метод из этой таблицы и вызывают его на указателе данных. Что приятно — DynMetadata имеет одно и то же представление вне зависимости от подставленного типа... И да, именно на это мы и будем полагаться (еееее, больше ансейфа!).

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

type ErasedVtable = DynMetadata<()>;

struct ErasedDyn {
    data: *mut (),
    vtable: ErasedVtable,
}

impl ErasedDyn {
    fn new<T, F>(f: F) -> Self
    where
        T: 'static,
        F: FnMut(&mut T) + 'static,
    {
        let boxed = Box::new(f) as Box<dyn FnMut(&mut T)>;
        let (data, vtable) = Box::into_raw(boxed).to_raw_parts();
        Self {
            data,
            vtable: unsafe { std::mem::transmute(vtable) },
        }
    }

    unsafe fn downcast_mut_unchecked<T: 'static>(&mut self) -> &mut dyn FnMut(&mut T) {
        let data_ptr = self.data;
        let vtable = std::mem::transmute(self.vtable);
        &mut *std::ptr::from_raw_parts_mut(data_ptr, vtable)
    }
}

#[derive(Default)]
struct CallbackMap {
    map: HashMap<TypeId, ErasedDyn>,
}

impl CallbackMap {
    fn insert<T: 'static>(&mut self, f: impl FnMut(&mut T) + 'static) {
        self.map.insert(TypeId::of::<T>(), ErasedDyn::new(f));
    }

    fn get<T: 'static>(&mut self) -> Option<&mut dyn FnMut(&mut T)> {
        self.map
            .get_mut(&TypeId::of::<T>())
            .map(|cb| unsafe { cb.downcast_mut_unchecked() })
    }
}

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

Есть только одна небольшая проблема.

Мы не можем нормально дропать ErasedDyn.

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

Кажется, мы зашли в тупик. Мы можем легко получить TypeId по T, но не наоборот. Возможно, можно было бы хранить деструктор вместе с ErasedDyn, но если мы добавим возможность удалять коллбеки из мапы, то у нас не выйдет автоматически дропать его вне мапы. Можно добавить деструктор и в сам ErasedDyn, но тогда он будет занимать размер трёх usize против двух usize у нативных толстых указателей. Видимо, выход один: хранить деструктор в памяти, на который указывает ErasedDyn.

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

type Destructor = unsafe fn(*mut (), ErasedVtable);

fn make_drop_fn<T>() -> Destructor {
    |data, vtable| unsafe {
        let dyn_ptr: *mut ??? =
            std::ptr::from_raw_parts_mut(data, std::mem::transmute(vtable));
        drop(Box::from_raw(dyn_ptr))
    }
}

Возникает небольшой вопрос, что подставить вместо ???. Это не может быть dyn FnMut(&mut T), так как нам ещё нужно как-то добавить деструктор. Что ж, сделаем под это тип:

#[repr(C)]
struct WithDrop<T: ?Sized> {
    destructor: Destructor,
    data: T,
}

В make_drop_fn вместо ??? подставим WithDrop<dyn FnMut(&mut T)>. Очевидно, конструктор ErasedDyn придётся немного поменять:

impl ErasedDyn {
    fn new<T, F>(f: F) -> Self
    where
        T: 'static,
        F: FnMut(&mut T) + 'static,
    {
        let with_drop = WithDrop {
            destructor: make_drop_fn::<T>(),
            data: f,
        };
        let boxed = Box::new(with_drop) as Box<WithDrop<dyn FnMut(&mut T)>>;
        let (data, vtable) = Box::into_raw(boxed).to_raw_parts();
        Self {
            data,
            vtable: unsafe { std::mem::transmute(vtable) },
        }
    }
}

Ну и, раз уж мы сделали деструктор, то, очевидно, надо сделать надлежащую реализацию Drop:

impl Drop for ErasedDyn {
    fn drop(&mut self) {
        unsafe {
            let destructor = self.data.cast::<Destructor>().read();
            destructor(self.data, self.vtable);
        }
    }
}

Каст из data в указатель на Destructor валиден за счёт того, что WithDrop имеет фиксированную раскладку в памяти — #[repr(C)] — и адрес первого поля совпадает с адресом структуры в целом. Окей, с дропом мы разобрались, но теперь нам нужно поменять метод для даункаста. Теперь нам надо каким-то образом получить из указателя на WithDrop<dyn FnMut(&mut T)> указатель на data. Кажется, достаточно просто сместить указатель на размер Destructor?..

impl ErasedDyn {
    unsafe fn downcast_mut_unchecked<T: 'static>(&mut self) -> &mut dyn FnMut(&mut T) {
        let data_ptr = self
            .data
            .byte_offset(std::mem::size_of::<Destructor>() as isize);
        let vtable = std::mem::transmute(self.vtable);
        &mut *std::ptr::from_raw_parts_mut(data_ptr, vtable)
    }
}

А вот и нет.

Именно, этот unsound код полагается на то, что data располагается в памяти сразу после destructor, без промежутков. И это абсолютно необоснованное предположение! Именно, unsized coercion добавляет vtable к указателям на данные, но, разумеется, никак не трогает данные за указателем. Раскладка в памяти данных за указателем зависит от исходного типа, из которого произошло unsized coercion — и конкретно в случае WithData тип поля data вполне мог иметь выравнивание большее, чем у Destructor! Как следствие, между destructor и data может быть паддинг, и, как несложно понять, подсчитать его статически после стирания типов нельзя.

Покажу наглядно, как ломается текущий подход:

fn main() {
    #[repr(align(16))]
    struct OverAligned(i32);

    let mut map = CallbackMap::default();
    let multiplier = OverAligned(7);
    // захватываем multiplier целиком, а не только поле
    let f = move |x: &mut i32| *x *= { &multiplier }.0;
    map.insert(f);
    let mut x = 6;
    map.get().unwrap()(&mut x);
    assert_eq!(x, 42);
}

Неправильная реализация downcast_mut выдаёт указатель куда-то в паддинг, в неинициализированные данные. Попытка их прочитать приводит к неопределённому поведению. На практике это означает, что ассерт у вас почти наверняка выстрелит (а в дебажной сборке — может и свалиться с паникой из-за переполнения при умножении). Если запустить этот пример в miri, то он выдаёт ошибку на момент создания ссылки — некорректное выравнивание.

Так как же всё починить? Очевидно, нам нужно хранить не только деструктор, но и смещение до данных. Поменяем WithDrop на WithHeader:

#[derive(Clone, Copy)]
struct Header {
    data_offset: isize,
    destructor: Destructor,
}

#[repr(C)]
struct WithHeader<T: ?Sized> {
    header: Header,
    data: T,
}

В конструкторе ErasedDyn нужно также дополнительно высчитывать смещение до собственно данных:

impl ErasedDyn {
    fn new<T, F>(f: F) -> Self
    where
        F: FnMut(&mut T) + 'static,
    {
        let mut with_hdr = WithHeader {
            header: Header {
                data_offset: 0,
                destructor: make_drop_fn::<T>(),
            },
            data: f,
        };
        unsafe {
            with_hdr.header.data_offset =
                <*const _>::byte_offset_from(&with_hdr.data, &with_hdr);
        }
        let boxed = Box::new(with_hdr) as Box<WithHeader<dyn FnMut(&mut T)>>;
        let (data, vtable) = Box::into_raw(boxed).to_raw_parts();
        Self {
            data,
            vtable: unsafe { std::mem::transmute(vtable) },
        }
    }
}

Деструктор почти не поменялся, только вместо вычитывания деструктора напрямую по указателю указатель сначала кастуется в Header и из него уже достаётся destructor. Немного сложнее даункаст:

impl ErasedDyn {
    unsafe fn downcast_mut_unchecked<T>(&mut self) -> &mut dyn FnMut(&mut T) {
        let data_offset = (*self.data.cast::<Header>()).data_offset;
        let data_ptr = self.data.byte_offset(data_offset);
        let vtable = std::mem::transmute(self.vtable);
        &mut *std::ptr::from_raw_parts_mut(data_ptr, vtable)
    }
}

И вот это уж корректно работает. miri не жалуется даже на это:

fn main() {
    #[repr(align(128))]
    struct A(String);
    
    impl A {
        fn get(&self) -> &str {
            &self.0
        }
    }

    let mut a = A(String::from("a"));
    let f = move |x: &mut i32| {
        *x += a.get().len() as i32;
        a.0 += "a";
    };
    let mut map = CallbackMap::default();
    map.insert(f);
    let mut x = 4;
    map.get().unwrap()(&mut x);
    assert_eq!(x, 5);
    map.get().unwrap()(&mut x);
    assert_eq!(x, 7);
}

===============

БОНУС: если мы хотим обойтись одной аллокацией и не хотим быть менее эффективным, чем нативные толстые указатели, но готовы немного поступиться типобезопасностью, то хранить коллбеки со стиранием типа можно и сильно проще:

#[derive(Default)]
struct ErasedCallbackMap(HashMap<TypeId, Box<dyn FnMut(&mut dyn Any)>>);

fn wrap<T: 'static>(mut f: impl FnMut(&mut T) + 'static) -> Box<dyn FnMut(&mut dyn Any)> {
    Box::new(move |arg| f(arg.downcast_mut().unwrap()))
}

impl ErasedCallbackMap {
    fn insert<T: 'static>(&mut self, f: impl FnMut(&mut T) + 'static) {
        self.0.insert(TypeId::of::<T>(), wrap(f));
    }

    fn get<T: 'static>(&mut self) -> Option<&mut (dyn FnMut(&mut dyn Any) + 'static)> {
        self.0.get_mut(&TypeId::of::<T>()).map(|f| &mut **f)
    }
}

fn main() {
    let mut map = ErasedCallbackMap::default();
    let multiplier = 7;
    let f = move |x: &mut i32| *x *= multiplier;
    map.insert(f);
    let mut x = 6;
    // аргумент приводится к &mut dyn Any
    map.get::<i32>().unwrap()(&mut x);
    assert_eq!(x, 42);
}

МЕЖДУНАРОДНОЕ ОБЩЕСТВЕННОЕ ДВИЖЕНИЕ ЛГБТ

Сегодня одному из моих подписчиков исполнилось 18 лет.

С днём рождения, папищек.

Как указано в исковом заявлении, если в аббревиатуре ЛГБТ передвинуть по одной букве, получится слово КВАС. Поэтому движение угрожает основам культуры России.

https://t.iss.one/ia_panorama/14790

Новость от 24 февраля, на минуточку.

Я тупо умный

#game

В Steam сейчас раздача слонов первой Half-Life:
https://t.iss.one/welovegames/18558

Дело о запрете "международного общественного движения ЛГБТ" будет слушаться в закрытом судебном заседании. Значит, кому именно и какие претензии предъявляются мы не узнаем

Валидация на кассовом чеке

Это не очень хороший год, но следующий будет хуже