1use proc_macro2::{Span, TokenStream};
4use syn::{
5 parse_quote, Data, DataEnum, DataStruct, DataUnion, Error, Expr, ExprLit, ExprUnary, Lit, UnOp,
6 WherePredicate,
7};
8
9use crate::{
10 derive::try_from_bytes::derive_try_from_bytes,
11 repr::{CompoundRepr, EnumRepr, Repr, Spanned},
12 util::{enum_size_from_repr, Ctx, FieldBounds, ImplBlockBuilder, Trait, TraitBound},
13};
14pub(crate) fn find_zero_variant(enm: &DataEnum) -> Result<usize, bool> {
20 let mut next_negative_discriminant = Some(0);
28
29 let mut has_unknown_discriminants = false;
36
37 for (i, v) in enm.variants.iter().enumerate() {
38 match v.discriminant.as_ref() {
39 None => {
41 match next_negative_discriminant.as_mut() {
42 Some(0) => return Ok(i),
43 Some(n) => *n -= 1,
45 None => (),
46 }
47 }
48 Some((_, Expr::Lit(ExprLit { lit: Lit::Int(int), .. }))) => {
50 match int.base10_parse::<u128>().ok() {
51 Some(0) => return Ok(i),
52 Some(_) => next_negative_discriminant = None,
53 None => {
54 has_unknown_discriminants = true;
56 next_negative_discriminant = None;
57 }
58 }
59 }
60 Some((_, Expr::Unary(ExprUnary { op: UnOp::Neg(_), expr, .. }))) => match &**expr {
62 Expr::Lit(ExprLit { lit: Lit::Int(int), .. }) => {
63 match int.base10_parse::<u128>().ok() {
64 Some(0) => return Ok(i),
65 Some(x) => next_negative_discriminant = Some(x - 1),
67 None => {
68 has_unknown_discriminants = true;
71 next_negative_discriminant = None;
72 }
73 }
74 }
75 _ => {
77 has_unknown_discriminants = true;
78 next_negative_discriminant = None;
79 }
80 },
81 _ => {
83 has_unknown_discriminants = true;
84 next_negative_discriminant = None;
85 }
86 }
87 }
88
89 Err(has_unknown_discriminants)
90}
91pub(crate) fn derive_from_zeros(ctx: &Ctx, top_level: Trait) -> Result<TokenStream, Error> {
92 if let Some(span) = ctx.invariant_span {
93 return ctx.error_or_skip(Error::new(
94 span,
95 "cannot derive `FromZeros` for a type with invariants",
96 ));
97 }
98 let try_from_bytes = derive_try_from_bytes(ctx, top_level)?;
99 let from_zeros = match &ctx.ast.data {
100 Data::Struct(strct) => derive_from_zeros_struct(ctx, strct),
101 Data::Enum(enm) => derive_from_zeros_enum(ctx, enm)?,
102 Data::Union(unn) => derive_from_zeros_union(ctx, unn),
103 };
104 Ok(IntoIterator::into_iter([try_from_bytes, from_zeros]).collect())
105}
106pub(crate) fn derive_from_bytes(ctx: &Ctx, top_level: Trait) -> Result<TokenStream, Error> {
107 if let Some(span) = ctx.invariant_span {
108 return ctx.error_or_skip(Error::new(
109 span,
110 "cannot derive `FromBytes` for a type with invariants",
111 ));
112 }
113 let from_zeros = derive_from_zeros(ctx, top_level)?;
114 let from_bytes = match &ctx.ast.data {
115 Data::Struct(strct) => derive_from_bytes_struct(ctx, strct),
116 Data::Enum(enm) => derive_from_bytes_enum(ctx, enm)?,
117 Data::Union(unn) => derive_from_bytes_union(ctx, unn),
118 };
119
120 Ok(IntoIterator::into_iter([from_zeros, from_bytes]).collect())
121}
122fn derive_from_zeros_struct(ctx: &Ctx, strct: &DataStruct) -> TokenStream {
123 ImplBlockBuilder::new(ctx, strct, Trait::FromZeros, FieldBounds::ALL_SELF).build()
124}
125fn derive_from_zeros_enum(ctx: &Ctx, enm: &DataEnum) -> Result<TokenStream, Error> {
126 let repr = EnumRepr::from_attrs(&ctx.ast.attrs)?;
127
128 match repr {
131 Repr::Compound(Spanned { t: CompoundRepr::C | CompoundRepr::Primitive(_), span: _ }, _) => {
132 }
133 Repr::Transparent(_) | Repr::Compound(Spanned { t: CompoundRepr::Rust, span: _ }, _) => {
134 return ctx.error_or_skip(
135 Error::new(
136 Span::call_site(),
137 "must have #[repr(C)] or #[repr(Int)] attribute in order to guarantee this type's memory layout",
138 ),
139 );
140 }
141 }
142
143 let zero_variant = match find_zero_variant(enm) {
144 Ok(index) => enm.variants.iter().nth(index).unwrap(),
145 Err(true) => {
147 return ctx.error_or_skip(Error::new_spanned(
148 &ctx.ast,
149 "FromZeros only supported on enums with a variant that has a discriminant of `0`\n\
150 help: This enum has discriminants which are not literal integers. One of those may \
151 define or imply which variant has a discriminant of zero. Use a literal integer to \
152 define or imply the variant with a discriminant of zero.",
153 ));
154 }
155 Err(false) => {
157 return ctx.error_or_skip(Error::new_spanned(
158 &ctx.ast,
159 "FromZeros only supported on enums with a variant that has a discriminant of `0`",
160 ));
161 }
162 };
163
164 let zerocopy_crate = &ctx.zerocopy_crate;
165 let explicit_bounds = zero_variant
166 .fields
167 .iter()
168 .map(|field| {
169 let ty = &field.ty;
170 parse_quote! { #ty: #zerocopy_crate::FromZeros }
171 })
172 .collect::<Vec<WherePredicate>>();
173
174 Ok(ImplBlockBuilder::new(ctx, enm, Trait::FromZeros, FieldBounds::Explicit(explicit_bounds))
175 .build())
176}
177fn derive_from_zeros_union(ctx: &Ctx, unn: &DataUnion) -> TokenStream {
178 let field_type_trait_bounds = FieldBounds::All(&[TraitBound::Slf]);
179 ImplBlockBuilder::new(ctx, unn, Trait::FromZeros, field_type_trait_bounds).build()
180}
181fn derive_from_bytes_struct(ctx: &Ctx, strct: &DataStruct) -> TokenStream {
182 ImplBlockBuilder::new(ctx, strct, Trait::FromBytes, FieldBounds::ALL_SELF).build()
183}
184fn derive_from_bytes_enum(ctx: &Ctx, enm: &DataEnum) -> Result<TokenStream, Error> {
185 let repr = EnumRepr::from_attrs(&ctx.ast.attrs)?;
186
187 let variants_required = 1usize << enum_size_from_repr(&repr)?;
188 if enm.variants.len() != variants_required {
189 return ctx.error_or_skip(Error::new_spanned(
190 &ctx.ast,
191 format!(
192 "FromBytes only supported on {} enum with {} variants",
193 repr.repr_type_name(),
194 variants_required
195 ),
196 ));
197 }
198
199 Ok(ImplBlockBuilder::new(ctx, enm, Trait::FromBytes, FieldBounds::ALL_SELF).build())
200}
201fn derive_from_bytes_union(ctx: &Ctx, unn: &DataUnion) -> TokenStream {
202 let field_type_trait_bounds = FieldBounds::All(&[TraitBound::Slf]);
203 ImplBlockBuilder::new(ctx, unn, Trait::FromBytes, field_type_trait_bounds).build()
204}