Desugars contract into the internal AST extensions

Check ensures on early return due to Try / Yeet

Expand these two expressions to include a call to contract checking

Author: pnkfelix

compiler/rustc_ast_lowering/src/expr.rs

@@ -314,21 +314,8 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     hir::ExprKind::Continue(self.lower_jump_destination(e.id, *opt_label))
                 }
                 ExprKind::Ret(e) => {
-                    let mut e = e.as_ref().map(|x| self.lower_expr(x));
-                    if let Some(Some((span, fresh_ident))) = self
-                        .contract
-                        .as_ref()
-                        .map(|c| c.ensures.as_ref().map(|e| (e.expr.span, e.fresh_ident)))
-                    {
-                        let checker_fn = self.expr_ident(span, fresh_ident.0, fresh_ident.2);
-                        let args = if let Some(e) = e {
-                            std::slice::from_ref(e)
-                        } else {
-                            std::slice::from_ref(self.expr_unit(span))
-                        };
-                        e = Some(self.expr_call(span, checker_fn, args));
-                    }
-                    hir::ExprKind::Ret(e)
+                    let expr = e.as_ref().map(|x| self.lower_expr(x));
+                    self.checked_return(expr)
                 }
                 ExprKind::Yeet(sub_expr) => self.lower_expr_yeet(e.span, sub_expr.as_deref()),
                 ExprKind::Become(sub_expr) => {
@@ -395,6 +382,32 @@ impl<'hir> LoweringContext<'_, 'hir> {
         })
     }
 
+    /// Create an `ExprKind::Ret` that is preceded by a call to check contract ensures clause.
+    fn checked_return(&mut self, opt_expr: Option<&'hir hir::Expr<'hir>>) -> hir::ExprKind<'hir> {
+        let checked_ret = if let Some(Some((span, fresh_ident))) =
+            self.contract.as_ref().map(|c| c.ensures.as_ref().map(|e| (e.expr.span, e.fresh_ident)))
+        {
+            let expr = opt_expr.unwrap_or_else(|| self.expr_unit(span));
+            Some(self.inject_ensures_check(expr, span, fresh_ident.0, fresh_ident.2))
+        } else {
+            opt_expr
+        };
+        hir::ExprKind::Ret(checked_ret)
+    }
+
+    /// Wraps an expression with a call to the ensures check before it gets returned.
+    pub(crate) fn inject_ensures_check(
+        &mut self,
+        expr: &'hir hir::Expr<'hir>,
+        span: Span,
+        check_ident: Ident,
+        check_hir_id: HirId,
+    ) -> &'hir hir::Expr<'hir> {
+        let checker_fn = self.expr_ident(span, check_ident, check_hir_id);
+        let span = self.mark_span_with_reason(DesugaringKind::Contract, span, None);
+        self.expr_call(span, checker_fn, std::slice::from_ref(expr))
+    }
+
     pub(crate) fn lower_const_block(&mut self, c: &AnonConst) -> hir::ConstBlock {
         self.with_new_scopes(c.value.span, |this| {
             let def_id = this.local_def_id(c.id);
@@ -1983,7 +1996,8 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     ),
                 ))
             } else {
-                self.arena.alloc(self.expr(try_span, hir::ExprKind::Ret(Some(from_residual_expr))))
+                let ret_expr = self.checked_return(Some(from_residual_expr));
+                self.arena.alloc(self.expr(try_span, ret_expr))
             };
             self.lower_attrs(ret_expr.hir_id, &attrs);
 
@@ -2032,7 +2046,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             let target_id = Ok(catch_id);
             hir::ExprKind::Break(hir::Destination { label: None, target_id }, Some(from_yeet_expr))
         } else {
-            hir::ExprKind::Ret(Some(from_yeet_expr))
+            self.checked_return(Some(from_yeet_expr))
         }
     }
 

compiler/rustc_ast_lowering/src/item.rs

@@ -215,7 +215,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     if let Some(contract) = contract {
                         let requires = contract.requires.clone();
                         let ensures = contract.ensures.clone();
-                        let ensures = if let Some(ens) = ensures {
+                        let ensures = ensures.map(|ens| {
                             // FIXME: this needs to be a fresh (or illegal) identifier to prevent
                             // accidental capture of a parameter or global variable.
                             let checker_ident: Ident =
@@ -226,13 +226,11 @@ impl<'hir> LoweringContext<'_, 'hir> {
                                 hir::BindingMode::NONE,
                             );
 
-                            Some(crate::FnContractLoweringEnsures {
+                            crate::FnContractLoweringEnsures {
                                 expr: ens,
                                 fresh_ident: (checker_ident, checker_pat, checker_hir_id),
-                            })
-                        } else {
-                            None
-                        };
+                            }
+                        });
 
                         // Note: `with_new_scopes` will reinstall the outer
                         // item's contract (if any) after its callback finishes.
