Tweak alloc_error_handler desugaring

Author: estebank

compiler/rustc_builtin_macros/src/alloc_error_handler.rs

@@ -10,7 +10,7 @@ use thin_vec::{thin_vec, ThinVec};
 
 pub fn expand(
     ecx: &mut ExtCtxt<'_>,
-    span: Span,
+    _span: Span,
     meta_item: &ast::MetaItem,
     item: Annotatable,
 ) -> Vec<Annotatable> {
@@ -20,29 +20,29 @@ pub fn expand(
 
     // Allow using `#[alloc_error_handler]` on an item statement
     // FIXME - if we get deref patterns, use them to reduce duplication here
-    let (item, is_stmt, sig_span) =
+    let (item, is_stmt, sig) =
         if let Annotatable::Item(item) = &item
             && let ItemKind::Fn(fn_kind) = &item.kind
         {
-            (item, false, ecx.with_def_site_ctxt(fn_kind.sig.span))
+            (item, false, &fn_kind.sig)
         } else if let Annotatable::Stmt(stmt) = &item
             && let StmtKind::Item(item) = &stmt.kind
             && let ItemKind::Fn(fn_kind) = &item.kind
         {
-            (item, true, ecx.with_def_site_ctxt(fn_kind.sig.span))
+            (item, true, &fn_kind.sig)
         } else {
             ecx.sess.parse_sess.span_diagnostic.span_err(item.span(), "alloc_error_handler must be a function");
             return vec![orig_item];
         };
 
     // Generate a bunch of new items using the AllocFnFactory
-    let span = ecx.with_def_site_ctxt(span);
+    let span = ecx.with_def_site_ctxt(item.span);
 
     // Generate item statements for the allocator methods.
-    let stmts = thin_vec![generate_handler(ecx, item.ident, span, sig_span)];
+    let stmts = thin_vec![generate_handler(ecx, item.ident, span, sig)];
 
     // Generate anonymous constant serving as container for the allocator methods.
-    let const_ty = ecx.ty(sig_span, TyKind::Tup(ThinVec::new()));
+    let const_ty = ecx.ty(ecx.with_def_site_ctxt(sig.span), TyKind::Tup(ThinVec::new()));
     let const_body = ecx.expr_block(ecx.block(span, stmts));
     let const_item = ecx.item_const(span, Ident::new(kw::Underscore, span), const_ty, const_body);
     let const_item = if is_stmt {
@@ -59,27 +59,31 @@ pub fn expand(
 // unsafe fn __rg_oom(size: usize, align: usize) -> ! {
 //     handler(core::alloc::Layout::from_size_align_unchecked(size, align))
 // }
-fn generate_handler(cx: &ExtCtxt<'_>, handler: Ident, span: Span, sig_span: Span) -> Stmt {
+fn generate_handler(cx: &ExtCtxt<'_>, handler: Ident, span: Span, sig: &FnSig) -> Stmt {
     let usize = cx.path_ident(span, Ident::new(sym::usize, span));
     let ty_usize = cx.ty_path(usize);
     let size = Ident::from_str_and_span("size", span);
     let align = Ident::from_str_and_span("align", span);
 
+    let sig_span = cx.with_def_site_ctxt(sig.span);
+    let ret_sp = cx.with_def_site_ctxt(sig.decl.output.span());
+
+    // core::alloc::Layout::from_size_align_unchecked(size, align)
     let layout_new = cx.std_path(&[sym::alloc, sym::Layout, sym::from_size_align_unchecked]);
     let layout_new = cx.expr_path(cx.path(span, layout_new));
     let layout = cx.expr_call(
-        span,
+        cx.with_def_site_ctxt(sig.decl.inputs.get(0).map_or(span, |p| p.span)),
         layout_new,
         thin_vec![cx.expr_ident(span, size), cx.expr_ident(span, align)],
     );
-
-    let call = cx.expr_call_ident(sig_span, handler, thin_vec![layout]);
+    let call =
+        cx.expr(ret_sp, ast::ExprKind::Call(cx.expr_ident(span, handler), thin_vec![layout]));
 
     let never = ast::FnRetTy::Ty(cx.ty(span, TyKind::Never));
     let params = thin_vec![cx.param(span, size, ty_usize.clone()), cx.param(span, align, ty_usize)];
     let decl = cx.fn_decl(params, never);
     let header = FnHeader { unsafety: Unsafe::Yes(span), ..FnHeader::default() };
-    let sig = FnSig { decl, header, span: span };
+    let sig = FnSig { decl, header, span };
 
     let body = Some(cx.block_expr(call));
     let kind = ItemKind::Fn(Box::new(Fn {

tests/ui/alloc-error/alloc-error-handler-bad-signature-1.rs

@@ -6,10 +6,10 @@
 
 use core::alloc::Layout;
 
-#[alloc_error_handler] //~ ERROR mismatched types
-fn oom( //~ ERROR mismatched types
-    info: &Layout,
-) -> ()
+#[alloc_error_handler]
+fn oom(
+    info: &Layout, //~ ERROR mismatched types
+) -> () //~ ERROR mismatched types
 {
     loop {}
 }

tests/ui/alloc-error/alloc-error-handler-bad-signature-1.stderr

@@ -1,12 +1,16 @@
 error[E0308]: mismatched types
-  --> $DIR/alloc-error-handler-bad-signature-1.rs:9:1
+  --> $DIR/alloc-error-handler-bad-signature-1.rs:11:5
    |
 LL |   #[alloc_error_handler]
-   |   ^^^^^^^^^^^^^^^^^^^^^^ expected `&Layout`, found `Layout`
+   |   ---------------------- in this procedural macro expansion
 LL | / fn oom(
 LL | |     info: &Layout,
+   | |     ^^^^^^^^^^^^^ expected `&Layout`, found `Layout`
 LL | | ) -> ()
-   | |_______- arguments to this function are incorrect
+LL | | {
+LL | |     loop {}
+LL | | }
+   | |_- arguments to this function are incorrect
    |
 note: function defined here
   --> $DIR/alloc-error-handler-bad-signature-1.rs:10:4
@@ -18,17 +22,18 @@ LL |     info: &Layout,
    = note: this error originates in the attribute macro `alloc_error_handler` (in Nightly builds, run with -Z macro-backtrace for more info)
 
 error[E0308]: mismatched types
-  --> $DIR/alloc-error-handler-bad-signature-1.rs:10:1
+  --> $DIR/alloc-error-handler-bad-signature-1.rs:12:6
    |
 LL |   #[alloc_error_handler]
-   |   ----------------------
-   |   |
-   |   expected `!` because of return type
-   |   in this procedural macro expansion
+   |   ---------------------- in this procedural macro expansion
 LL | / fn oom(
 LL | |     info: &Layout,
 LL | | ) -> ()
-   | |_______^ expected `!`, found `()`
+   | |      ^^ expected `!`, found `()`
+LL | | {
+LL | |     loop {}
+LL | | }
+   | |_- expected `!` because of return type
    |
    = note:   expected type `!`
            found unit type `()`

tests/ui/alloc-error/alloc-error-handler-bad-signature-2.rs

@@ -6,10 +6,10 @@
 
 struct Layout;
 
-#[alloc_error_handler] //~ ERROR mismatched types
-fn oom( //~ ERROR mismatched types
-    info: Layout,
-) {
+#[alloc_error_handler]
+fn oom(
+    info: Layout, //~ ERROR mismatched types
+) { //~ ERROR mismatched types
     loop {}
 }
 

tests/ui/alloc-error/alloc-error-handler-bad-signature-2.stderr

@@ -1,11 +1,14 @@
 error[E0308]: mismatched types
-  --> $DIR/alloc-error-handler-bad-signature-2.rs:9:1
+  --> $DIR/alloc-error-handler-bad-signature-2.rs:11:5
    |
 LL |   #[alloc_error_handler]
-   |   ^^^^^^^^^^^^^^^^^^^^^^ expected `Layout`, found `core::alloc::Layout`
+   |   ---------------------- in this procedural macro expansion
 LL | / fn oom(
 LL | |     info: Layout,
+   | |     ^^^^^^^^^^^^ expected `Layout`, found `core::alloc::Layout`
 LL | | ) {
+LL | |     loop {}
+LL | | }
    | |_- arguments to this function are incorrect
    |
    = note: `core::alloc::Layout` and `Layout` have similar names, but are actually distinct types
@@ -26,17 +29,17 @@ LL |     info: Layout,
    = note: this error originates in the attribute macro `alloc_error_handler` (in Nightly builds, run with -Z macro-backtrace for more info)
 
 error[E0308]: mismatched types
-  --> $DIR/alloc-error-handler-bad-signature-2.rs:10:1
+  --> $DIR/alloc-error-handler-bad-signature-2.rs:12:3
    |
 LL |   #[alloc_error_handler]
-   |   ----------------------
-   |   |
-   |   expected `!` because of return type
-   |   in this procedural macro expansion
+   |   ---------------------- in this procedural macro expansion
 LL | / fn oom(
 LL | |     info: Layout,
 LL | | ) {
-   | |_^ expected `!`, found `()`
+   | |   ^ expected `!`, found `()`
+LL | |     loop {}
+LL | | }
+   | |_- expected `!` because of return type
    |
    = note:   expected type `!`
            found unit type `()`

tests/ui/alloc-error/alloc-error-handler-bad-signature-3.stderr

@@ -1,13 +1,12 @@
 error[E0061]: this function takes 0 arguments but 1 argument was supplied
   --> $DIR/alloc-error-handler-bad-signature-3.rs:10:1
    |
-LL | #[alloc_error_handler]
-   | ----------------------
-   | |
-   | unexpected argument of type `core::alloc::Layout`
-   | in this procedural macro expansion
-LL | fn oom() -> ! {
-   | ^^^^^^^^^^^^^
+LL |   #[alloc_error_handler]
+   |   ---------------------- in this procedural macro expansion
+LL | / fn oom() -> ! {
+LL | |     loop {}
+LL | | }
+   | |_^ unexpected argument of type `core::alloc::Layout`
    |
 note: function defined here
   --> $DIR/alloc-error-handler-bad-signature-3.rs:10:4