Something went wrong. Try again.
forked from verus-lang/verus
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288use crate::cfg::{self, DocCfg};use crate::{file, full, gen};use anyhow::Result;use proc_macro2::{Ident, Span, TokenStream};use quote::{format_ident, quote};use syn::Index;use syn_codegen::{Data, Definitions, Features, Node, Type};
const FOLD_SRC: &str = "src/gen/fold.rs";
fn method_name(item: &str) -> Ident { let ident = gen::under_name(item); format_ident!("fold_{}", ident)}
fn simple_fold(item: &str, name: &TokenStream) -> TokenStream { let method = method_name(item); quote! { f.#method(#name) }}
fn visit( ty: &Type, features: &Features, defs: &Definitions, name: &TokenStream,) -> Option<TokenStream> { match ty { Type::Box(t) => { let res = visit(t, features, defs, "e!(*#name))?; Some(quote! { Box::new(#res) }) } Type::Vec(t) => { let Type::Syn(t) = &**t else { unimplemented!() }; if t == "Attribute" { Some(quote! { f.fold_attributes(#name) }) } else { let method = method_name(t); Some(quote! { fold_vec(#name, f, F::#method) }) } } Type::Punctuated(p) => { let t = &*p.element; let Type::Syn(t) = t else { unimplemented!() }; let method = method_name(t); Some(quote! { crate::punctuated::fold(#name, f, F::#method) }) } Type::Option(t) => { let it = quote!(it); let val = visit(t, features, defs, &it)?; Some(quote! { (#name).map(|it| #val) }) } Type::Tuple(t) => { let mut code = TokenStream::new(); for (i, elem) in t.iter().enumerate() { let i = Index::from(i); let it = quote!((#name).#i); let val = visit(elem, features, defs, &it).unwrap_or(it); code.extend(val); code.extend(quote!(,)); } Some(quote! { (#code) }) } Type::Syn(t) => { fn requires_full(features: &Features) -> bool { features.any.contains("full") && features.any.len() == 1 } let mut res = simple_fold(t, name); let target = defs.types.iter().find(|ty| ty.ident == *t).unwrap(); if requires_full(&target.features) && !requires_full(features) { res = quote!(full!(#res)); } Some(res) } Type::Ext(t) if gen::TERMINAL_TYPES.contains(&&t[..]) => Some(simple_fold(t, name)), Type::Ext(_) | Type::Std(_) | Type::Token(_) | Type::Group(_) => None, }}
fn node(traits: &mut TokenStream, impls: &mut TokenStream, s: &Node, defs: &Definitions) { let under_name = gen::under_name(&s.ident); let ident = Ident::new(&s.ident, Span::call_site()); let ty = if gen::TERMINAL_TYPES.contains(&s.ident.as_str()) { quote!(proc_macro2::#ident) } else { quote!(crate::#ident) }; let fold_fn = format_ident!("fold_{}", under_name);
let mut fold_impl = TokenStream::new();
match &s.data { Data::Enum(variants) => { let mut fold_variants = TokenStream::new();
for (variant, fields) in variants { let variant_ident = Ident::new(variant, Span::call_site());
if fields.is_empty() { fold_variants.extend(quote! { #ty::#variant_ident => { #ty::#variant_ident } }); } else { let mut bind_fold_fields = TokenStream::new(); let mut fold_fields = TokenStream::new();
for (idx, ty) in fields.iter().enumerate() { let binding = format_ident!("_binding_{}", idx);
bind_fold_fields.extend(quote! { #binding, });
let owned_binding = quote!(#binding);
fold_fields.extend( visit(ty, &s.features, defs, &owned_binding).unwrap_or(owned_binding), );
fold_fields.extend(quote!(,)); }
fold_variants.extend(quote! { #ty::#variant_ident(#bind_fold_fields) => { #ty::#variant_ident( #fold_fields ) } }); } }
fold_impl.extend(quote! { match node { #fold_variants } }); } Data::Struct(fields) => { let mut fold_fields = TokenStream::new();
for (field, ty) in fields { let id = Ident::new(field, Span::call_site()); let ref_toks = quote!(node.#id);
let fold = visit(ty, &s.features, defs, &ref_toks).unwrap_or(ref_toks);
fold_fields.extend(quote! { #id: #fold, }); }
if fields.is_empty() { if s.ident == "Ident" { fold_impl.extend(quote! { let mut node = node; let span = f.fold_span(node.span()); node.set_span(span); }); } fold_impl.extend(quote! { node }); } else { fold_impl.extend(quote! { #ty { #fold_fields } }); } } Data::Private => { if s.ident == "Ident" { fold_impl.extend(quote! { let mut node = node; let span = f.fold_span(node.span()); node.set_span(span); }); } fold_impl.extend(quote! { node }); } }
let fold_span_only = s.data == Data::Private && !gen::TERMINAL_TYPES.contains(&s.ident.as_str()); if fold_span_only { fold_impl = quote! { let span = f.fold_span(node.span()); let mut node = node; node.set_span(span); node }; }
let traits_body = if s.ident == "Span" || s.ident == "TokenStream" { quote!(i) } else { quote!(#fold_fn(self, i)) };
traits.extend(quote! { fn #fold_fn(&mut self, i: #ty) -> #ty { #traits_body } });
if s.ident != "TokenStream" { impls.extend(quote! { pub fn #fold_fn<F>(f: &mut F, node: #ty) -> #ty where F: Fold + ?Sized, { #fold_impl } }); }
if s.ident == "Attribute" { let features = cfg::features(&s.features, DocCfg::Ordinary); traits.extend(quote! { #features fn fold_attributes(&mut self, i: Vec<crate::Attribute>) -> Vec<crate::Attribute> { fold_vec(i, self, Self::fold_attribute) } }); }}
pub fn generate(defs: &Definitions) -> Result<()> { let (traits, impls) = gen::traverse(defs, node); let full_macro = full::get_macro(); file::write( FOLD_SRC, quote! { // Unreachable code is generated sometimes without the full feature. #![allow(unreachable_code, unused_variables)] #![allow( clippy::match_wildcard_for_single_variants, clippy::needless_match, clippy::needless_pass_by_ref_mut, )]
#[cfg(any(feature = "derive", feature = "full"))] use alloc::boxed::Box; #[cfg(any(feature = "derive", feature = "full"))] use alloc::vec::Vec;
#full_macro
/// Syntax tree traversal to transform the nodes of an owned syntax tree. /// /// See the [module documentation] for details. /// /// [module documentation]: self pub trait Fold { #traits }
#impls
#[cfg(any(feature = "derive", feature = "full"))] fn fold_vec<T, V, F>(vec: Vec<T>, fold: &mut V, mut f: F) -> Vec<T> where V: ?Sized, F: FnMut(&mut V, T) -> T, { vec.into_iter().map(|it| f(fold, it)).collect() } }, )?; Ok(())}