Suggest expressions that look like const generic arguments should be enclosed in brackets

Co-Authored-By: Esteban Kuber <github@kuber.com.ar>

4 files changed, 190 insertions, 22 deletions

diff --git a/compiler/rustc_parse/src/parser/diagnostics.rs b/compiler/rustc_parse/src/parser/diagnostics.rs
index 1ea01d95a13..52cbba9d2bf 100644
--- a/compiler/rustc_parse/src/parser/diagnostics.rs
+++ b/compiler/rustc_parse/src/parser/diagnostics.rs
@@ -1,13 +1,14 @@
 use super::ty::AllowPlus;
-use super::{BlockMode, Parser, PathStyle, SemiColonMode, SeqSep, TokenExpectType, TokenType};
+use super::TokenType;
+use super::{BlockMode, Parser, PathStyle, Restrictions, SemiColonMode, SeqSep, TokenExpectType};
 
 use rustc_ast::ptr::P;
 use rustc_ast::token::{self, Lit, LitKind, TokenKind};
 use rustc_ast::util::parser::AssocOp;
 use rustc_ast::{
-    self as ast, AngleBracketedArgs, AttrVec, BinOpKind, BindingMode, Block, BlockCheckMode, Expr,
-    ExprKind, Item, ItemKind, Mutability, Param, Pat, PatKind, Path, PathSegment, QSelf, Ty,
-    TyKind,
+    self as ast, AngleBracketedArg, AngleBracketedArgs, AnonConst, AttrVec, BinOpKind, BindingMode,
+    Block, BlockCheckMode, Expr, ExprKind, GenericArg, Item, ItemKind, Mutability, Param, Pat,
+    PatKind, Path, PathSegment, QSelf, Ty, TyKind,
 };
 use rustc_ast_pretty::pprust;
 use rustc_data_structures::fx::FxHashSet;
@@ -1774,4 +1775,142 @@ impl<'a> Parser<'a> {
             }
         }
     }
