| use proc_macro2::{Span, TokenStream}; |
| use quote::ToTokens; |
| use syn::{ |
| visit_mut::{self, VisitMut}, |
| *, |
| }; |
| |
| use crate::utils::*; |
| |
| pub(crate) fn attribute(args: &TokenStream, input: Stmt, mutability: Mutability) -> TokenStream { |
| parse_as_empty(args) |
| .and_then(|()| parse(input, mutability)) |
| .unwrap_or_else(|e| e.to_compile_error()) |
| } |
| |
| fn parse(mut stmt: Stmt, mutability: Mutability) -> Result<TokenStream> { |
| match &mut stmt { |
| Stmt::Expr(Expr::Match(expr)) | Stmt::Semi(Expr::Match(expr), _) => { |
| Context::new(mutability).replace_expr_match(expr) |
| } |
| Stmt::Local(local) => Context::new(mutability).replace_local(local)?, |
| Stmt::Item(Item::Fn(item)) => replace_item_fn(item, mutability)?, |
| Stmt::Item(Item::Impl(item)) => replace_item_impl(item, mutability), |
| Stmt::Item(Item::Use(item)) => replace_item_use(item, mutability)?, |
| _ => {} |
| } |
| |
| Ok(stmt.into_token_stream()) |
| } |
| |
| struct Context { |
| register: Option<(Ident, usize)>, |
| replaced: bool, |
| mutability: Mutability, |
| } |
| |
| impl Context { |
| fn new(mutability: Mutability) -> Self { |
| Self { register: None, replaced: false, mutability } |
| } |
| |
| fn update(&mut self, ident: &Ident, len: usize) { |
| if self.register.is_none() { |
| self.register = Some((ident.clone(), len)); |
| } |
| } |
| |
| fn compare_paths(&self, ident: &Ident, len: usize) -> bool { |
| match &self.register { |
| Some((i, l)) => *l == len && ident == i, |
| None => false, |
| } |
| } |
| |
| fn replace_local(&mut self, local: &mut Local) -> Result<()> { |
| if let Some(Expr::Match(expr)) = local.init.as_mut().map(|(_, expr)| &mut **expr) { |
| self.replace_expr_match(expr); |
| } |
| |
| if self.replaced { |
| if is_replaceable(&local.pat, false) { |
| return Err(error!( |
| local.pat, |
| "Both initializer expression and pattern are replaceable, \ |
| you need to split the initializer expression into separate let bindings \ |
| to avoid ambiguity" |
| )); |
| } |
| } else { |
| self.replace_pat(&mut local.pat, false); |
| } |
| |
| Ok(()) |
| } |
| |
| fn replace_expr_match(&mut self, expr: &mut ExprMatch) { |
| expr.arms.iter_mut().for_each(|Arm { pat, .. }| self.replace_pat(pat, true)) |
| } |
| |
| fn replace_pat(&mut self, pat: &mut Pat, allow_pat_path: bool) { |
| match pat { |
| Pat::Ident(PatIdent { subpat: Some((_, pat)), .. }) |
| | Pat::Reference(PatReference { pat, .. }) |
| | Pat::Box(PatBox { pat, .. }) |
| | Pat::Type(PatType { pat, .. }) => self.replace_pat(pat, allow_pat_path), |
| |
| Pat::Or(PatOr { cases, .. }) => { |
| cases.iter_mut().for_each(|pat| self.replace_pat(pat, allow_pat_path)) |
| } |
| |
| Pat::Struct(PatStruct { path, .. }) | Pat::TupleStruct(PatTupleStruct { path, .. }) => { |
| self.replace_path(path) |
| } |
| Pat::Path(PatPath { qself: None, path, .. }) if allow_pat_path => { |
| self.replace_path(path) |
| } |
| _ => {} |
| } |
| } |
| |
| fn replace_path(&mut self, path: &mut Path) { |
| let len = match path.segments.len() { |
| // 1: struct |
| // 2: enum |
| len @ 1 | len @ 2 => len, |
| // other path |
| _ => return, |
| }; |
| |
| if self.register.is_none() || self.compare_paths(&path.segments[0].ident, len) { |
| self.update(&path.segments[0].ident, len); |
| self.replaced = true; |
| replace_ident(&mut path.segments[0].ident, self.mutability); |
| } |
| } |
| } |
| |
| fn is_replaceable(pat: &Pat, allow_pat_path: bool) -> bool { |
| match pat { |
| Pat::Ident(PatIdent { subpat: Some((_, pat)), .. }) |
| | Pat::Reference(PatReference { pat, .. }) |
| | Pat::Box(PatBox { pat, .. }) |
| | Pat::Type(PatType { pat, .. }) => is_replaceable(pat, allow_pat_path), |
| |
| Pat::Or(PatOr { cases, .. }) => cases.iter().any(|pat| is_replaceable(pat, allow_pat_path)), |
| |
| Pat::Struct(_) | Pat::TupleStruct(_) => true, |
| Pat::Path(PatPath { qself: None, .. }) => allow_pat_path, |
| _ => false, |
| } |
| } |
| |
| fn replace_item_impl(item: &mut ItemImpl, mutability: Mutability) { |
| let PathSegment { ident, arguments } = match &mut *item.self_ty { |
| Type::Path(TypePath { qself: None, path }) => path.segments.last_mut().unwrap(), |
| _ => return, |
| }; |
| |
| replace_ident(ident, mutability); |
| |
| let mut lifetime_name = String::from(DEFAULT_LIFETIME_NAME); |
| determine_lifetime_name(&mut lifetime_name, &item.generics.params); |
| item.items |
| .iter_mut() |
| .filter_map(|i| if let ImplItem::Method(i) = i { Some(i) } else { None }) |
| .for_each(|item| determine_lifetime_name(&mut lifetime_name, &item.sig.generics.params)); |
| let lifetime = Lifetime::new(&lifetime_name, Span::call_site()); |
| |
| insert_lifetime(&mut item.generics, lifetime.clone()); |
| |
| match arguments { |
| PathArguments::None => { |
| *arguments = PathArguments::AngleBracketed(syn::parse_quote!(<#lifetime>)); |
| } |
| PathArguments::AngleBracketed(args) => { |
| args.args.insert(0, syn::parse_quote!(#lifetime)); |
| } |
| PathArguments::Parenthesized(_) => unreachable!(), |
| } |
| } |
| |
| fn replace_item_fn(item: &mut ItemFn, mutability: Mutability) -> Result<()> { |
| let mut visitor = FnVisitor { res: Ok(()), mutability }; |
| visitor.visit_block_mut(&mut item.block); |
| visitor.res |
| } |
| |
| fn replace_item_use(item: &mut ItemUse, mutability: Mutability) -> Result<()> { |
| let mut visitor = UseTreeVisitor { res: Ok(()), mutability }; |
| visitor.visit_item_use_mut(item); |
| visitor.res |
| } |
| |
| fn replace_ident(ident: &mut Ident, mutability: Mutability) { |
| *ident = proj_ident(ident, mutability); |
| } |
| |
| // ================================================================================================= |
| // visitor |
| |
| struct FnVisitor { |
| res: Result<()>, |
| mutability: Mutability, |
| } |
| |
| impl FnVisitor { |
| /// Returns the attribute name. |
| fn name(&self) -> &str { |
| if self.mutability == Mutable { "project" } else { "project_ref" } |
| } |
| |
| fn visit_stmt(&mut self, node: &mut Stmt) -> Result<()> { |
| let attr = match node { |
| Stmt::Expr(Expr::Match(expr)) | Stmt::Semi(Expr::Match(expr), _) => { |
| expr.attrs.find_remove(self.name())? |
| } |
| Stmt::Local(local) => local.attrs.find_remove(self.name())?, |
| _ => return Ok(()), |
| }; |
| if let Some(attr) = attr { |
| parse_as_empty(&attr.tokens)?; |
| match node { |
| Stmt::Expr(Expr::Match(expr)) | Stmt::Semi(Expr::Match(expr), _) => { |
| Context::new(self.mutability).replace_expr_match(expr) |
| } |
| Stmt::Local(local) => Context::new(self.mutability).replace_local(local)?, |
| _ => unreachable!(), |
| } |
| } |
| Ok(()) |
| } |
| } |
| |
| impl VisitMut for FnVisitor { |
| fn visit_stmt_mut(&mut self, node: &mut Stmt) { |
| if self.res.is_err() { |
| return; |
| } |
| |
| visit_mut::visit_stmt_mut(self, node); |
| |
| if let Err(e) = self.visit_stmt(node) { |
| self.res = Err(e) |
| } |
| } |
| |
| fn visit_item_mut(&mut self, _: &mut Item) { |
| // Do not recurse into nested items. |
| } |
| } |
| |
| struct UseTreeVisitor { |
| res: Result<()>, |
| mutability: Mutability, |
| } |
| |
| impl VisitMut for UseTreeVisitor { |
| fn visit_use_tree_mut(&mut self, node: &mut UseTree) { |
| if self.res.is_err() { |
| return; |
| } |
| |
| match node { |
| // Desugar `use tree::<name>` into `tree::__<name>Projection`. |
| UseTree::Name(name) => replace_ident(&mut name.ident, self.mutability), |
| UseTree::Glob(glob) => { |
| self.res = |
| Err(error!(glob, "#[project] attribute may not be used on glob imports")); |
| } |
| UseTree::Rename(rename) => { |
| // TODO: Consider allowing the projected type to be renamed by `#[project] use Foo as Bar`. |
| self.res = |
| Err(error!(rename, "#[project] attribute may not be used on renamed imports")); |
| } |
| node @ UseTree::Path(_) | node @ UseTree::Group(_) => { |
| visit_mut::visit_use_tree_mut(self, node) |
| } |
| } |
| } |
| } |