Implement #[define_opaque] attribute for functions.

Author: oli-obk

Cargo.lock

@@ -4529,6 +4529,7 @@ version = "0.0.0"
 dependencies = [
  "itertools",
  "rustc_abi",
+ "rustc_attr_parsing",
  "rustc_data_structures",
  "rustc_errors",
  "rustc_fluent_macro",

compiler/rustc_ast/src/ast.rs

@@ -3466,6 +3466,7 @@ pub struct Fn {
     pub generics: Generics,
     pub sig: FnSig,
     pub contract: Option<P<FnContract>>,
+    pub define_opaque: Option<ThinVec<(NodeId, Path)>>,
     pub body: Option<P<Block>>,
 }
 
@@ -3763,7 +3764,7 @@ mod size_asserts {
     static_assert_size!(Block, 32);
     static_assert_size!(Expr, 72);
     static_assert_size!(ExprKind, 40);
-    static_assert_size!(Fn, 168);
+    static_assert_size!(Fn, 176);
     static_assert_size!(ForeignItem, 88);
     static_assert_size!(ForeignItemKind, 16);
     static_assert_size!(GenericArg, 24);

compiler/rustc_ast/src/mut_visit.rs

@@ -973,7 +973,14 @@ fn walk_fn<T: MutVisitor>(vis: &mut T, kind: FnKind<'_>) {
             _ctxt,
             _ident,
             _vis,
-            Fn { defaultness, generics, contract, body, sig: FnSig { header, decl, span } },
+            Fn {
+                defaultness,
+                generics,
+                contract,
+                body,
+                sig: FnSig { header, decl, span },
+                define_opaque,
+            },
         ) => {
             // Identifier and visibility are visited as a part of the item.
             visit_defaultness(vis, defaultness);
@@ -987,6 +994,11 @@ fn walk_fn<T: MutVisitor>(vis: &mut T, kind: FnKind<'_>) {
                 vis.visit_block(body);
             }
             vis.visit_span(span);
+
+            for (id, path) in define_opaque.iter_mut().flatten() {
+                vis.visit_id(id);
+                vis.visit_path(path)
+            }
         }
         FnKind::Closure(binder, coroutine_kind, decl, body) => {
             vis.visit_closure_binder(binder);

compiler/rustc_ast/src/visit.rs

@@ -892,14 +892,24 @@ pub fn walk_fn<'a, V: Visitor<'a>>(visitor: &mut V, kind: FnKind<'a>) -> V::Resu
             _ctxt,
             _ident,
             _vis,
-            Fn { defaultness: _, sig: FnSig { header, decl, span: _ }, generics, contract, body },
+            Fn {
+                defaultness: _,
+                sig: FnSig { header, decl, span: _ },
+                generics,
+                contract,
+                body,
+                define_opaque,
+            },
         ) => {
             // Identifier and visibility are visited as a part of the item.
             try_visit!(visitor.visit_fn_header(header));
             try_visit!(visitor.visit_generics(generics));
             try_visit!(visitor.visit_fn_decl(decl));
             visit_opt!(visitor, visit_contract, contract);
             visit_opt!(visitor, visit_block, body);
+            for (id, path) in define_opaque.iter().flatten() {
+                try_visit!(visitor.visit_path(path, *id))
+            }
         }
         FnKind::Closure(binder, coroutine_kind, decl, body) => {
             try_visit!(visitor.visit_closure_binder(binder));
@@ -1203,7 +1213,7 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) -> V
                 FnKind::Closure(binder, coroutine_kind, fn_decl, body),
                 *span,
                 *id
-            ))
+            ));
         }
         ExprKind::Block(block, opt_label) => {
             visit_opt!(visitor, visit_label, opt_label);

compiler/rustc_ast_lowering/src/item.rs

@@ -3,10 +3,9 @@ use rustc_ast::ptr::P;
 use rustc_ast::visit::AssocCtxt;
 use rustc_ast::*;
 use rustc_errors::ErrorGuaranteed;
-use rustc_hir as hir;
-use rustc_hir::PredicateOrigin;
 use rustc_hir::def::{DefKind, Res};
 use rustc_hir::def_id::{CRATE_DEF_ID, LocalDefId};
+use rustc_hir::{self as hir, HirId, PredicateOrigin};
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::span_bug;
 use rustc_middle::ty::{ResolverAstLowering, TyCtxt};
@@ -209,6 +208,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 generics,
                 body,
                 contract,
+                define_opaque,
                 ..
             }) => {
                 self.with_new_scopes(*fn_sig_span, |this| {
@@ -237,6 +237,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                         header: this.lower_fn_header(*header, hir::Safety::Safe, attrs),
                         span: this.lower_span(*fn_sig_span),
                     };
+                    this.lower_define_opaque(hir_id, &define_opaque);
                     hir::ItemKind::Fn { sig, generics, body: body_id, has_body: body.is_some() }
                 })
             }
@@ -779,7 +780,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 );
                 (generics, kind, expr.is_some())
             }
