From 8f516d73d2d33988f6cdb9367244c11bc36ede22 Mon Sep 17 00:00:00 2001
From: Ary Borenszweig <asterite@gmail.com>
Date: Wed, 30 Oct 2024 17:50:50 -0300
Subject: [PATCH] feat: improve malformed test attribute error (#6414)

---
 compiler/noirc_frontend/src/lexer/errors.rs   |  8 ++++++++
 compiler/noirc_frontend/src/lexer/lexer.rs    |  8 +-------
 compiler/noirc_frontend/src/lexer/token.rs    | 20 ++++++++++++-------
 .../lsp/src/requests/completion/builtins.rs   |  9 +++++++++
 4 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/compiler/noirc_frontend/src/lexer/errors.rs b/compiler/noirc_frontend/src/lexer/errors.rs
index 66f79bd444b..8d799ef35d1 100644
--- a/compiler/noirc_frontend/src/lexer/errors.rs
+++ b/compiler/noirc_frontend/src/lexer/errors.rs
@@ -20,6 +20,8 @@ pub enum LexerErrorKind {
     IntegerLiteralTooLarge { span: Span, limit: String },
     #[error("{:?} is not a valid attribute", found)]
     MalformedFuncAttribute { span: Span, found: String },
+    #[error("Malformed test attribute")]
+    MalformedTestAttribute { span: Span },
     #[error("{:?} is not a valid inner attribute", found)]
     InvalidInnerAttribute { span: Span, found: String },
     #[error("Logical and used instead of bitwise and")]
@@ -61,6 +63,7 @@ impl LexerErrorKind {
             LexerErrorKind::InvalidIntegerLiteral { span, .. } => *span,
             LexerErrorKind::IntegerLiteralTooLarge { span, .. } => *span,
             LexerErrorKind::MalformedFuncAttribute { span, .. } => *span,
+            LexerErrorKind::MalformedTestAttribute { span, .. } => *span,
             LexerErrorKind::InvalidInnerAttribute { span, .. } => *span,
             LexerErrorKind::LogicalAnd { span } => *span,
             LexerErrorKind::UnterminatedBlockComment { span } => *span,
@@ -109,6 +112,11 @@ impl LexerErrorKind {
                 format!(" {found} is not a valid attribute"),
                 *span,
             ),
+            LexerErrorKind::MalformedTestAttribute { span } => (
+                "Malformed test attribute".to_string(),
+                "The test attribute can be written in one of these forms: `#[test]`, `#[test(should_fail)]` or `#[test(should_fail_with = \"message\")]`".to_string(),
+                *span,
+            ),
             LexerErrorKind::InvalidInnerAttribute { span, found } => (
                 "Invalid inner attribute".to_string(),
                 format!(" {found} is not a valid inner attribute"),
diff --git a/compiler/noirc_frontend/src/lexer/lexer.rs b/compiler/noirc_frontend/src/lexer/lexer.rs
index 904ce41fbf0..91ae544ddf0 100644
--- a/compiler/noirc_frontend/src/lexer/lexer.rs
+++ b/compiler/noirc_frontend/src/lexer/lexer.rs
@@ -935,13 +935,7 @@ mod tests {
             Err(err) => err,
         };
 
-        // Check if error is MalformedFuncAttribute and found is "foo"
-        let sub_string = match err {
-            LexerErrorKind::MalformedFuncAttribute { found, .. } => found,
-            _ => panic!("expected malformed func attribute error"),
-        };
-
-        assert_eq!(sub_string, "test(invalid_scope)");
+        assert!(matches!(err, LexerErrorKind::MalformedTestAttribute { .. }));
     }
 
     #[test]
diff --git a/compiler/noirc_frontend/src/lexer/token.rs b/compiler/noirc_frontend/src/lexer/token.rs
index daf59445982..593e78d01a0 100644
--- a/compiler/noirc_frontend/src/lexer/token.rs
+++ b/compiler/noirc_frontend/src/lexer/token.rs
@@ -751,10 +751,18 @@ impl Attribute {
         contents_span: Span,
         is_tag: bool,
     ) -> Result<Attribute, LexerErrorKind> {
-        let word_segments: Vec<&str> = word
-            .split(|c| c == '(' || c == ')')
-            .filter(|string_segment| !string_segment.is_empty())
-            .collect();
+        // See if we can parse the word into "name ( contents )".
+        // We first split into "first_segment ( rest".
+        let word_segments = if let Some((first_segment, rest)) = word.trim().split_once('(') {
+            // Now we try to remove the final ")" (it must be at the end, if it exists)
+            if let Some(middle) = rest.strip_suffix(')') {
+                vec![first_segment.trim(), middle.trim()]
+            } else {
+                vec![word]
+            }
+        } else {
+            vec![word]
+        };
 
         let validate = |slice: &str| {
             let is_valid = slice
@@ -799,11 +807,9 @@ impl Attribute {
             ["inline_always"] => Attribute::Function(FunctionAttribute::InlineAlways),
             ["test", name] => {
                 validate(name)?;
-                let malformed_scope =
-                    LexerErrorKind::MalformedFuncAttribute { span, found: word.to_owned() };
                 match TestScope::lookup_str(name) {
                     Some(scope) => Attribute::Function(FunctionAttribute::Test(scope)),
-                    None => return Err(malformed_scope),
+                    None => return Err(LexerErrorKind::MalformedTestAttribute { span }),
                 }
             }
             ["field", name] => {
diff --git a/tooling/lsp/src/requests/completion/builtins.rs b/tooling/lsp/src/requests/completion/builtins.rs
index 078e2faf036..c2c561ced32 100644
--- a/tooling/lsp/src/requests/completion/builtins.rs
+++ b/tooling/lsp/src/requests/completion/builtins.rs
@@ -107,6 +107,15 @@ impl<'a> NodeFinder<'a> {
                 let one_argument_attributes = &["abi", "field", "foreign", "oracle"];
                 self.suggest_one_argument_attributes(prefix, one_argument_attributes);
 
+                if name_matches("test", prefix) || name_matches("should_fail", prefix) {
+                    self.completion_items.push(snippet_completion_item(
+                        "test(should_fail)",
+                        CompletionItemKind::METHOD,
+                        "test(should_fail)",
+                        None,
+                    ));
+                }
+
                 if name_matches("test", prefix) || name_matches("should_fail_with", prefix) {
                     self.completion_items.push(snippet_completion_item(
                         "test(should_fail_with = \"...\")",