@@ -1095,73 +1093,56 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
             // { body }
             // ==>
-            // { rustc_contract_requires(PRECOND); { body } }
-            let result: hir::Expr<'hir> = if let Some(contract) = opt_contract {
-                let lit_unit = |this: &mut LoweringContext<'_, 'hir>| {
-                    this.expr(contract.span, hir::ExprKind::Tup(&[]))
-                };
-
-                let precond: hir::Stmt<'hir> = if let Some(req) = contract.requires {
-                    let lowered_req = this.lower_expr_mut(&req);
-                    let precond = this.expr_call_lang_item_fn_mut(
-                        req.span,
-                        hir::LangItem::ContractCheckRequires,
-                        &*arena_vec![this; lowered_req],
-                    );
-                    this.stmt_expr(req.span, precond)
-                } else {
-                    let u = lit_unit(this);
-                    this.stmt_expr(contract.span, u)
-                };
+            // { contract_requires(PRECOND); { body } }
+            let Some(contract) = opt_contract else { return (params, result) };
+            let result_ref = this.arena.alloc(result);
+            let lit_unit = |this: &mut LoweringContext<'_, 'hir>| {
+                this.expr(contract.span, hir::ExprKind::Tup(&[]))
+            };
 
-                let (postcond_checker, result) = if let Some(ens) = contract.ensures {
-                    let crate::FnContractLoweringEnsures { expr: ens, fresh_ident } = ens;
-                    let lowered_ens: hir::Expr<'hir> = this.lower_expr_mut(&ens);
-                    let postcond_checker = this.expr_call_lang_item_fn(
-                        ens.span,
-                        hir::LangItem::ContractBuildCheckEnsures,
-                        &*arena_vec![this; lowered_ens],
-                    );
-                    let checker_binding_pat = fresh_ident.1;
-                    (
-                        this.stmt_let_pat(
-                            None,
-                            ens.span,
-                            Some(postcond_checker),
-                            this.arena.alloc(checker_binding_pat),
-                            hir::LocalSource::Contract,
-                        ),
-                        {
-                            let checker_fn =
-                                this.expr_ident(ens.span, fresh_ident.0, fresh_ident.2);
-                            let span = this.mark_span_with_reason(
-                                DesugaringKind::Contract,
-                                ens.span,
-                                None,
-                            );
-                            this.expr_call_mut(
-                                span,
-                                checker_fn,
-                                std::slice::from_ref(this.arena.alloc(result)),
-                            )
-                        },
-                    )
-                } else {
-                    let u = lit_unit(this);
-                    (this.stmt_expr(contract.span, u), result)
-                };
+            let precond: hir::Stmt<'hir> = if let Some(req) = contract.requires {
+                let lowered_req = this.lower_expr_mut(&req);
+                let precond = this.expr_call_lang_item_fn_mut(
+                    req.span,
+                    hir::LangItem::ContractCheckRequires,
+                    &*arena_vec![this; lowered_req],
+                );
+                this.stmt_expr(req.span, precond)
+            } else {
+                let u = lit_unit(this);
+                this.stmt_expr(contract.span, u)
+            };
 
-                let block = this.block_all(
-                    contract.span,
-                    arena_vec![this; precond, postcond_checker],
-                    Some(this.arena.alloc(result)),
+            let (postcond_checker, result) = if let Some(ens) = contract.ensures {
+                let crate::FnContractLoweringEnsures { expr: ens, fresh_ident } = ens;
+                let lowered_ens: hir::Expr<'hir> = this.lower_expr_mut(&ens);
+                let postcond_checker = this.expr_call_lang_item_fn(
+                    ens.span,
+                    hir::LangItem::ContractBuildCheckEnsures,
+                    &*arena_vec![this; lowered_ens],
                 );
-                this.expr_block(block)
+                let checker_binding_pat = fresh_ident.1;
+                (
+                    this.stmt_let_pat(
+                        None,
+                        ens.span,
+                        Some(postcond_checker),
+                        this.arena.alloc(checker_binding_pat),
+                        hir::LocalSource::Contract,
+                    ),
+                    this.inject_ensures_check(result_ref, ens.span, fresh_ident.0, fresh_ident.2),
+                )
             } else {
-                result
+                let u = lit_unit(this);
+                (this.stmt_expr(contract.span, u), &*result_ref)
             };
 
-            (params, result)
+            let block = this.block_all(
+                contract.span,
+                arena_vec![this; precond, postcond_checker],
+                Some(result),
+            );
+            (params, this.expr_block(block))
         })
     }
 