-            AssocItemKind::Fn(box Fn { sig, generics, body: None, .. }) => {
+            AssocItemKind::Fn(box Fn { sig, generics, body: None, define_opaque, .. }) => {
                 // FIXME(contracts): Deny contract here since it won't apply to
                 // any impl method or callees.
                 let names = self.lower_fn_params_to_names(&sig.decl);
@@ -791,9 +792,22 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     sig.header.coroutine_kind,
                     attrs,
                 );
+                if define_opaque.is_some() {
+                    self.dcx().span_err(
+                        i.span,
+                        "only trait methods with default bodies can define opaque types",
+                    );
+                }
                 (generics, hir::TraitItemKind::Fn(sig, hir::TraitFn::Required(names)), false)
             }
-            AssocItemKind::Fn(box Fn { sig, generics, body: Some(body), contract, .. }) => {
+            AssocItemKind::Fn(box Fn {
+                sig,
+                generics,
+                body: Some(body),
+                contract,
+                define_opaque,
+                ..
+            }) => {
                 let body_id = self.lower_maybe_coroutine_body(
                     sig.span,
                     i.span,
@@ -812,6 +826,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     sig.header.coroutine_kind,
                     attrs,
                 );
+                self.lower_define_opaque(hir_id, &define_opaque);
                 (generics, hir::TraitItemKind::Fn(sig, hir::TraitFn::Provided(body_id)), true)
             }
             AssocItemKind::Type(box TyAlias { generics, where_clauses, bounds, ty, .. }) => {
@@ -911,7 +926,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     hir::ImplItemKind::Const(ty, body)
                 },
             ),
-            AssocItemKind::Fn(box Fn { sig, generics, body, contract, .. }) => {
+            AssocItemKind::Fn(box Fn { sig, generics, body, contract, define_opaque, .. }) => {
                 let body_id = self.lower_maybe_coroutine_body(
                     sig.span,
                     i.span,
@@ -930,6 +945,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     sig.header.coroutine_kind,
                     attrs,
                 );
+                self.lower_define_opaque(hir_id, &define_opaque);
 
                 (generics, hir::ImplItemKind::Fn(sig, body_id))
             }
@@ -1657,6 +1673,31 @@ impl<'hir> LoweringContext<'_, 'hir> {
         (lowered_generics, res)
     }
 
