blob: 20a338923a6517bea808fc161e7e06f732d1bbae [file] [log] [blame]
// Copyright © 2023 Collabora, Ltd.
// SPDX-License-Identifier: MIT
use proc_macro::TokenStream;
use proc_macro2::{Span, TokenStream as TokenStream2};
use syn::*;
fn expr_as_usize(expr: &syn::Expr) -> usize {
let lit = match expr {
syn::Expr::Lit(lit) => lit,
_ => panic!("Expected a literal, found an expression"),
};
let lit_int = match &lit.lit {
syn::Lit::Int(i) => i,
_ => panic!("Expected a literal integer"),
};
assert!(lit.attrs.is_empty());
lit_int
.base10_parse()
.expect("Failed to parse integer literal")
}
fn count_type(ty: &Type, slice_type: &str) -> usize {
match ty {
syn::Type::Array(a) => {
let elems = count_type(a.elem.as_ref(), slice_type);
if elems > 0 {
elems * expr_as_usize(&a.len)
} else {
0
}
}
syn::Type::Path(p) => {
if p.qself.is_none() && p.path.is_ident(slice_type) {
1
} else {
0
}
}
_ => 0,
}
}
fn get_attr(field: &Field, attr_name: &str) -> Option<String> {
for attr in &field.attrs {
if let Meta::List(ml) = &attr.meta {
if ml.path.is_ident(attr_name) {
return Some(format!("{}", ml.tokens));
}
}
}
None
}
pub fn derive_as_slice(
input: TokenStream,
slice_type: &str,
attr_name: &str,
attr_type: &str,
) -> TokenStream {
let DeriveInput {
attrs, ident, data, ..
} = parse_macro_input!(input);
match data {
Data::Struct(s) => {
let mut has_repr_c = false;
for attr in attrs {
match attr.meta {
Meta::List(ml) => {
if ml.path.is_ident("repr")
&& format!("{}", ml.tokens) == "C"
{
has_repr_c = true;
}
}
_ => (),
}
}
assert!(has_repr_c, "Struct must be declared #[repr(C)]");
let mut first = None;
let mut count = 0_usize;
let mut found_last = false;
let mut attrs = TokenStream2::new();
if let Fields::Named(named) = s.fields {
for f in named.named {
let f_count = count_type(&f.ty, slice_type);
let f_attr = get_attr(&f, &attr_name);
if f_count > 0 {
assert!(
!found_last,
"All fields of type {slice_type} must be consecutive",
);
let attr_type =
Ident::new(attr_type, Span::call_site());
let f_attr = if let Some(s) = f_attr {
let s = syn::parse_str::<Ident>(&s).unwrap();
quote! { #attr_type::#s, }
} else {
quote! { #attr_type::DEFAULT, }
};
first.get_or_insert(f.ident);
for _ in 0..f_count {
attrs.extend(f_attr.clone());
}
count += f_count;
} else {
assert!(
f_attr.is_none(),
"{attr_name} attribute is only allowed on {slice_type}"
);
if !first.is_none() {
found_last = true;
}
}
}
} else {
panic!("Fields are not named");
}
let slice_type = Ident::new(slice_type, Span::call_site());
let attr_type = Ident::new(attr_type, Span::call_site());
if let Some(first) = first {
quote! {
impl compiler::as_slice::AsSlice<#slice_type> for #ident {
type Attr = #attr_type;
fn as_slice(&self) -> &[#slice_type] {
unsafe {
let first = &self.#first as *const #slice_type;
std::slice::from_raw_parts(first, #count)
}
}
fn as_mut_slice(&mut self) -> &mut [#slice_type] {
unsafe {
let first =
&mut self.#first as *mut #slice_type;
std::slice::from_raw_parts_mut(first, #count)
}
}
fn attrs(&self) -> AttrList<Self::Attr> {
static ATTRS: [#attr_type; #count] = [#attrs];
AttrList::Array(&ATTRS)
}
}
}
} else {
quote! {
impl compiler::as_slice::AsSlice<#slice_type> for #ident {
type Attr = #attr_type;
fn as_slice(&self) -> &[#slice_type] {
&[]
}
fn as_mut_slice(&mut self) -> &mut [#slice_type] {
&mut []
}
fn attrs(&self) -> AttrList<Self::Attr> {
AttrList::Uniform(#attr_type::DEFAULT)
}
}
}
}
.into()
}
Data::Enum(e) => {
let mut as_slice_cases = TokenStream2::new();
let mut as_mut_slice_cases = TokenStream2::new();
let mut types_cases = TokenStream2::new();
let slice_type = Ident::new(slice_type, Span::call_site());
let attr_type = Ident::new(attr_type, Span::call_site());
for v in e.variants {
let case = v.ident;
as_slice_cases.extend(quote! {
#ident::#case(x) => compiler::as_slice::AsSlice::<#slice_type>::as_slice(x),
});
as_mut_slice_cases.extend(quote! {
#ident::#case(x) => compiler::as_slice::AsSlice::<#slice_type>::as_mut_slice(x),
});
types_cases.extend(quote! {
#ident::#case(x) => compiler::as_slice::AsSlice::<#slice_type>::attrs(x),
});
}
quote! {
impl compiler::as_slice::AsSlice<#slice_type> for #ident {
type Attr = #attr_type;
fn as_slice(&self) -> &[#slice_type] {
match self {
#as_slice_cases
}
}
fn as_mut_slice(&mut self) -> &mut [#slice_type] {
match self {
#as_mut_slice_cases
}
}
fn attrs(&self) -> AttrList<Self::Attr> {
match self {
#types_cases
}
}
}
}
.into()
}
_ => panic!("Not a struct type"),
}
}