compiler/rustc_builtin_macros/src/contracts.rs

@@ -0,0 +1,172 @@
+#![allow(unused_imports, unused_variables)]
+
+use rustc_ast::token;
+use rustc_ast::tokenstream::{DelimSpacing, DelimSpan, Spacing, TokenStream, TokenTree};
+use rustc_errors::ErrorGuaranteed;
+use rustc_expand::base::{AttrProcMacro, ExtCtxt};
+use rustc_span::Span;
+use rustc_span::symbol::{Ident, Symbol, kw, sym};
+
+pub(crate) struct ExpandRequires;
+
+pub(crate) struct ExpandEnsures;
+
+impl AttrProcMacro for ExpandRequires {
+    fn expand<'cx>(
+        &self,
+        ecx: &'cx mut ExtCtxt<'_>,
+        span: Span,
+        annotation: TokenStream,
+        annotated: TokenStream,
+    ) -> Result<TokenStream, ErrorGuaranteed> {
+        expand_requires_tts(ecx, span, annotation, annotated)
+    }
+}
+
+impl AttrProcMacro for ExpandEnsures {
+    fn expand<'cx>(
+        &self,
+        ecx: &'cx mut ExtCtxt<'_>,
+        span: Span,
+        annotation: TokenStream,
+        annotated: TokenStream,
+    ) -> Result<TokenStream, ErrorGuaranteed> {
+        expand_ensures_tts(ecx, span, annotation, annotated)
+    }
+}
+
+fn expand_injecting_circa_where_clause(
+    _ecx: &mut ExtCtxt<'_>,
+    attr_span: Span,
+    annotated: TokenStream,
+    inject: impl FnOnce(&mut Vec<TokenTree>) -> Result<(), ErrorGuaranteed>,
+) -> Result<TokenStream, ErrorGuaranteed> {
+    let mut new_tts = Vec::with_capacity(annotated.len());
+    let mut cursor = annotated.into_trees();
+
+    // 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_ref() {
+        new_tts.push(tt.clone());
+        if let TokenTree::Token(tok, _) = tt
+            && tok.is_ident_named(kw::Fn)
+        {
+            break;
+        }
+    }
+
+    // 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_ref() {
+        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.");
+        }
+    }
+
+    // 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_ref();
+    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;
+        }
+
+        // for anything else, transcribe the tt and keep looking.
+        new_tts.push(next_tt.clone());
+        opt_next_tt = cursor.next_ref();
+        continue;
+    }
+
+    // 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
+    // token trees.
+    if let Some(tt) = opt_next_tt {
+        new_tts.push(tt.clone());
+    }
+    while let Some(tt) = cursor.next_ref() {
+        new_tts.push(tt.clone());
+    }
+
+    Ok(TokenStream::new(new_tts))
+}
+
+fn expand_requires_tts(
+    _ecx: &mut ExtCtxt<'_>,
+    attr_span: Span,
+    annotation: TokenStream,
+    annotated: TokenStream,
+) -> Result<TokenStream, ErrorGuaranteed> {
+    expand_injecting_circa_where_clause(_ecx, attr_span, annotated, |new_tts| {
+        new_tts.push(TokenTree::Token(
+            token::Token::from_ast_ident(Ident::new(kw::RustcContractRequires, attr_span)),
+            Spacing::Joint,
+        ));
+        new_tts.push(TokenTree::Token(
+            token::Token::new(token::TokenKind::OrOr, attr_span),
+            Spacing::Alone,
+        ));
+        new_tts.push(TokenTree::Delimited(
+            DelimSpan::from_single(attr_span),
+            DelimSpacing::new(Spacing::JointHidden, Spacing::JointHidden),
+            token::Delimiter::Parenthesis,
+            annotation,
+        ));
+        Ok(())
+    })
+}
+
+fn expand_ensures_tts(
+    _ecx: &mut ExtCtxt<'_>,
+    attr_span: Span,
+    annotation: TokenStream,
+    annotated: TokenStream,
+) -> Result<TokenStream, ErrorGuaranteed> {
+    expand_injecting_circa_where_clause(_ecx, attr_span, annotated, |new_tts| {
+        new_tts.push(TokenTree::Token(
+            token::Token::from_ast_ident(Ident::new(kw::RustcContractEnsures, attr_span)),
+            Spacing::Joint,
+        ));
+        new_tts.push(TokenTree::Delimited(
+            DelimSpan::from_single(attr_span),
+            DelimSpacing::new(Spacing::JointHidden, Spacing::JointHidden),
+            token::Delimiter::Parenthesis,
+            annotation,
+        ));
+        Ok(())
+    })
+}