+
+    /// Handle encountering a symbol in a generic argument list that is not a `,` or `>`. In this
+    /// case, we emit an error and try to suggest enclosing a const argument in braces if it looks
+    /// like the user has forgotten them.
+    pub fn handle_ambiguous_unbraced_const_arg(
+        &mut self,
+        args: &mut Vec<AngleBracketedArg>,
+    ) -> PResult<'a, bool> {
+        // If we haven't encountered a closing `>`, then the argument is malformed.
+        // It's likely that the user has written a const expression without enclosing it
+        // in braces, so we try to recover here.
+        let arg = args.pop().unwrap();
+        // FIXME: for some reason using `unexpected` or `expected_one_of_not_found` has
+        // adverse side-effects to subsequent errors and seems to advance the parser.
+        // We are causing this error here exclusively in case that a `const` expression
+        // could be recovered from the current parser state, even if followed by more
+        // arguments after a comma.
+        let mut err = self.struct_span_err(
+            self.token.span,
+            &format!("expected one of `,` or `>`, found {}", super::token_descr(&self.token)),
+        );
+        err.span_label(self.token.span, "expected one of `,` or `>`");
+        match self.recover_const_arg(arg.span(), err) {
+            Ok(arg) => {
+                args.push(AngleBracketedArg::Arg(arg));
+                if self.eat(&token::Comma) {
+                    return Ok(true); // Continue
+                }
+            }
+            Err(mut err) => {
+                args.push(arg);
+                // We will emit a more generic error later.
+                err.delay_as_bug();
+            }
+        }
+        return Ok(false); // Don't continue.
+    }
+
+    /// Handle a generic const argument that had not been enclosed in braces, and suggest enclosing
+    /// it braces. In this situation, unlike in `handle_ambiguous_unbraced_const_arg`, this is
+    /// almost certainly a const argument, so we always offer a suggestion.
+    pub fn handle_unambiguous_unbraced_const_arg(&mut self) -> PResult<'a, P<Expr>> {
+        let start = self.token.span;
+        let expr = self.parse_expr_res(Restrictions::CONST_EXPR, None).map_err(|mut err| {
+            err.span_label(
+                start.shrink_to_lo(),
+                "while parsing a const generic argument starting here",
+            );
+            err
+        })?;
+        if !self.expr_is_valid_const_arg(&expr) {
+            self.struct_span_err(
+                expr.span,
+                "expressions must be enclosed in braces to be used as const generic \
+                    arguments",
+            )
+            .multipart_suggestion(
+                "enclose the `const` expression in braces",
+                vec![
+                    (expr.span.shrink_to_lo(), "{ ".to_string()),
+                    (expr.span.shrink_to_hi(), " }".to_string()),
+                ],
+                Applicability::MachineApplicable,
+            )
+            .emit();
+        }
+        Ok(expr)
+    }
+
+    /// Try to recover from possible generic const argument without `{` and `}`.
+    ///
+    /// When encountering code like `foo::< bar + 3 >` or `foo::< bar - baz >` we suggest
+    /// `foo::<{ bar + 3 }>` and `foo::<{ bar - baz }>`, respectively. We only provide a suggestion
+    /// if we think that that the resulting expression would be well formed.
+    pub fn recover_const_arg(
+        &mut self,
+        start: Span,
+        mut err: DiagnosticBuilder<'a>,
+    ) -> PResult<'a, GenericArg> {
+        let is_op = AssocOp::from_token(&self.token)
+            .and_then(|op| {
+                if let AssocOp::Greater
+                | AssocOp::Less
+                | AssocOp::ShiftRight
+                | AssocOp::GreaterEqual
+                // Don't recover from `foo::<bar = baz>`, because this could be an attempt to
+                // assign a value to a defaulted generic parameter.
+                | AssocOp::Assign
+                | AssocOp::AssignOp(_) = op
+                {
+                    None
+                } else {
+                    Some(op)
+                }
+            })
+            .is_some();
+        // This will be true when a trait object type `Foo +` or a path which was a `const fn` with
+        // type params has been parsed.
+        let was_op =
+            matches!(self.prev_token.kind, token::BinOp(token::Plus | token::Shr) | token::Gt);
+        if !is_op && !was_op {
+            // We perform these checks and early return to avoid taking a snapshot unnecessarily.
+            return Err(err);
+        }
+        let snapshot = self.clone();
+        if is_op {
+            self.bump();
+        }
+        match self.parse_expr_res(Restrictions::CONST_EXPR, None) {
+            Ok(expr) => {
+                if token::Comma == self.token.kind || self.token.kind.should_end_const_arg() {
+                    // Avoid the following output by checking that we consumed a full const arg:
+                    // help: expressions must be enclosed in braces to be used as const generic
+                    //       arguments
+                    //    |
+                    // LL |     let sr: Vec<{ (u32, _, _) = vec![] };
+                    //    |                 ^                      ^
+                    err.multipart_suggestion(
+                        "expressions must be enclosed in braces to be used as const generic \
+                         arguments",
+                        vec![
+                            (start.shrink_to_lo(), "{ ".to_string()),
+                            (expr.span.shrink_to_hi(), " }".to_string()),
+                        ],
+                        Applicability::MaybeIncorrect,
+                    );
+                    let value = self.mk_expr_err(start.to(expr.span));
+                    err.emit();
+                    return Ok(GenericArg::Const(AnonConst { id: ast::DUMMY_NODE_ID, value }));
+                }
+            }
+            Err(mut err) => {
+                err.cancel();
+            }
+        }
+        *self = snapshot;
+        Err(err)
+    }
 }
