Расширяем возможности процедурных макросов с помощью WASM

от автора

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

  • Сложность с поддержкой таких макросов в IDE. По сути дела нужно как-то научить анализатор кода самостоятельно компилировать, загружать и исполнять эти самые макросы с учетом всех особенностей. Это весьма нетривиальная задача.
  • Так как макросы самодостаточные и ничего не знают друг о друге, то нет никакой возможности делать композицию макросов, что иногда могло бы быть полезным.

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

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

Пусть у нас имеется некоторый макрос TextMessage, который выводит для заданного типа трейты ToString и FromStr используя в качестве текстового представления некоторый кодек. У разных типов сообщений может быть различный кодек, причем их полный список со временем может расширятся, а у каждого кодека может быть свой уникальный набор атрибутов.

#[derive(Debug, Serialize, Deserialize, PartialEq, TextMessage)] #[text_message(codec = "serde_json", params(pretty))] struct FooMessage {     name: String,     description: String,     value: u64, }

Чтобы сделать такой макрос возможным, мы должны динамически подгружать реализации кодеков в процессе выполнения макроса. Можно вынести кодеки в подключаемую библиотеку и просто загружать их через libloading, но это очень неудобно и еще больше отдалит нас от возможности поддержки макросов в IDE. Вместо этого теоретически возможно написать такой вот кодек на динамическом языке типа Питона, но тогда нам придется писать для Питона аналоги syn и quote, что будет больше напоминать Сизифов труд, чем реальное решение проблемы.
Наиболее же простым и удобным видится вариант скомпилировать кодек в WASM модуль, объединив плюсы обоих подходов. Именно таким путем я предлагаю и пойти.

Выбираем подход к реализации

На первый взгляд кажется, что проблема уже решена в рамках watt и можно просто использовать его для загрузки и выполнения WASM модулей, но с этим подходом есть один весьма неприятный недостаток. Для своей работы watt использует модифицированный крейт proc-macro2, что частенько приводит к непонятным или трудноуловимым проблемам. Например, у меня не компилировался darling или если я забывал подменять proc-macro2, то получал в рантайме неочевидные ошибки.

В результате я решил, что лучше уж пользоваться ванильным proc-macro2, а в качестве WASM
рантайма взять какой-нибудь из самых популярных. В результате, мой субъективный выбор пал на wasmtime, этот рантайм разрабатывается сообществом bytecodealliance, в состав которого входят такие гиганты, как Mozilla, Intel и RedHat. И хотя wasmtime сейчас выглядит еще достаточно сырым, в нем не хватает документации, хороших примеров, но развивается он очень быстро и улучшается прямо на глазах

Взаимодействие хостом и таргетом

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

Погнали!

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

В самом простом виде интерфейс плагина для нашего процедурного макроса должен представлять из себя вариацию на тему:

pub fn implement_codec(input: TokenStream) -> TokenStream;

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

pub fn implement_codec(input: &str) -> String;

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

В целях обеспечения безопасности и большей стабильности, Республика будет реорганизована нами в первую Галактическую Империю, во имя сохранности и во имя блага общества! память WASM рантайма отделена от хостовой, с точки зрения хоста это просто плоский массив байт, в котором находится код программы, глобальные переменные, стек и куча.
Есть возможность сделать так, чтобы память была расширяемой, то если если при очередном выделении памяти нам не хватает места, то верхняя граница памяти автоматически увеличивается. Индекс ячейки в этом самом массиве используется в качестве указателя внутри таргета, но мы не можем просто взять и записать строчку в случайный участок памяти и отдать таргету индекс его начала, потому что снаружи мы не знаем то, как таргет в реальности использует память, где у него находится стек, а где куча. Но мы можем пойти на хитрость: с хоста обратиться к менеджеру памяти таргета и попросить у него выделить нам участок памяти.

#[no_mangle] pub unsafe extern "C" fn toy_alloc(size: i32) -> i32 {     let size_bytes: [u8; 4] = size.to_le_bytes();     let mut buf: Vec<u8> = Vec::with_capacity(size as usize + size_bytes.len());     // Первые 4 байта - это длина общая куска памяти, она нам еще понадобится в      // дальнейшем.     buf.extend(size_bytes.iter());     to_host_ptr(buf) }  unsafe fn to_host_ptr(mut buf: Vec<u8>) -> i32 {     let ptr = buf.as_mut_ptr();     // Просто забываем о выделенном участке памяти, позволяя ему "утечь", таким      // образом мы передаем его во владение хосту.     mem::forget(buf);     ptr as *mut c_void as usize as i32 }  #[no_mangle] pub unsafe extern "C" fn toy_free(ptr: i32) {     let ptr = ptr as usize as *mut u8;     let mut size_bytes = [0u8; 4];     ptr.copy_to(size_bytes.as_mut_ptr(), 4);     // Вычитываем общую длину куска памяти для того, чтобы корректно выполнить      // его очистку.     let size = u32::from_le_bytes(size_bytes) as usize;     // Собираем вектор, о котором мы ранее "забыли" в методе `to_host_ptr` и      // таким образом даем его деструктору вызваться нормальным образом и очистить     // ранее выделенный участок памяти.     Vec::from_raw_parts(ptr, size, size); }

