derive_insert_impl/
lib.rs

1// Copyright 2021 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use ::proc_macro::TokenStream;
16use ::proc_macro2::TokenStream as TokenStream2;
17use ::quote::quote;
18use ::syn::{parse, Data, DeriveInput, Field, Fields, Ident, Variant};
19
20#[proc_macro_derive(GetOrInsert)]
21pub fn derive_get_or_insert(input: TokenStream) -> TokenStream {
22    implementation(input).unwrap().into()
23}
24
25fn implementation(input: TokenStream) -> Result<TokenStream2, String> {
26    let ast: DeriveInput = parse(input).unwrap();
27
28    let data = match ast.data {
29        Data::Enum(data) => Ok(data),
30        _ => Err("GetOrInsert can only be derived for enums"),
31    }?;
32
33    let enum_ident: Ident = ast.ident.clone();
34
35    let variants_data = data.variants.iter().collect::<Vec<_>>();
36
37    let variants_to_tokens = variants_data
38        .iter()
39        .filter_map(|v| {
40            let ident = v.ident.clone();
41            VariantData::try_from_syn_variant(v, &enum_ident)
42                .map_err(|_| {
43                    format!(
44                        "Error in {enum_ident}::{ident}: \
45                        GetOrInsert only supports a tuple-like, single item enum variants."
46                    )
47                })
48                .transpose()
49        })
50        .collect::<Result<Vec<_>, _>>()?;
51
52    let variant_tokens = variants_to_tokens.iter().map(VariantData::to_tokens);
53
54    Ok(quote! {
55        #(#variant_tokens)*
56    })
57}
58
59struct VariantData {
60    ident: Ident,
61    enum_ident: Ident,
62    first_unnamed_field: Field,
63}
64
65impl VariantData {
66    fn try_from_syn_variant(variant: &Variant, enum_ident: &Ident) -> Result<Option<Self>, ()> {
67        let ident = variant.ident.clone();
68        let enum_ident = enum_ident.clone();
69        let first_unnamed_field = match &variant.fields {
70            Fields::Unit => return Ok(None),
71            Fields::Named(_) => return Err(()),
72            Fields::Unnamed(ref fields_unnamed) => {
73                let mut items_iter = fields_unnamed.unnamed.iter();
74                let Some(first_field) = items_iter.next() else {
75                    return Err(());
76                };
77                if items_iter.next().is_some() {
78                    return Err(());
79                }
80                first_field.clone()
81            }
82        };
83        Ok(Some(VariantData {
84            ident,
85            enum_ident,
86            first_unnamed_field,
87        }))
88    }
89
90    fn to_tokens(&self) -> TokenStream2 {
91        let ident = &self.ident;
92        let enum_ident = &self.enum_ident;
93        let first_field = &self.first_unnamed_field;
94        quote! {
95            impl ::derive_insert::GetOrInsert<#first_field> for #enum_ident {
96                fn insert(&mut self, value: #first_field) -> &mut #first_field {
97                    *self = #enum_ident::#ident(value);
98                    match self {
99                        #enum_ident::#ident(inner) => inner,
100                        _ => unreachable!(),
101                    }
102                }
103                fn get_or_insert_with<F: FnOnce() -> #first_field>(&mut self, f: F) -> &mut #first_field {
104                    match self {
105                        #enum_ident::#ident(inner) => inner,
106                        _ => self.insert(f()),
107                    }
108                }
109            }
110        }
111    }
112}
derive_insert_impl/lib.rs

derive_insert_impl/
lib.rs
