| 1 | use std::borrow::Cow; |
| 2 | |
| 3 | use proc_macro2::TokenStream; |
| 4 | use quote::{format_ident, quote, ToTokens}; |
| 5 | use syn::{spanned::Spanned, DeriveInput, Error, Field, Ident, Result, Type}; |
| 6 | |
| 7 | use crate::get_fields; |
| 8 | |
| 9 | #[derive(Debug, Default)] |
| 10 | enum ConversionMode { |
| 11 | #[default] |
| 12 | None, |
| 13 | Omit, |
| 14 | Into(Type), |
| 15 | TryInto(Type), |
| 16 | ToMigrationState, |
| 17 | } |
| 18 | |
| 19 | impl ConversionMode { |
| 20 | fn target_type(&self, original_type: &Type) -> TokenStream { |
| 21 | match self { |
| 22 | ConversionMode::Into(ty) | ConversionMode::TryInto(ty) => ty.to_token_stream(), |
| 23 | ConversionMode::ToMigrationState => { |
| 24 | quote! { <#original_type as ToMigrationState>::Migrated } |
| 25 | } |
| 26 | _ => original_type.to_token_stream(), |
| 27 | } |
| 28 | } |
| 29 | } |
| 30 | |
| 31 | #[derive(Debug, Default)] |
| 32 | struct ContainerAttrs { |
| 33 | rename: Option<Ident>, |
| 34 | } |
| 35 | |
| 36 | impl ContainerAttrs { |
| 37 | fn parse_from(&mut self, attrs: &[syn::Attribute]) -> Result<()> { |
| 38 | use attrs::{set, with, Attrs}; |
| 39 | Attrs::new() |
| 40 | .once("rename", with::eq(set::parse(&mut self.rename))) |
| 41 | .parse_attrs("migration_state", attrs)?; |
| 42 | Ok(()) |
| 43 | } |
| 44 | |
| 45 | fn parse(attrs: &[syn::Attribute]) -> Result<Self> { |
| 46 | let mut container_attrs = Self::default(); |
| 47 | container_attrs.parse_from(attrs)?; |
| 48 | Ok(container_attrs) |
| 49 | } |
| 50 | } |
| 51 | |
| 52 | #[derive(Debug, Default)] |
| 53 | struct FieldAttrs { |
| 54 | conversion: ConversionMode, |
| 55 | clone: bool, |
| 56 | } |
| 57 | |
| 58 | impl FieldAttrs { |
| 59 | fn parse_from(&mut self, attrs: &[syn::Attribute]) -> Result<()> { |
| 60 | let mut omit_flag = false; |
| 61 | let mut into_type: Option<Type> = None; |
| 62 | let mut try_into_type: Option<Type> = None; |
| 63 | |
| 64 | use attrs::{set, with, Attrs}; |
| 65 | Attrs::new() |
| 66 | .once("omit", set::flag(&mut omit_flag)) |
| 67 | .once("into", with::paren(set::parse(&mut into_type))) |
| 68 | .once("try_into", with::paren(set::parse(&mut try_into_type))) |
| 69 | .once("clone", set::flag(&mut self.clone)) |
| 70 | .parse_attrs("migration_state", attrs)?; |
| 71 | |
| 72 | self.conversion = match (omit_flag, into_type, try_into_type, self.clone) { |
| 73 | // Valid combinations of attributes first... |
| 74 | (true, None, None, false) => ConversionMode::Omit, |
| 75 | (false, Some(ty), None, _) => ConversionMode::Into(ty), |
| 76 | (false, None, Some(ty), _) => ConversionMode::TryInto(ty), |
| 77 | (false, None, None, true) => ConversionMode::None, // clone without conversion |
| 78 | (false, None, None, false) => ConversionMode::ToMigrationState, // default behavior |
| 79 | |
| 80 | // ... then the error cases |
| 81 | (true, _, _, _) => { |
| 82 | return Err(Error::new( |
| 83 | attrs[0].span(), |
| 84 | "ToMigrationState: omit cannot be used with other attributes", |
| 85 | )); |
| 86 | } |
| 87 | (_, Some(_), Some(_), _) => { |
| 88 | return Err(Error::new( |
| 89 | attrs[0].span(), |
| 90 | "ToMigrationState: into and try_into attributes cannot be used together", |
| 91 | )); |
| 92 | } |
| 93 | }; |
| 94 | |
| 95 | Ok(()) |
| 96 | } |
| 97 | |
| 98 | fn parse(attrs: &[syn::Attribute]) -> Result<Self> { |
| 99 | let mut field_attrs = Self::default(); |
| 100 | field_attrs.parse_from(attrs)?; |
| 101 | Ok(field_attrs) |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | #[derive(Debug)] |
| 106 | struct MigrationStateField { |
| 107 | name: Ident, |
| 108 | original_type: Type, |
| 109 | attrs: FieldAttrs, |
| 110 | } |
| 111 | |
| 112 | impl MigrationStateField { |
| 113 | fn maybe_clone(&self, mut value: TokenStream) -> TokenStream { |
| 114 | if self.attrs.clone { |
| 115 | value = quote! { #value.clone() }; |
| 116 | } |
| 117 | value |
| 118 | } |
| 119 | |
| 120 | fn generate_migration_state_field(&self) -> TokenStream { |
| 121 | let name = &self.name; |
| 122 | let field_type = self.attrs.conversion.target_type(&self.original_type); |
| 123 | |
| 124 | quote! { |
| 125 | pub #name: #field_type, |
| 126 | } |
| 127 | } |
| 128 | |
| 129 | fn generate_snapshot_field(&self) -> TokenStream { |
| 130 | let name = &self.name; |
| 131 | let value = self.maybe_clone(quote! { self.#name }); |
| 132 | |
| 133 | match &self.attrs.conversion { |
| 134 | ConversionMode::Omit => { |
| 135 | unreachable!("Omitted fields are filtered out during processing") |
| 136 | } |
| 137 | ConversionMode::None => quote! { |
| 138 | target.#name = #value; |
| 139 | }, |
| 140 | ConversionMode::Into(_) => quote! { |
| 141 | target.#name = #value.into(); |
| 142 | }, |
| 143 | ConversionMode::TryInto(_) => quote! { |
| 144 | target.#name = #value.try_into().map_err(|_| migration::InvalidError)?; |
| 145 | }, |
| 146 | ConversionMode::ToMigrationState => quote! { |
| 147 | self.#name.snapshot_migration_state(&mut target.#name)?; |
| 148 | }, |
| 149 | } |
| 150 | } |
| 151 | |
| 152 | fn generate_restore_field(&self) -> TokenStream { |
| 153 | let name = &self.name; |
| 154 | |
| 155 | match &self.attrs.conversion { |
| 156 | ConversionMode::Omit => { |
| 157 | unreachable!("Omitted fields are filtered out during processing") |
| 158 | } |
| 159 | ConversionMode::None => quote! { |
| 160 | self.#name = #name; |
| 161 | }, |
| 162 | ConversionMode::Into(_) => quote! { |
| 163 | self.#name = #name.into(); |
| 164 | }, |
| 165 | ConversionMode::TryInto(_) => quote! { |
| 166 | self.#name = #name.try_into().map_err(|_| migration::InvalidError)?; |
| 167 | }, |
| 168 | ConversionMode::ToMigrationState => quote! { |
| 169 | self.#name.restore_migrated_state_mut(#name, _version_id)?; |
| 170 | }, |
| 171 | } |
| 172 | } |
| 173 | } |
| 174 | |
| 175 | #[derive(Debug)] |
| 176 | pub struct MigrationStateDerive { |
| 177 | input: DeriveInput, |
| 178 | fields: Vec<MigrationStateField>, |
| 179 | container_attrs: ContainerAttrs, |
| 180 | } |
| 181 | |
| 182 | impl MigrationStateDerive { |
| 183 | fn parse(input: DeriveInput) -> Result<Self> { |
| 184 | let container_attrs = ContainerAttrs::parse(&input.attrs)?; |
| 185 | let fields = get_fields(&input, "ToMigrationState")?; |
| 186 | let fields = Self::process_fields(fields)?; |
| 187 | |
| 188 | Ok(Self { |
| 189 | input, |
| 190 | fields, |
| 191 | container_attrs, |
| 192 | }) |
| 193 | } |
| 194 | |
| 195 | fn process_fields( |
| 196 | fields: &syn::punctuated::Punctuated<Field, syn::token::Comma>, |
| 197 | ) -> Result<Vec<MigrationStateField>> { |
| 198 | let processed = fields |
| 199 | .iter() |
| 200 | .map(|field| { |
| 201 | let attrs = FieldAttrs::parse(&field.attrs)?; |
| 202 | Ok((field, attrs)) |
| 203 | }) |
| 204 | .collect::<Result<Vec<_>>>()? |
| 205 | .into_iter() |
| 206 | .filter(|(_, attrs)| !matches!(attrs.conversion, ConversionMode::Omit)) |
| 207 | .map(|(field, attrs)| MigrationStateField { |
| 208 | name: field.ident.as_ref().unwrap().clone(), |
| 209 | original_type: field.ty.clone(), |
| 210 | attrs, |
| 211 | }) |
| 212 | .collect(); |
| 213 | |
| 214 | Ok(processed) |
| 215 | } |
| 216 | |
| 217 | fn migration_state_name(&self) -> Cow<'_, Ident> { |
| 218 | match &self.container_attrs.rename { |
| 219 | Some(rename) => Cow::Borrowed(rename), |
| 220 | None => Cow::Owned(format_ident!("{}Migration", &self.input.ident)), |
| 221 | } |
| 222 | } |
| 223 | |
| 224 | fn generate_migration_state_struct(&self) -> TokenStream { |
| 225 | let name = self.migration_state_name(); |
| 226 | let fields = self |
| 227 | .fields |
| 228 | .iter() |
| 229 | .map(MigrationStateField::generate_migration_state_field); |
| 230 | |
| 231 | quote! { |
| 232 | #[derive(Default)] |
| 233 | pub struct #name { |
| 234 | #(#fields)* |
| 235 | } |
| 236 | } |
| 237 | } |
| 238 | |
| 239 | fn generate_snapshot_migration_state(&self) -> TokenStream { |
| 240 | let fields = self |
| 241 | .fields |
| 242 | .iter() |
| 243 | .map(MigrationStateField::generate_snapshot_field); |
| 244 | |
| 245 | quote! { |
| 246 | fn snapshot_migration_state(&self, target: &mut Self::Migrated) -> Result<(), migration::InvalidError> { |
| 247 | #(#fields)* |
| 248 | Ok(()) |
| 249 | } |
| 250 | } |
| 251 | } |
| 252 | |
| 253 | fn generate_restore_migrated_state(&self) -> TokenStream { |
| 254 | let names: Vec<_> = self.fields.iter().map(|f| &f.name).collect(); |
| 255 | let fields = self |
| 256 | .fields |
| 257 | .iter() |
| 258 | .map(MigrationStateField::generate_restore_field); |
| 259 | |
| 260 | // version_id could be used or not depending on conversion attributes |
| 261 | quote! { |
| 262 | #[allow(clippy::used_underscore_binding)] |
| 263 | fn restore_migrated_state_mut(&mut self, source: Self::Migrated, _version_id: u8) -> Result<(), migration::InvalidError> { |
| 264 | let Self::Migrated { #(#names),* } = source; |
| 265 | #(#fields)* |
| 266 | Ok(()) |
| 267 | } |
| 268 | } |
| 269 | } |
| 270 | |
| 271 | fn generate(&self) -> TokenStream { |
| 272 | let struct_name = &self.input.ident; |
| 273 | let generics = &self.input.generics; |
| 274 | |
| 275 | let (impl_generics, ty_generics, where_clause) = generics.split_for_impl(); |
| 276 | let name = self.migration_state_name(); |
| 277 | let migration_state_struct = self.generate_migration_state_struct(); |
| 278 | let snapshot_impl = self.generate_snapshot_migration_state(); |
| 279 | let restore_impl = self.generate_restore_migrated_state(); |
| 280 | |
| 281 | quote! { |
| 282 | #migration_state_struct |
| 283 | |
| 284 | impl #impl_generics ToMigrationState for #struct_name #ty_generics #where_clause { |
| 285 | type Migrated = #name; |
| 286 | |
| 287 | #snapshot_impl |
| 288 | |
| 289 | #restore_impl |
| 290 | } |
| 291 | } |
| 292 | } |
| 293 | |
| 294 | pub fn expand(input: DeriveInput) -> Result<TokenStream> { |
| 295 | let tokens = Self::parse(input)?.generate(); |
| 296 | Ok(tokens) |
| 297 | } |
| 298 | } |