use diesel_attribute_parser::AttributeSpanWrapper;
use proc_macro2::{Span, TokenStream};
use syn::{Data, Ident, Result, spanned::Spanned};
use crate::util::wrap_in_dummy_mod;
const ERROR_MESSAGE: &str = "this derive can only be used on enums with exclusively unit-variants";
pub fn derive(item: DeriveEnumInput) -> Result<TokenStream> {
let mut to_sql_impls = Vec::new();
let mut from_sql_impls = Vec::new();
for tpe in &item.sql_type_attrs {
let span = tpe.attribute_span;
let tpe = &tpe.item;
let sql_type = quote::quote_spanned! {span=> #tpe};
let enum_name = &item.ident;
let has_explicit_discriminant = item.has_explicit_discriminants;
let variants = item
.variants
.iter()
.map(EnumVariant::as_diesel_enum_variant)
.collect::<Vec<_>>();
let variant_constructor = item.variants.iter().enumerate().map(|(idx, v)| {
let span = v.span;
let ident = &v.rust_name;
quote::quote_spanned! {span=> #idx => Ok(Self::#ident)}
});
let variant_to_enum_variant_mapping = item.variants.iter().map(|v| {
let span = v.span;
let ident = &v.rust_name;
let enum_variant = v.as_diesel_enum_variant();
quote::quote_spanned! {span=> Self::#ident => &#enum_variant}
});
from_sql_impls.push(quote::quote! {
impl<__DB> diesel::deserialize::FromSql<#sql_type, __DB> for #enum_name
where
__DB: diesel::backend::Backend,
#sql_type: diesel::sql_types::EnumSqlType<#has_explicit_discriminant, __DB>,
<#sql_type as diesel::sql_types::EnumSqlType<#has_explicit_discriminant, __DB>>::Strategy: diesel::internal::derives::enum_::EnumMapping<__DB>,
{
fn from_sql(value: <__DB as diesel::backend::Backend>::RawValue<'_>) -> diesel::deserialize::Result<Self> {
const VARIANTS: &[diesel::internal::derives::enum_::EnumVariant] = &[#(#variants,)*];
let idx = <<#sql_type as diesel::sql_types::EnumSqlType<#has_explicit_discriminant, __DB>>::Strategy as diesel::internal::derives::enum_::EnumMapping<__DB>>::map_from_database_value(
value,
stringify!(#enum_name),
VARIANTS
)?;
match idx {
#(#variant_constructor,)*
_ => unreachable!("We construct all relevant variants"),
}
}
}
});
to_sql_impls.push(quote::quote! {
impl<__DB> diesel::serialize::ToSql<#sql_type, __DB> for #enum_name
where
__DB: diesel::backend::Backend,
#sql_type: diesel::sql_types::EnumSqlType<#has_explicit_discriminant, __DB>,
<#sql_type as diesel::sql_types::EnumSqlType<#has_explicit_discriminant, __DB>>::Strategy: diesel::internal::derives::enum_::EnumMapping<__DB>,
{
fn to_sql<'b>(&'b self, output: &mut diesel::serialize::Output<'b, '_, __DB>) -> diesel::serialize::Result {
let variant = match self {
#(#variant_to_enum_variant_mapping,)*
};
<<#sql_type as diesel::sql_types::EnumSqlType<#has_explicit_discriminant, __DB>>::Strategy as diesel::internal::derives::enum_::EnumMapping<__DB>>::map_to_database_value(
output,
variant
)
}
}
});
}
let struct_ty = syn::Type::Path(syn::TypePath {
qself: None,
path: item.ident.into(),
attrs: Vec::new(),
});
let sql_types = item
.sql_type_attrs
.iter()
.map(|v| syn::Type::Path(v.item.clone()))
.collect::<Vec<_>>();
let as_expression_impl = super::as_expression::derive_inner(
sql_types,
syn::Generics::default(),
struct_ty.clone(),
false,
false,
)?;
let from_sql_row_impl = super::from_sql_row::derive_inner(struct_ty, syn::Generics::default())?;
Ok(wrap_in_dummy_mod(quote::quote! {
#(#from_sql_impls)*
#(#to_sql_impls)*
#as_expression_impl
#from_sql_row_impl
}))
}
pub struct DeriveEnumInput {
sql_type_attrs: Vec<AttributeSpanWrapper<syn::TypePath>>,
ident: syn::Ident,
has_explicit_discriminants: bool,
variants: Vec<EnumVariant>,
}
pub struct EnumVariant {
span: Span,
discriminant: i128,
rust_name: Ident,
sql_name: String,
}
impl EnumVariant {
fn as_diesel_enum_variant(&self) -> TokenStream {
let Self {
span,
discriminant,
rust_name,
sql_name,
} = self;
let span = *span;
quote::quote_spanned! {span=>
diesel::internal::derives::enum_::EnumVariant {
discriminant: #discriminant,
rust_name: stringify!(#rust_name),
sql_name: #sql_name,
}
}
}
}
impl syn::parse::Parse for DeriveEnumInput {
fn parse(input: syn::parse::ParseStream) -> Result<Self> {
let input = input.parse::<syn::DeriveInput>()?;
let input_span = input.span();
if !input.generics.params.is_empty() {
return Err(syn::Error::new(input.span(), ERROR_MESSAGE));
}
let enum_ = match input.data {
Data::Enum(data_enum) => data_enum,
_ => {
return Err(syn::Error::new(input.span(), ERROR_MESSAGE));
}
};
let attrs = diesel_attribute_parser::parse_attributes::<diesel_attribute_parser::StructAttr>(
&input.attrs,
)?;
let rename_all = attrs.iter().find_map(|a| {
if let diesel_attribute_parser::StructAttr::RenameAll(_, r) = &a.item {
Some(r)
} else {
None
}
});
let mut has_explicit_discriminants = true;
let mut variants = Vec::with_capacity(enum_.variants.len());
for (idx, v) in enum_.variants.iter().enumerate() {
if !v.fields.is_empty() {
return Err(syn::Error::new(v.span(), ERROR_MESSAGE));
}
has_explicit_discriminants = has_explicit_discriminants && v.discriminant.is_some();
let discriminant = v
.discriminant
.as_ref()
.map(|(_, v)| {
let (f, l) = match v {
syn::Expr::Lit(l) => (1, l),
syn::Expr::Unary(syn::ExprUnary {
op: syn::UnOp::Neg(_),
expr,
..
}) => {
if let syn::Expr::Lit(l) = &**expr {
(-1, l)
} else {
return Err(syn::Error::new(
v.span(),
"expected a literal expression, but got something else",
));
}
}
_ => {
return Err(syn::Error::new(
v.span(),
"expected a literal expression, but got something else",
));
}
};
let syn::Lit::Int(i) = &l.lit else {
return Err(syn::Error::new(
l.span(),
"expected a integer literal expression, but got something else",
));
};
Ok(f * i.base10_parse::<i128>()?)
})
.transpose()?
.unwrap_or(idx as i128);
let rust_name = v.ident.clone();
let attrs = diesel_attribute_parser::parse_attributes::<
diesel_attribute_parser::FieldAttr,
>(&v.attrs)?;
let rename_attr = attrs.iter().find_map(|a| {
if let diesel_attribute_parser::FieldAttr::Rename(_, r) = &a.item {
Some(r.value())
} else {
None
}
});
let sql_name = rename_attr
.or_else(|| Some(rename_all?.apply_case_to_enum_variant(rust_name.to_string())))
.unwrap_or_else(|| rust_name.to_string());
variants.push(EnumVariant {
span: v.span(),
discriminant,
rust_name,
sql_name,
});
}
let sql_type_attrs = attrs
.into_iter()
.filter_map(
|a: AttributeSpanWrapper<diesel_attribute_parser::StructAttr>| {
if let diesel_attribute_parser::StructAttr::SqlType(_, path) = a.item {
Some(diesel_attribute_parser::AttributeSpanWrapper {
item: path,
attribute_span: a.attribute_span,
ident_span: a.ident_span,
})
} else {
None
}
},
)
.collect::<Vec<_>>();
if sql_type_attrs.is_empty() {
return Err(syn::Error::new(
input_span,
"no `#[diesel(sql_type = ...)]` attribute provided",
));
}
Ok(Self {
sql_type_attrs,
ident: input.ident,
has_explicit_discriminants,
variants,
})
}
}