refactor: Simplify attribute handling in rule_code_prefix

if_all_same(codes.values().cloned()).unwrap_or_default()

was quite unreadable because it wasn't obvious that codes.values() are
the prefixes. It's better to introduce another Map rather than having
Maps within Maps.
This commit is contained in:
Martin Fischer 2023-02-10 09:15:48 +01:00 committed by Charlie Marsh
parent fa191cceeb
commit 9011456aa1

View file

@ -1,4 +1,4 @@
use std::collections::BTreeMap; use std::collections::{BTreeMap, BTreeSet};
use proc_macro2::Span; use proc_macro2::Span;
use quote::quote; use quote::quote;
@ -11,10 +11,10 @@ pub fn expand<'a>(
variant_name: impl Fn(&str) -> &'a Ident, variant_name: impl Fn(&str) -> &'a Ident,
) -> proc_macro2::TokenStream { ) -> proc_macro2::TokenStream {
// Build up a map from prefix to matching RuleCodes. // Build up a map from prefix to matching RuleCodes.
let mut prefix_to_codes: BTreeMap<String, BTreeMap<String, Vec<Attribute>>> = let mut prefix_to_codes: BTreeMap<String, BTreeSet<String>> = BTreeMap::default();
BTreeMap::default(); let mut attributes: BTreeMap<String, &[Attribute]> = BTreeMap::default();
let mut pl_codes = BTreeMap::new(); let mut pl_codes = BTreeSet::new();
for (variant, attr) in variants { for (variant, attr) in variants {
let code_str = variant.to_string(); let code_str = variant.to_string();
@ -28,26 +28,32 @@ pub fn expand<'a>(
prefix_to_codes prefix_to_codes
.entry(prefix) .entry(prefix)
.or_default() .or_default()
.entry(code_str.clone()) .insert(code_str.clone());
.or_insert_with(|| attr.clone());
} }
if code_str.starts_with("PL") { if code_str.starts_with("PL") {
pl_codes.insert(code_str, attr.clone()); pl_codes.insert(code_str.clone());
} }
attributes.insert(code_str, attr);
} }
prefix_to_codes.insert("PL".to_string(), pl_codes); prefix_to_codes.insert("PL".to_string(), pl_codes);
let prefix_variants = prefix_to_codes.iter().map(|(prefix, codes)| { let prefix_variants = prefix_to_codes.iter().map(|(prefix, codes)| {
let prefix = Ident::new(prefix, Span::call_site()); let prefix = Ident::new(prefix, Span::call_site());
let attr = if_all_same(codes.values().cloned()).unwrap_or_default(); let attrs = attributes_for_prefix(codes, &attributes);
quote! { quote! {
#(#attr)* #attrs
#prefix #prefix
} }
}); });
let prefix_impl = generate_impls(rule_type, prefix_ident, &prefix_to_codes, variant_name); let prefix_impl = generate_impls(
rule_type,
prefix_ident,
&prefix_to_codes,
variant_name,
&attributes,
);
quote! { quote! {
#[derive( #[derive(
@ -76,21 +82,23 @@ pub fn expand<'a>(
fn generate_impls<'a>( fn generate_impls<'a>(
rule_type: &Ident, rule_type: &Ident,
prefix_ident: &Ident, prefix_ident: &Ident,
prefix_to_codes: &BTreeMap<String, BTreeMap<String, Vec<Attribute>>>, prefix_to_codes: &BTreeMap<String, BTreeSet<String>>,
variant_name: impl Fn(&str) -> &'a Ident, variant_name: impl Fn(&str) -> &'a Ident,
attributes: &BTreeMap<String, &[Attribute]>,
) -> proc_macro2::TokenStream { ) -> proc_macro2::TokenStream {
let into_iter_match_arms = prefix_to_codes.iter().map(|(prefix_str, codes)| { let into_iter_match_arms = prefix_to_codes.iter().map(|(prefix_str, codes)| {
let prefix = Ident::new(prefix_str, Span::call_site()); let prefix = Ident::new(prefix_str, Span::call_site());
let attr = if_all_same(codes.values().cloned()).unwrap_or_default(); let attrs = attributes_for_prefix(codes, attributes);
let codes = codes.iter().map(|(code, attr)| { let codes = codes.iter().map(|code| {
let rule_variant = variant_name(code); let rule_variant = variant_name(code);
let attrs = attributes[code];
quote! { quote! {
#(#attr)* #(#attrs)*
#rule_type::#rule_variant #rule_type::#rule_variant
} }
}); });
quote! { quote! {
#(#attr)* #attrs
#prefix_ident::#prefix => vec![#(#codes),*].into_iter(), #prefix_ident::#prefix => vec![#(#codes),*].into_iter(),
} }
}); });
@ -110,9 +118,9 @@ fn generate_impls<'a>(
5 => quote! { Specificity::Code5Chars }, 5 => quote! { Specificity::Code5Chars },
_ => panic!("Invalid prefix: {prefix}"), _ => panic!("Invalid prefix: {prefix}"),
}; };
let attr = if_all_same(codes.values().cloned()).unwrap_or_default(); let attrs = attributes_for_prefix(codes, attributes);
quote! { quote! {
#(#attr)* #attrs
#prefix_ident::#prefix => #suffix_len, #prefix_ident::#prefix => #suffix_len,
} }
}); });
@ -143,6 +151,16 @@ fn generate_impls<'a>(
} }
} }
fn attributes_for_prefix(
codes: &BTreeSet<String>,
attributes: &BTreeMap<String, &[Attribute]>,
) -> proc_macro2::TokenStream {
match if_all_same(codes.iter().map(|code| attributes[code])) {
Some(attr) => quote!(#(#attr)*),
None => quote!(),
}
}
/// If all values in an iterator are the same, return that value. Otherwise, /// If all values in an iterator are the same, return that value. Otherwise,
/// return `None`. /// return `None`.
fn if_all_same<T: PartialEq>(iter: impl Iterator<Item = T>) -> Option<T> { fn if_all_same<T: PartialEq>(iter: impl Iterator<Item = T>) -> Option<T> {