parasol_cpu_macros/
lib.rs

1#![deny(missing_docs)]
2#![deny(rustdoc::broken_intra_doc_links)]
3//! Derive macros used with the Parasol processor.
4
5extern crate proc_macro;
6use proc_macro::TokenStream;
7use proc_macro2::TokenStream as TokenStream2;
8use quote::quote;
9use syn::{Data, DeriveInput, Fields, parse_macro_input};
10
11#[proc_macro_derive(IntoBytes)]
12/// Allows you to `#[derive(IntoBytes)]` on structures where each member impls
13/// `IntoBytes`
14pub fn derive_into_bytes(item: TokenStream) -> TokenStream {
15    let input = parse_macro_input!(item as DeriveInput);
16
17    derive_into_bytes_impl(input).into()
18}
19
20fn derive_into_bytes_impl(item: DeriveInput) -> TokenStream2 {
21    let ident = item.ident;
22    let data = item.data;
23
24    let fields = match data {
25        Data::Struct(x) => x.fields,
26        _ => {
27            return quote! {
28                compile_error!("Only structs are supported for derive IntoBytes")
29            };
30        }
31    };
32
33    let field_idents = match &fields {
34        Fields::Unnamed(x) => x
35            .unnamed
36            .iter()
37            .enumerate()
38            .map(|(i, _)| quote! { #i })
39            .collect::<Vec<_>>(),
40        Fields::Named(x) => x
41            .named
42            .iter()
43            .map(|f| {
44                let ident = f.ident.clone().unwrap();
45
46                quote! { #ident }
47            })
48            .collect::<Vec<_>>(),
49        Fields::Unit => {
50            return quote! {
51                compile_error("Structs must have at least one member.")
52            };
53        }
54    };
55
56    let field_types = fields.iter().cloned().map(|x| x.ty).collect::<Vec<_>>();
57
58    let mut prev_field_size = vec![quote! { 0usize }];
59
60    for f in field_types.iter().take(field_types.len() - 1) {
61        prev_field_size.push(quote! {
62            <#f as tfhe_cpu::IntoBytes>::size()
63        })
64    }
65
66    let last_field_type = field_types.iter().last().unwrap();
67
68    quote! {
69        impl tfhe_cpu::IntoBytes for #ident {
70            #[inline(always)]
71            fn alignment() -> usize {
72                let mut alignment = 0usize;
73
74                #(alignment = alignment.max(<#field_types as tfhe_cpu::IntoBytes>::alignment());)*
75
76                alignment
77            }
78
79            #[inline(always)]
80            fn size() -> usize {
81                let mut cur_offset = 0;
82
83                #(
84                    cur_offset += #prev_field_size;
85                    cur_offset = cur_offset.next_multiple_of(<#field_types as tfhe_cpu::IntoBytes>::alignment());
86                )*
87
88                cur_offset += <#last_field_type as tfhe_cpu::IntoBytes>::size();
89
90                cur_offset
91            }
92
93            fn try_into_bytes(&self, data: &mut [u8]) -> tfhe_cpu::Result<()> {
94                if data.len() != Self::size() {
95                    return Err(tfhe_cpu::Error::buffer_size_mismatch());
96                }
97
98                let mut cur_offset = 0;
99
100                #(
101                    cur_offset += #prev_field_size;
102                    cur_offset = cur_offset.next_multiple_of(<#field_types as tfhe_cpu::IntoBytes>::alignment());
103
104                    let byte_slice = &mut data[cur_offset..cur_offset + <#field_types as tfhe_cpu::IntoBytes>::size()];
105                    <#field_types as tfhe_cpu::IntoBytes>::try_into_bytes(&self. #field_idents, byte_slice)?;
106                )*
107
108                Ok(())
109            }
110
111            fn try_from_bytes(data: &[u8]) -> tfhe_cpu::Result<Self> {
112                if data.len() != Self::size() {
113                    return Err(tfhe_cpu::Error::buffer_size_mismatch());
114                }
115
116                let mut cur_offset = 0;
117
118                #(
119                    cur_offset += #prev_field_size;
120                    cur_offset = cur_offset.next_multiple_of(<#field_types as tfhe_cpu::IntoBytes>::alignment());
121
122                    let byte_slice = &data[cur_offset..cur_offset + <#field_types as tfhe_cpu::IntoBytes>::size()];
123                    let #field_idents = <#field_types as tfhe_cpu::IntoBytes>::try_from_bytes( byte_slice)?;
124                )*
125
126
127                Ok(Self {
128                    #(#field_idents,)*
129                })
130            }
131        }
132    }
133}
parasol_cpu_macros/lib.rs

parasol_cpu_macros/
lib.rs