В принципе, ничего хитрого на самом деле в этом нет, примерно этим же занимается wasm_bindgen.

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

#[no_mangle] pub unsafe extern "C" fn implement_codec(     item_ptr: i32,     item_len: i32, ) -> i32 {     let item = str_from_raw_parts(item_ptr, item_len);     let item = TokenStream::from_str(&item).expect("Unable to parse item");      // Здесь уже вызывается типичная функция, реализующая процедурный макрос.     // `fn(item: TokenStream) -> TokenStream`     let tokens = codec::implement_codec(item);     let out = tokens.to_string();      to_host_buf(out) }  pub unsafe fn str_from_raw_parts<'a>(ptr: i32, len: i32) -> &'a str {     let slice = std::slice::from_raw_parts(ptr as *const u8, len as usize);     std::str::from_utf8(slice).unwrap() }

Код хостовой части состоит из двух основных компонент, первым из которых является загрузчик WASM модуля.

 pub struct WasmMacro {     module: Module, }  impl WasmMacro {     // Конструктор нашего макроса расширения.     pub fn from_file(file: impl AsRef<Path>) -> anyhow::Result<Self> {         // Загружаем и компилируем WASM модуль, находящийся по заданному пути.         let store = Store::default();         let module = Module::from_file(&store, file)?;         Ok(Self { module })     }      // Вызываем метод с именем `fun` внутри нашего модуля, в котором содержится     // основная логика преобразования входного TokenStream в выходной.     pub fn proc_macro_derive(         &self,         fun: &str,         item: TokenStream,     ) -> anyhow::Result<TokenStream> {         // Как уже описывалось ранее, чтобы передавать TokenStream между средами,         // нам необходимо преобразовать его в строку.         let item = item.to_string();          // Создаем конкретный экземпляр модуля, с которым и будем работать.         let instance = Instance::new(&self.module, &[])?;         // Получаем указатель на нужную нам функцию, в данном случае это          // описанная выше `implement_codec`.         let proc_macro_attribute_fn = instance             .get_export(fun)             .ok_or_else(|| anyhow!("Unable to find `{}` method in the export table", fun))?             .func()             .ok_or_else(|| anyhow!("export {} is not a function", fun))?             .get2::<i32, i32, i32,>()?;          // Для передачи данных строки внутрь WASM модуля используем специальную         // обертку, о которой я подробнее расскажу ниже.         let item_buf = WasmBuf::from_host_buf(&instance, item);         // Получим из обертки указатель на начало строки и ее длину в байтах         let (item_ptr, item_len) = item_buf.raw_parts();         // А теперь вызываем искомый метод и в результате получаем указатель         // на начало строки с выходным TokenStream.          let ptr = proc_macro_attribute_fn(item_ptr, item_len).unwrap();         // Оборачиваем сырой указатель и читаем получившуюся строку.         let res = WasmBuf::from_raw_ptr(&instance, ptr);         let res_str = std::str::from_utf8(res.as_ref())?;         // В заключительном этапе парсим строку в TokenStream и возращаем выше.         TokenStream::from_str(&res_str)             .map_err(|_| anyhow!("Unable to parse token stream"))     } }

Теперь давайте чуть подробнее рассмотрим WasmBuf модуль: по сути дела это умный указатель,
который владеет некоторой частью памяти, выделенной при помощи toy_alloc. Рассмотрим
самые интересные его части, а остальной код можно посмотреть в репозитории.

struct WasmBuf<'a> {     // Индекс начала выделенного буфера, проще говоря, указатель на его начало.     offset: usize,     // Длина буфера в байтах.     len: usize,     // Ссылка на инстанс модуля, в котором выделялась память     instance: &'a Instance,     // Ссылка на всю память, связанную с этим инстансом.     memory: &'a Memory, }  const WASM_PTR_LEN: usize = 4;  impl<'a> WasmBuf<'a> {     // Самый простой конструктор буфера: мы просто при помощи `toy_alloc`     // запрашиваем искомое число байт.     pub fn new(instance: &'a Instance, len: usize) -> Self {         let memory = Self::get_memory(instance);         // Выделяем память и получаем на нее указатель.         let offset = Self::toy_alloc(instance, len);          Self {             offset: offset as usize,             len,             instance,             memory,         }     }      // Намного удобнее не просто запрашивать буфер, а потом руками заполнять его,     // а сразу передать ссылку на байты, которые мы хотим в него записать.     pub fn from_host_buf(instance: &'a Instance, bytes: impl AsRef<[u8]>) -> Self {         let bytes = bytes.as_ref();         let len = bytes.len();          let mut wasm_buf = Self::new(instance, len);         // Копируем байты с хостового буфера в буфер таргета.         wasm_buf.as_mut().copy_from_slice(bytes);         wasm_buf     }      // Если же буфер был выделен внутри таргета, то все становится несколько     // сложнее. Так как получить мы можем лишь указатель на начало буфера и      // непонятно каким же образом мы получим размер выделенной памяти.     // Но мы не зря написали `toy_alloc` таким образом, чтобы первые его 4     // байта содержали размер выделенного буфера.     pub fn from_raw_ptr(instance: &'a Instance, offset: i32) -> Self {         let offset = offset as usize;         let memory = Self::get_memory(instance);          let len = unsafe {             // Получаем сырой указатель на память инстанса.             let buf = memory.data_unchecked();              let mut len_bytes = [0; WASM_PTR_LEN];             // Читаем байты с размером выделенного буфера.             len_bytes.copy_from_slice(&buf[offset..offset + WASM_PTR_LEN]);             u32::from_le_bytes(len_bytes)         };          Self {             offset,             len: len as usize,             memory,             instance,         }     }      // Методы для чтения и записи данных являются весьма тривиальными.     // Важно лишь помнить про то, что нужно читать со смещением в 4 байта.      pub fn as_ref(&self) -> &[u8] {         unsafe {             let begin = self.offset + WASM_PTR_LEN;             let end = begin + self.len;              &self.memory.data_unchecked()[begin..end]         }     }      pub fn as_mut(&mut self) -> &mut [u8] {         unsafe {             let begin = self.offset + WASM_PTR_LEN;             let end = begin + self.len;              &mut self.memory.data_unchecked_mut()[begin..end]         }     }     }

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

impl Drop for WasmBuf<'_> {     fn drop(&mut self) {         Self::toy_free(self.instance, self.len);     } }

Собираем все вместе

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

#[proc_macro_derive(TextMessage, attributes(text_message))] pub fn text_message(input: TokenStream) -> TokenStream {     let input: DeriveInput = parse_macro_input!(input);      let attrs = TextMessageAttrs::from_raw(&input.attrs)         .expect("Unable to parse text message attributes.");      // Для простоты будем грузить модули из директории codecs, которые имеют      // особым образом сформированное имя.      let codec_dir = Path::new(&std::env::var("CARGO_MANIFEST_DIR")         .unwrap())         .join("codecs");     let plugin_name = format!("{}_text_codec.wasm", attrs.codec);     let codec_path = codec_dir.join(plugin_name);      let wasm_macro = WasmMacro::from_file(codec_path)         .expect("Unable to load wasm module");      wasm_macro         .proc_macro_derive(             "implement_codec",             input.into_token_stream().into(),         )         .expect("Unable to apply proc_macro_attribute") }

В репозитории есть готовый пример с демонстрацией работы. Вы можете убедиться, что кодек действительно загружается из WASM модуля, а не компилируется вместе с макросом.

#[derive(Debug, Serialize, Deserialize, PartialEq, TextMessage)] // Что особенно хорошо, каждый WASM плагин может иметь свои произвольные атрибуты. #[text_message(codec = "serde_json", params(pretty))] struct FooMessage {     name: String,     description: String,     value: u64, }  fn main() {     let msg = FooMessage {         name: "Linus Torvalds".to_owned(),         description: "The Linux founder.".to_owned(),         value: 1,     };      let text = msg.to_string();     println!("{}", text);     let msg2 = text.parse().unwrap();      assert_eq!(msg, msg2); }

Выводы

Пока это больше похоже на троллейбус из буханки хлеба, но с другой стороны это небольшая, но прекрасная демонстрация самого принципа. Такие макросы становятся открытыми для расширения. У нас больше нет необходимости в переписывании исходного процедурного макроса, чтобы изменить или расширить его поведение. А если же воспользоваться реестром модулей для WASM, то можно будет распространять подобные модули подобно крейтам cargo.

ссылка на оригинал статьи https://habr.com/ru/post/497916/


Комментарии

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *