Fix ICE for invalid return activity and proper error handling

3 files changed, 25 insertions, 1 deletions

diff --git a/compiler/rustc_builtin_macros/messages.ftl b/compiler/rustc_builtin_macros/messages.ftl
index 4cac7cb93f5..9dac7a8aaf4 100644
--- a/compiler/rustc_builtin_macros/messages.ftl
+++ b/compiler/rustc_builtin_macros/messages.ftl
@@ -75,8 +75,10 @@ builtin_macros_autodiff_mode = unknown Mode: `{$mode}`. Use `Forward` or `Revers
 builtin_macros_autodiff_mode_activity = {$act} can not be used in {$mode} Mode
 builtin_macros_autodiff_not_build = this rustc version does not support autodiff
 builtin_macros_autodiff_number_activities = expected {$expected} activities, but found {$found}
+builtin_macros_autodiff_ret_activity = invalid return activity {$act} in {$mode} Mode
 builtin_macros_autodiff_ty_activity = {$act} can not be used for this type
 
+
 builtin_macros_autodiff_unknown_activity = did not recognize Activity: `{$act}`
 builtin_macros_bad_derive_target = `derive` may only be applied to `struct`s, `enum`s and `union`s
     .label = not applicable here
diff --git a/compiler/rustc_builtin_macros/src/autodiff.rs b/compiler/rustc_builtin_macros/src/autodiff.rs
index 6d9c3575657..6707ab8edde 100644
--- a/compiler/rustc_builtin_macros/src/autodiff.rs
+++ b/compiler/rustc_builtin_macros/src/autodiff.rs
@@ -8,7 +8,8 @@ mod llvm_enzyme {
     use std::string::String;
 
     use rustc_ast::expand::autodiff_attrs::{
-        AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ty_for_activity,
+        AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ret_activity,
+        valid_ty_for_activity,
     };
     use rustc_ast::ptr::P;
     use rustc_ast::token::{Token, TokenKind};
@@ -632,10 +633,22 @@ mod llvm_enzyme {
                 errors = true;
             }
         }
+
+        if has_ret && !valid_ret_activity(x.mode, x.ret_activity) {
+            dcx.emit_err(errors::AutoDiffInvalidRetAct {
+                span,
+                mode: x.mode.to_string(),
+                act: x.ret_activity.to_string(),
+            });
+            // We don't set `errors = true` to avoid annoying type errors relative
+            // to the expanded macro type signature
+        }
+
         if errors {
             // This is not the right signature, but we can continue parsing.
             return (sig.clone(), new_inputs, idents, true);
         }
+
         let unsafe_activities = x
             .input_activity
             .iter()
diff --git a/compiler/rustc_builtin_macros/src/errors.rs b/compiler/rustc_builtin_macros/src/errors.rs
index ab1e0d8ee89..30597944124 100644
--- a/compiler/rustc_builtin_macros/src/errors.rs
+++ b/compiler/rustc_builtin_macros/src/errors.rs
@@ -186,6 +186,15 @@ mod autodiff {
     }
 
     #[derive(Diagnostic)]
+    #[diag(builtin_macros_autodiff_ret_activity)]
+    pub(crate) struct AutoDiffInvalidRetAct {
+        #[primary_span]
+        pub(crate) span: Span,
+        pub(crate) mode: String,
+        pub(crate) act: String,
+    }
+
+    #[derive(Diagnostic)]
     #[diag(builtin_macros_autodiff_mode)]
     pub(crate) struct AutoDiffInvalidMode {
         #[primary_span]