1
Fork 0

fix usage of autodiff macro with inner functions

- fix errors caused by the move of `ast::Item::ident` (see #138740)
- move the logic of getting `sig`, `vis`, and `ident` from two seperate
  `match` statements into one (less repetition especially with the
  nested `match`)
This commit is contained in:
HaeNoe 2025-03-09 22:55:07 +01:00
parent 175dcc7773
commit 13bf79cd87
No known key found for this signature in database

View file

@ -199,27 +199,46 @@ mod llvm_enzyme {
return vec![item];
}
let dcx = ecx.sess.dcx();
// first get the annotable item:
let (primal, sig, is_impl): (Ident, FnSig, bool) = match &item {
// first get information about the annotable item:
let (sig, vis, primal) = match &item {
Annotatable::Item(iitem) => {
let (ident, sig) = match &iitem.kind {
ItemKind::Fn(box ast::Fn { ident, sig, .. }) => (ident, sig),
let (sig, ident) = match &iitem.kind {
ItemKind::Fn(box ast::Fn { sig, ident, .. }) => (sig, ident),
_ => {
dcx.emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
return vec![item];
}
};
(*ident, sig.clone(), false)
(sig.clone(), iitem.vis.clone(), ident.clone())
}
Annotatable::AssocItem(assoc_item, Impl { of_trait: false }) => {
let (ident, sig) = match &assoc_item.kind {
ast::AssocItemKind::Fn(box ast::Fn { ident, sig, .. }) => (ident, sig),
let (sig, ident) = match &assoc_item.kind {
ast::AssocItemKind::Fn(box ast::Fn { sig, ident, .. }) => (sig, ident),
_ => {
dcx.emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
return vec![item];
}
};
(*ident, sig.clone(), true)
(sig.clone(), assoc_item.vis.clone(), ident.clone())
}
Annotatable::Stmt(stmt) => {
let (sig, vis, ident) = match &stmt.kind {
ast::StmtKind::Item(iitem) => match &iitem.kind {
ast::ItemKind::Fn(box ast::Fn { sig, ident, .. }) => {
(sig.clone(), iitem.vis.clone(), ident.clone())
}
_ => {
dcx.emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
return vec![item];
}
},
_ => {
dcx.emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
return vec![item];
}
};
(sig, vis, ident)
}
_ => {
dcx.emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
@ -238,15 +257,6 @@ mod llvm_enzyme {
let has_ret = has_ret(&sig.decl.output);
let sig_span = ecx.with_call_site_ctxt(sig.span);
let vis = match &item {
Annotatable::Item(iitem) => iitem.vis.clone(),
Annotatable::AssocItem(assoc_item, _) => assoc_item.vis.clone(),
_ => {
dcx.emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
return vec![item];
}
};
// create TokenStream from vec elemtents:
// meta_item doesn't have a .tokens field
let mut ts: Vec<TokenTree> = vec![];
@ -379,6 +389,22 @@ mod llvm_enzyme {
}
Annotatable::AssocItem(assoc_item.clone(), i)
}
Annotatable::Stmt(ref mut stmt) => {
match stmt.kind {
ast::StmtKind::Item(ref mut iitem) => {
if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &attr.kind)) {
iitem.attrs.push(attr);
}
if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind))
{
iitem.attrs.push(inline_never.clone());
}
}
_ => unreachable!("stmt kind checked previously"),
};
Annotatable::Stmt(stmt.clone())
}
_ => {
unreachable!("annotatable kind checked previously")
}
@ -389,22 +415,40 @@ mod llvm_enzyme {
delim: rustc_ast::token::Delimiter::Parenthesis,
tokens: ts,
});
let d_attr = outer_normal_attr(&rustc_ad_attr, new_id, span);
let d_annotatable = if is_impl {
let assoc_item: AssocItemKind = ast::AssocItemKind::Fn(asdf);
let d_fn = P(ast::AssocItem {
attrs: thin_vec![d_attr, inline_never],
id: ast::DUMMY_NODE_ID,
span,
vis,
kind: assoc_item,
tokens: None,
});
Annotatable::AssocItem(d_fn, Impl { of_trait: false })
} else {
let mut d_fn = ecx.item(span, thin_vec![d_attr, inline_never], ItemKind::Fn(asdf));
d_fn.vis = vis;
Annotatable::Item(d_fn)
let d_annotatable = match &item {
Annotatable::AssocItem(_, _) => {
let assoc_item: AssocItemKind = ast::AssocItemKind::Fn(asdf);
let d_fn = P(ast::AssocItem {
attrs: thin_vec![d_attr, inline_never],
id: ast::DUMMY_NODE_ID,
span,
vis,
kind: assoc_item,
tokens: None,
});
Annotatable::AssocItem(d_fn, Impl { of_trait: false })
}
Annotatable::Item(_) => {
let mut d_fn = ecx.item(span, thin_vec![d_attr, inline_never], ItemKind::Fn(asdf));
d_fn.vis = vis;
Annotatable::Item(d_fn)
}
Annotatable::Stmt(_) => {
let mut d_fn = ecx.item(span, thin_vec![d_attr, inline_never], ItemKind::Fn(asdf));
d_fn.vis = vis;
Annotatable::Stmt(P(ast::Stmt {
id: ast::DUMMY_NODE_ID,
kind: ast::StmtKind::Item(d_fn),
span,
}))
}
_ => {
unreachable!("item kind checked previously")
}
};
return vec![orig_annotatable, d_annotatable];