+    pub(super) fn lower_define_opaque(
+        &mut self,
+        hir_id: HirId,
+        define_opaque: &Option<ThinVec<(NodeId, Path)>>,
+    ) {
+        assert_eq!(self.define_opaque, None);
+        assert!(hir_id.is_owner());
+        let Some(define_opaque) = define_opaque.as_ref() else {
+            return;
+        };
+        let define_opaque = define_opaque
+            .iter()
+            // TODO: error reporting for non-local items being mentioned and tests that go through these code paths
+            .map(|(id, _path)| {
+                self.resolver
+                    .get_partial_res(*id)
+                    .unwrap()
+                    .expect_full_res()
+                    .def_id()
+                    .expect_local()
+            });
+        let define_opaque = self.arena.alloc_from_iter(define_opaque);
+        self.define_opaque = Some(define_opaque);
+    }
+
     pub(super) fn lower_generic_bound_predicate(
         &mut self,
         ident: Ident,

compiler/rustc_ast_lowering/src/lib.rs

@@ -35,6 +35,7 @@
 #![doc(rust_logo)]
 #![feature(assert_matches)]
 #![feature(box_patterns)]
+#![feature(exact_size_is_empty)]
 #![feature(if_let_guard)]
 #![feature(let_chains)]
 #![feature(rustdoc_internals)]
@@ -97,6 +98,8 @@ struct LoweringContext<'a, 'hir> {
 
     /// Bodies inside the owner being lowered.
     bodies: Vec<(hir::ItemLocalId, &'hir hir::Body<'hir>)>,
+    /// `#[define_opaque]` attributes
+    define_opaque: Option<&'hir [LocalDefId]>,
     /// Attributes inside the owner being lowered.
     attrs: SortedMap<hir::ItemLocalId, &'hir [hir::Attribute]>,
     /// Collect items that were created by lowering the current owner.
@@ -154,6 +157,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
             // HirId handling.
             bodies: Vec::new(),
+            define_opaque: None,
             attrs: SortedMap::default(),
             children: Vec::default(),
             contract_ensures: None,
@@ -546,6 +550,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
         let current_attrs = std::mem::take(&mut self.attrs);
         let current_bodies = std::mem::take(&mut self.bodies);
+        let current_define_opaque = std::mem::take(&mut self.define_opaque);
         let current_ident_and_label_to_local_id =
             std::mem::take(&mut self.ident_and_label_to_local_id);
 
@@ -579,6 +584,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
         self.attrs = current_attrs;
         self.bodies = current_bodies;
+        self.define_opaque = current_define_opaque;
         self.ident_and_label_to_local_id = current_ident_and_label_to_local_id;
 
         #[cfg(debug_assertions)]
@@ -598,6 +604,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
     fn make_owner_info(&mut self, node: hir::OwnerNode<'hir>) -> &'hir hir::OwnerInfo<'hir> {
         let attrs = std::mem::take(&mut self.attrs);
         let mut bodies = std::mem::take(&mut self.bodies);
+        let define_opaque = std::mem::take(&mut self.define_opaque);
         let trait_map = std::mem::take(&mut self.trait_map);
 
         #[cfg(debug_assertions)]
@@ -617,7 +624,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         let num_nodes = self.item_local_id_counter.as_usize();
         let (nodes, parenting) = index::index_hir(self.tcx, node, &bodies, num_nodes);
         let nodes = hir::OwnerNodes { opt_hash_including_bodies, nodes, bodies };
-        let attrs = hir::AttributeMap { map: attrs, opt_hash: attrs_hash };
+        let attrs = hir::AttributeMap { map: attrs, opt_hash: attrs_hash, define_opaque };
 
         self.arena.alloc(hir::OwnerInfo { nodes, parenting, attrs, trait_map })
     }

compiler/rustc_ast_passes/src/ast_validation.rs

@@ -917,7 +917,10 @@ impl<'a> Visitor<'a> for AstValidator<'a> {
                 walk_list!(self, visit_attribute, &item.attrs);
                 return; // Avoid visiting again.
             }
-            ItemKind::Fn(func @ box Fn { defaultness, generics: _, sig, contract: _, body }) => {
+            ItemKind::Fn(
+                func
+                @ box Fn { defaultness, generics: _, sig, contract: _, body, define_opaque: _ },
+            ) => {
                 self.check_defaultness(item.span, *defaultness);
 
                 let is_intrinsic =

compiler/rustc_ast_pretty/src/pprust/state.rs

@@ -24,7 +24,6 @@ use rustc_span::edition::Edition;
 use rustc_span::source_map::{SourceMap, Spanned};
 use rustc_span::symbol::IdentPrinter;
 use rustc_span::{BytePos, CharPos, DUMMY_SP, FileName, Ident, Pos, Span, Symbol, kw, sym};
-use thin_vec::ThinVec;
 
 use crate::pp::Breaks::{Consistent, Inconsistent};
 use crate::pp::{self, Breaks};
@@ -1978,15 +1977,7 @@ impl<'a> State<'a> {
     ) {
         self.ibox(INDENT_UNIT);
         self.print_formal_generic_params(generic_params);
-        let generics = ast::Generics {
-            params: ThinVec::new(),
-            where_clause: ast::WhereClause {
-                has_where_token: false,
-                predicates: ThinVec::new(),
-                span: DUMMY_SP,
-            },
-            span: DUMMY_SP,
-        };
+        let generics = ast::Generics::default();
         let header = ast::FnHeader { safety, ext, ..ast::FnHeader::default() };
         self.print_fn(decl, header, name, &generics);
         self.end();

compiler/rustc_ast_pretty/src/pprust/state/item.rs

@@ -650,7 +650,16 @@ impl<'a> State<'a> {
         attrs: &[ast::Attribute],
         func: &ast::Fn,
     ) {
-        let ast::Fn { defaultness, generics, sig, contract, body } = func;
+        let ast::Fn { defaultness, generics, sig, contract, body, define_opaque } = func;
+
+        if let Some(define_opaque) = define_opaque {
+            for (_, path) in define_opaque {
+                self.word("define opaques from ");
+                self.print_path(path, false, 0);
+                self.word(",");
+            }
+        }
+
         if body.is_some() {
             self.head("");
         }
@@ -698,7 +707,7 @@ impl<'a> State<'a> {
         }
         self.print_generic_params(&generics.params);
         self.print_fn_params_and_ret(decl, false);
-        self.print_where_clause(&generics.where_clause)
+        self.print_where_clause(&generics.where_clause);
     }
 
     pub(crate) fn print_fn_params_and_ret(&mut self, decl: &ast::FnDecl, is_closure: bool) {

compiler/rustc_borrowck/src/type_check/opaque_types.rs

@@ -180,6 +180,7 @@ pub(super) fn take_opaques_and_register_member_constraints<'tcx>(
 /// // Equivalent to:
 /// # mod dummy { use super::*;
 /// type FooReturn<'a, T> = impl Foo<'a>;
+/// #[define_opaque(FooReturn)]
 /// fn foo<'a, T>(x: &'a u32, y: T) -> FooReturn<'a, T> {
 ///   (x, y)
 /// }

compiler/rustc_builtin_macros/src/alloc_error_handler.rs

@@ -88,6 +88,7 @@ fn generate_handler(cx: &ExtCtxt<'_>, handler: Ident, span: Span, sig_span: Span
         generics: Generics::default(),
         contract: None,
         body,
+        define_opaque: None,
     }));
 
     let attrs = thin_vec![cx.attr_word(sym::rustc_std_internal_symbol, span)];

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -247,6 +247,7 @@ mod llvm_enzyme {
             generics: Generics::default(),
             contract: None,
             body: Some(d_body),
+            define_opaque: None,
         });
         let mut rustc_ad_attr =
             P(ast::NormalAttr::from_ident(Ident::with_dummy_span(sym::rustc_autodiff)));