diff --git a/compiler/rustc_parse/src/parser/expr.rs b/compiler/rustc_parse/src/parser/expr.rs
index 698a7e7d9cd..3bb0cfd81d7 100644
--- a/compiler/rustc_parse/src/parser/expr.rs
+++ b/compiler/rustc_parse/src/parser/expr.rs
@@ -359,6 +359,18 @@ impl<'a> Parser<'a> {
     /// Also performs recovery for `and` / `or` which are mistaken for `&&` and `||` respectively.
     fn check_assoc_op(&self) -> Option<Spanned<AssocOp>> {
         let (op, span) = match (AssocOp::from_token(&self.token), self.token.ident()) {
+            // When parsing const expressions, stop parsing when encountering `>`.
+            (
+                Some(
+                    AssocOp::ShiftRight
+                    | AssocOp::Greater
+                    | AssocOp::GreaterEqual
+                    | AssocOp::AssignOp(token::BinOpToken::Shr),
+                ),
+                _,
+            ) if self.restrictions.contains(Restrictions::CONST_EXPR) => {
+                return None;
+            }
             (Some(op), _) => (op, self.token.span),
             (None, Some((Ident { name: sym::and, span }, false))) => {
                 self.error_bad_logical_op("and", "&&", "conjunction");
@@ -1715,7 +1727,7 @@ impl<'a> Parser<'a> {
         let lo = self.prev_token.span;
         let pat = self.parse_top_pat(GateOr::No)?;
         self.expect(&token::Eq)?;
-        let expr = self.with_res(Restrictions::NO_STRUCT_LITERAL, |this| {
+        let expr = self.with_res(self.restrictions | Restrictions::NO_STRUCT_LITERAL, |this| {
             this.parse_assoc_expr_with(1 + prec_let_scrutinee_needs_par(), None.into())
         })?;
         let span = lo.to(expr.span);
diff --git a/compiler/rustc_parse/src/parser/mod.rs b/compiler/rustc_parse/src/parser/mod.rs
index fb825256d92..ca7a627ef18 100644
--- a/compiler/rustc_parse/src/parser/mod.rs
+++ b/compiler/rustc_parse/src/parser/mod.rs
@@ -36,6 +36,7 @@ bitflags::bitflags! {
     struct Restrictions: u8 {
         const STMT_EXPR         = 1 << 0;
         const NO_STRUCT_LITERAL = 1 << 1;
+        const CONST_EXPR        = 1 << 2;
     }
 }
 
diff --git a/compiler/rustc_parse/src/parser/path.rs b/compiler/rustc_parse/src/parser/path.rs
index 66ce015d02e..06760547eba 100644
--- a/compiler/rustc_parse/src/parser/path.rs
+++ b/compiler/rustc_parse/src/parser/path.rs
@@ -397,6 +397,13 @@ impl<'a> Parser<'a> {
         while let Some(arg) = self.parse_angle_arg()? {
             args.push(arg);
             if !self.eat(&token::Comma) {
+                if !self.token.kind.should_end_const_arg() {
+                    if self.handle_ambiguous_unbraced_const_arg(&mut args)? {
+                        // We've managed to (partially) recover, so continue trying to parse
+                        // arguments.
+                        continue;
+                    }
+                }
                 break;
             }
         }
@@ -476,41 +483,50 @@ impl<'a> Parser<'a> {
         Ok(self.mk_ty(span, ast::TyKind::Err))
     }
 
+    /// We do not permit arbitrary expressions as const arguments. They must be one of:
+    /// - An expression surrounded in `{}`.
+    /// - A literal.
+    /// - A numeric literal prefixed by `-`.
+    pub(super) fn expr_is_valid_const_arg(&self, expr: &P<rustc_ast::Expr>) -> bool {
+        match &expr.kind {
+            ast::ExprKind::Block(_, _) | ast::ExprKind::Lit(_) => true,
+            ast::ExprKind::Unary(ast::UnOp::Neg, expr) => match &expr.kind {
+                ast::ExprKind::Lit(_) => true,
+                _ => false,
+            },
+            _ => false,
+        }
+    }
+
     /// Parse a generic argument in a path segment.
     /// This does not include constraints, e.g., `Item = u8`, which is handled in `parse_angle_arg`.
     fn parse_generic_arg(&mut self) -> PResult<'a, Option<GenericArg>> {
+        let start = self.token.span;
         let arg = if self.check_lifetime() && self.look_ahead(1, |t| !t.is_like_plus()) {
             // Parse lifetime argument.
             GenericArg::Lifetime(self.expect_lifetime())
         } else if self.check_const_arg() {
             // Parse const argument.
-            let expr = if let token::OpenDelim(token::Brace) = self.token.kind {
+            let value = if let token::OpenDelim(token::Brace) = self.token.kind {
                 self.parse_block_expr(
                     None,
                     self.token.span,
                     BlockCheckMode::Default,
                     ast::AttrVec::new(),
                 )?
-            } else if self.token.is_ident() {
-                // FIXME(const_generics): to distinguish between idents for types and consts,
-                // we should introduce a GenericArg::Ident in the AST and distinguish when
-                // lowering to the HIR. For now, idents for const args are not permitted.
-                if self.token.is_bool_lit() {
-                    self.parse_literal_maybe_minus()?
-                } else {
-                    let span = self.token.span;
-                    let msg = "identifiers may currently not be used for const generics";
-                    self.struct_span_err(span, msg).emit();
-                    let block = self.mk_block_err(span);
-                    self.mk_expr(span, ast::ExprKind::Block(block, None), ast::AttrVec::new())
-                }
             } else {
-                self.parse_literal_maybe_minus()?
+                self.handle_unambiguous_unbraced_const_arg()?
             };
-            GenericArg::Const(AnonConst { id: ast::DUMMY_NODE_ID, value: expr })
+            GenericArg::Const(AnonConst { id: ast::DUMMY_NODE_ID, value })
         } else if self.check_type() {
             // Parse type argument.
-            GenericArg::Type(self.parse_ty()?)
+            match self.parse_ty() {
+                Ok(ty) => GenericArg::Type(ty),
+                Err(err) => {
+                    // Try to recover from possible `const` arg without braces.
+                    return self.recover_const_arg(start, err).map(Some);
+                }
+            }
         } else {
             return Ok(None);
         };