Refactor contract builtin macro + error handling

celinval · celinval · commit 804cce47d96d · 2025-02-03T13:55:15.000-08:00
Instead of parsing the different components of a function signature,
eagerly look for either the `where` keyword or the function body.

- Also address feedback to use `From` instead of `TryFrom` in cranelift
  contract and ubcheck codegen.
diff --git a/compiler/rustc_builtin_macros/src/contracts.rs b/compiler/rustc_builtin_macros/src/contracts.rs
@@ -35,98 +35,98 @@ impl AttrProcMacro for ExpandEnsures {
     }
 }
 
-fn expand_injecting_circa_where_clause(
+/// Expand the function signature to include the contract clause.
+///
+/// The contracts clause will be injected before the function body and the optional where clause.
+/// For that, we search for the body / where token, and invoke the `inject` callback to generate the
+/// contract clause in the right place.
+///
+// FIXME: this kind of manual token tree munging does not have significant precedent among
+// rustc builtin macros, probably because most builtin macros use direct AST manipulation to
+// accomplish similar goals. But since our attributes need to take arbitrary expressions, and
+// our attribute infrastructure does not yet support mixing a token-tree annotation with an AST
+// annotated, we end up doing token tree manipulation.
+fn expand_contract_clause(
     ecx: &mut ExtCtxt<'_>,
     attr_span: Span,
     annotated: TokenStream,
-    inject: impl FnOnce(&mut Vec<TokenTree>) -> Result<(), ErrorGuaranteed>,
+    inject: impl FnOnce(&mut TokenStream) -> Result<(), ErrorGuaranteed>,
 ) -> Result<TokenStream, ErrorGuaranteed> {
-    let mut new_tts = Vec::with_capacity(annotated.len());
+    let mut new_tts = TokenStream::default();
     let mut cursor = annotated.iter();
 
-    // Find the `fn name<G,...>(x:X,...)` and inject the AST contract forms right after
-    // the formal parameters (and return type if any).
-    while let Some(tt) = cursor.next() {
-        new_tts.push(tt.clone());
-        if let TokenTree::Token(tok, _) = tt
-            && tok.is_ident_named(kw::Fn)
-        {
-            break;
-        }
+    let is_kw = |tt: &TokenTree, sym: Symbol| {
+        if let TokenTree::Token(token, _) = tt { token.is_ident_named(sym) } else { false }
+    };
+
+    // Find the `fn` keyword to check if this is a function.
+    if cursor
+        .find(|tt| {
+            new_tts.push_tree((*tt).clone());
+            is_kw(tt, kw::Fn)
+        })
+        .is_none()
+    {
+        return Err(ecx
+            .sess
+            .dcx()
+            .span_err(attr_span, "contract annotations can only be used on functions"));
     }
 
-    // Found the `fn` keyword, now find the formal parameters.
-    //
-    // FIXME: can this fail if you have parentheticals in a generics list, like `fn foo<F: Fn(X) -> Y>` ?
-    while let Some(tt) = cursor.next() {
-        new_tts.push(tt.clone());
-
-        if let TokenTree::Delimited(_, _, token::Delimiter::Parenthesis, _) = tt {
-            break;
-        }
-        if let TokenTree::Token(token::Token { kind: token::TokenKind::Semi, .. }, _) = tt {
-            panic!("contract attribute applied to fn without parameter list.");
+    // Found the `fn` keyword, now find either the `where` token or the function body.
+    let next_tt = loop {
+        let Some(tt) = cursor.next() else {
+            return Err(ecx.sess.dcx().span_err(
+                attr_span,
+                "contract annotations is only supported in functions with bodies",
+            ));
+        };
+        // If `tt` is the last element. Check if it is the function body.
+        if cursor.peek().is_none() {
+            if let TokenTree::Delimited(_, _, token::Delimiter::Brace, _) = tt {
+                break tt;
+            } else {
+                return Err(ecx.sess.dcx().span_err(
+                    attr_span,
+                    "contract annotations is only supported in functions with bodies",
+                ));
+            }
         }
-    }
 
-    // There *might* be a return type declaration (and figuring out where that ends would require
-    // parsing an arbitrary type expression, e.g. `-> Foo<args ...>`
-    //
-    // Instead of trying to figure that out, scan ahead and look for the first occurence of a
-    // `where`, a `{ ... }`, or a `;`.
-    //
-    // FIXME: this might still fall into a trap for something like `-> Ctor<T, const { 0 }>`. I
-    // *think* such cases must be under a Delimited (e.g. `[T; { N }]` or have the braced form
-    // prefixed by e.g. `const`, so we should still be able to filter them out without having to
-    // parse the type expression itself. But rather than try to fix things with hacks like that,
-    // time might be better spent extending the attribute expander to suport tt-annotation atop
-    // ast-annotated, which would be an elegant way to sidestep all of this.
-    let mut opt_next_tt = cursor.next();
-    while let Some(next_tt) = opt_next_tt {
-        if let TokenTree::Token(tok, _) = next_tt
-            && tok.is_ident_named(kw::Where)
-        {
-            break;
-        }
-        if let TokenTree::Delimited(_, _, token::Delimiter::Brace, _) = next_tt {
-            break;
-        }
-        if let TokenTree::Token(token::Token { kind: token::TokenKind::Semi, .. }, _) = next_tt {
-            break;
+        if is_kw(tt, kw::Where) {
+            break tt;
         }
-
-        // for anything else, transcribe the tt and keep looking.
-        new_tts.push(next_tt.clone());
-        opt_next_tt = cursor.next();
-    }
+        new_tts.push_tree(tt.clone());
+    };
 
     // At this point, we've transcribed everything from the `fn` through the formal parameter list
     // and return type declaration, (if any), but `tt` itself has *not* been transcribed.
     //
     // Now inject the AST contract form.
     //
-    // FIXME: this kind of manual token tree munging does not have significant precedent among
-    // rustc builtin macros, probably because most builtin macros use direct AST manipulation to
-    // accomplish similar goals. But since our attributes need to take arbitrary expressions, and
-    // our attribute infrastructure does not yet support mixing a token-tree annotation with an AST
-    // annotated, we end up doing token tree manipulation.
     inject(&mut new_tts)?;
 
-    // Above we injected the internal AST requires/ensures contruct. Now copy over all the other
+    // Above we injected the internal AST requires/ensures construct. Now copy over all the other
     // token trees.
-    if let Some(tt) = opt_next_tt {
-        new_tts.push(tt.clone());
-    }
+    new_tts.push_tree(next_tt.clone());
     while let Some(tt) = cursor.next() {
-        new_tts.push(tt.clone());
+        new_tts.push_tree(tt.clone());
+        if cursor.peek().is_none()
+            && !matches!(tt, TokenTree::Delimited(_, _, token::Delimiter::Brace, _))
+        {
+            return Err(ecx.sess.dcx().span_err(
+                attr_span,
+                "contract annotations is only supported in functions with bodies",
+            ));
+        }
     }
 
     // Record the span as a contract attribute expansion.
     // This is used later to stop users from using the extended syntax directly
     // which is gated via `rustc_contracts_internals`.
     ecx.psess().contract_attribute_spans.push(attr_span);
 
-    Ok(TokenStream::new(new_tts))
+    Ok(new_tts)
 }
 
 fn expand_requires_tts(
@@ -135,16 +135,16 @@ fn expand_requires_tts(
     annotation: TokenStream,
     annotated: TokenStream,
 ) -> Result<TokenStream, ErrorGuaranteed> {
-    expand_injecting_circa_where_clause(_ecx, attr_span, annotated, |new_tts| {
-        new_tts.push(TokenTree::Token(
+    expand_contract_clause(_ecx, attr_span, annotated, |new_tts| {
+        new_tts.push_tree(TokenTree::Token(
             token::Token::from_ast_ident(Ident::new(kw::RustcContractRequires, attr_span)),
             Spacing::Joint,
         ));
-        new_tts.push(TokenTree::Token(
+        new_tts.push_tree(TokenTree::Token(
             token::Token::new(token::TokenKind::OrOr, attr_span),
             Spacing::Alone,
         ));
-        new_tts.push(TokenTree::Delimited(
+        new_tts.push_tree(TokenTree::Delimited(
             DelimSpan::from_single(attr_span),
             DelimSpacing::new(Spacing::JointHidden, Spacing::JointHidden),
             token::Delimiter::Parenthesis,
@@ -160,12 +160,12 @@ fn expand_ensures_tts(
     annotation: TokenStream,
     annotated: TokenStream,
 ) -> Result<TokenStream, ErrorGuaranteed> {
-    expand_injecting_circa_where_clause(_ecx, attr_span, annotated, |new_tts| {
-        new_tts.push(TokenTree::Token(
+    expand_contract_clause(_ecx, attr_span, annotated, |new_tts| {
+        new_tts.push_tree(TokenTree::Token(
             token::Token::from_ast_ident(Ident::new(kw::RustcContractEnsures, attr_span)),
             Spacing::Joint,
         ));
-        new_tts.push(TokenTree::Delimited(
+        new_tts.push_tree(TokenTree::Delimited(
             DelimSpan::from_single(attr_span),
             DelimSpacing::new(Spacing::JointHidden, Spacing::JointHidden),
             token::Delimiter::Parenthesis,
diff --git a/compiler/rustc_codegen_cranelift/src/base.rs b/compiler/rustc_codegen_cranelift/src/base.rs
@@ -868,7 +868,7 @@ fn codegen_stmt<'tcx>(
                         NullOp::UbChecks => {
                             let val = fx.tcx.sess.ub_checks();
                             let val = CValue::by_val(
-                                fx.bcx.ins().iconst(types::I8, i64::try_from(val).unwrap()),
+                                fx.bcx.ins().iconst(types::I8, i64::from(val)),
                                 fx.layout_of(fx.tcx.types.bool),
                             );
                             lval.write_cvalue(fx, val);
@@ -877,7 +877,7 @@ fn codegen_stmt<'tcx>(
                         NullOp::ContractChecks => {
                             let val = fx.tcx.sess.contract_checks();
                             let val = CValue::by_val(
-                                fx.bcx.ins().iconst(types::I8, i64::try_from(val).unwrap()),
+                                fx.bcx.ins().iconst(types::I8, i64::from(val)),
                                 fx.layout_of(fx.tcx.types.bool),
                             );
                             lval.write_cvalue(fx, val);
diff --git a/tests/ui/contracts/contract-annotation-limitations.rs b/tests/ui/contracts/contract-annotation-limitations.rs
@@ -0,0 +1,27 @@
+//! Test for some of the existing limitations and the current error messages.
+//! Some of these limitations may be removed in the future.
+
+#![feature(rustc_contracts)]
+#![allow(dead_code)]
+
+/// Represent a 5-star system.
+struct Stars(u8);
+
+impl Stars {
+    fn is_valid(&self) -> bool {
+        self.0 <= 5
+    }
+}
+
+trait ParseStars {
+    #[core::contracts::ensures(|ret| ret.is_none_or(Stars::is_valid))]
+    //~^ ERROR contract annotations is only supported in functions with bodies
+    fn parse_string(input: String) -> Option<Stars>;
+
+    #[core::contracts::ensures(|ret| ret.is_none_or(Stars::is_valid))]
+    //~^ ERROR contract annotations is only supported in functions with bodies
+    fn parse<T>(input: T) -> Option<Stars> where T: for<'a> Into<&'a str>;
+}
+
+fn main() {
+}
diff --git a/tests/ui/contracts/contract-annotation-limitations.stderr b/tests/ui/contracts/contract-annotation-limitations.stderr
@@ -0,0 +1,14 @@
+error: contract annotations is only supported in functions with bodies
+  --> $DIR/contract-annotation-limitations.rs:17:5
+   |
+LL |     #[core::contracts::ensures(|ret| ret.is_none_or(Stars::is_valid))]
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: contract annotations is only supported in functions with bodies
+  --> $DIR/contract-annotation-limitations.rs:21:5
+   |
+LL |     #[core::contracts::ensures(|ret| ret.is_none_or(Stars::is_valid))]
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: aborting due to 2 previous errors
+
diff --git a/tests/ui/contracts/contract-attributes-generics.rs b/tests/ui/contracts/contract-attributes-generics.rs
@@ -0,0 +1,70 @@
+//! Test that contracts can be applied to generic functions.
+
+//@ revisions: unchk_pass chk_pass chk_fail_pre chk_fail_post chk_const_fail
+//
+//@ [unchk_pass] run-pass
+//@ [chk_pass] run-pass
+//
+//@ [chk_fail_pre] run-fail
+//@ [chk_fail_post] run-fail
+//@ [chk_const_fail] run-fail
+//
+//@ [unchk_pass] compile-flags: -Zcontract-checks=no
+//
+//@ [chk_pass] compile-flags: -Zcontract-checks=yes
+//@ [chk_fail_pre] compile-flags: -Zcontract-checks=yes
+//@ [chk_fail_post] compile-flags: -Zcontract-checks=yes
+//@ [chk_const_fail] compile-flags: -Zcontract-checks=yes
+
+#![feature(rustc_contracts)]
+
+use std::ops::Sub;
+
+/// Dummy fn contract that precondition fails for val < 0, and post-condition fail for val == 1
+#[core::contracts::requires(val > 0u8.into())]
+#[core::contracts::ensures(|ret| *ret > 0u8.into())]
+fn decrement<T>(val: T) -> T
+where T: PartialOrd + Sub<Output=T> + From<u8>
+{
+    val - 1u8.into()
+}
+
+/// Create a structure that takes a constant parameter.
+#[allow(dead_code)]
+struct Capped<const MAX: usize>(usize);
+
+/// Now declare a function to create stars which shouldn't exceed 5 stars.
+// Add redundant braces to ensure the built-in macro can handle this syntax.
+#[allow(unused_braces)]
+#[core::contracts::requires(num <= 5)]
+unsafe fn stars_unchecked(num: usize) -> Capped<{ 5 }> {
+    Capped(num)
+}
+
+
+fn main() {
+    check_decrement();
+    check_stars();
+}
+
+fn check_stars() {
+    // This should always pass.
+    let _ = unsafe { stars_unchecked(3) };
+
+    // This violates the contract.
+    #[cfg(any(unchk_pass, chk_const_fail))]
+    let _ = unsafe { stars_unchecked(10) };
+}
+
+fn check_decrement() {
+    // This should always pass
+    assert_eq!(decrement(10u8), 9u8);
+
+    // This should fail requires but pass with no contract check.
+    #[cfg(any(unchk_pass, chk_fail_pre))]
+    assert_eq!(decrement(-2i128), -3i128);
+
+    // This should fail ensures but pass with no contract check.
+    #[cfg(any(unchk_pass, chk_fail_post))]
+    assert_eq!(decrement(1i32), 0i32);
+}
diff --git a/tests/ui/contracts/disallow-contract-annotation-on-non-fn.rs b/tests/ui/contracts/disallow-contract-annotation-on-non-fn.rs
@@ -0,0 +1,52 @@
+//! Checks for compilation errors related to adding contracts to non-function items.
+
+#![feature(rustc_contracts)]
+#![allow(dead_code)]
+
+#[core::contracts::requires(true)]
+//~^ ERROR contract annotations can only be used on functions
+struct Dummy(usize);
+
+#[core::contracts::ensures(|v| v == 100)]
+//~^ ERROR contract annotations can only be used on functions
+const MAX_VAL: usize = 100;
+
+// FIXME: Improve the error message here. The macro thinks this is a function.
+#[core::contracts::ensures(|v| v == 100)]
+//~^ ERROR contract annotations is only supported in functions with bodies
+type NewDummy = fn(usize) -> Dummy;
+
+#[core::contracts::ensures(|v| v == 100)]
+//~^ ERROR contract annotations is only supported in functions with bodies
+const NEW_DUMMY_FN : fn(usize) -> Dummy = || { Dummy(0) };
+
+#[core::contracts::requires(true)]
+//~^ ERROR contract annotations can only be used on functions
+impl Dummy {
+
+    // This should work
+    #[core::contracts::ensures(|ret| ret.0 == v)]
+    fn new(v: usize) -> Dummy {
+        Dummy(v)
+    }
+}
+
+#[core::contracts::ensures(|dummy| dummy.0 > 0)]
+//~^ ERROR contract annotations can only be used on functions
+impl From<usize> for Dummy {
+    // This should work.
+    #[core::contracts::ensures(|ret| ret.0 == v)]
+    fn from(value: usize) -> Self {
+        Dummy::new(value)
+    }
+}
+
+/// You should not be able to annotate a trait either.
+#[core::contracts::requires(true)]
+//~^ ERROR contract annotations can only be used on functions
+pub trait DummyBuilder {
+    fn build() -> Dummy;
+}
+
+fn main() {
+}
diff --git a/tests/ui/contracts/disallow-contract-annotation-on-non-fn.stderr b/tests/ui/contracts/disallow-contract-annotation-on-non-fn.stderr
