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() {
len @ 1 | len @ 2 => len,
_ => 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);
}
struct FnVisitor {
res: Result<()>,
mutability: Mutability,
}
impl FnVisitor {
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) {
}
}
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 {
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) => {
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)
}
}
}
}