From 244345cd6b07dd15582be583c63fa2386e09b1cf Mon Sep 17 00:00:00 2001
From: Maksym Hazevych <mhazevych@proton.me>
Date: Sat, 13 Jul 2024 02:09:14 +0300
Subject: [PATCH] feat: Consider functions to be failable when marked as such

---
 src/modules/function/declaration.rs | 16 +++++++---------
 src/tests/validity.rs               |  4 +---
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/src/modules/function/declaration.rs b/src/modules/function/declaration.rs
index a951908..a8c6c31 100644
--- a/src/modules/function/declaration.rs
+++ b/src/modules/function/declaration.rs
@@ -147,15 +147,13 @@ impl SyntaxModule<ParserMetadata> for FunctionDeclaration {
             }
 
             let mut type_tok = None;
-            // Semantic meaning: ValueWasSet<IsThere<aValue>>
-            let mut failable_mark: Option<Option<Option<Token>>> = None;
+            let mut marked_as_failable: Option<bool> = None;
             // Optionally parse the return type
             match token(meta, ":") {
                 Ok(_) => {
                     type_tok = Some(meta.get_current_token());
                     self.returns = parse_type(meta)?;
-                    let mark = meta.get_current_token();
-                    failable_mark = Some(token(meta, "?").ok().map(|_| mark));
+                    marked_as_failable = Some(token(meta, "?").is_ok());
                 },
                 Err(_) => self.returns = Type::Generic
             }
@@ -163,12 +161,12 @@ impl SyntaxModule<ParserMetadata> for FunctionDeclaration {
 
             // Parse the body
             token(meta, "{")?;
-            let (index_begin, index_end, is_failable) = skip_function_body(meta);
-            if let (Some(failable_mark), Some(type_tok)) = (failable_mark, type_tok) {
-                if is_failable && failable_mark.is_none() {
+            let (index_begin, index_end, mut is_failable) = skip_function_body(meta);
+            if let (Some(marked_as_failable), Some(type_tok)) = (marked_as_failable, type_tok) {
+                if is_failable && !marked_as_failable {
                     return error!(meta, type_tok, "Failable functions must be declared as such");
-                } else if !is_failable && failable_mark.is_some() {
-                    return error!(meta, failable_mark.unwrap(), "Non-failable functions cannot be marked as failable");
+                } else if !is_failable && marked_as_failable {
+                    is_failable = true;
                 }
             }
             // Create a new context with the function body
diff --git a/src/tests/validity.rs b/src/tests/validity.rs
index e981c5d..626b303 100644
--- a/src/tests/validity.rs
+++ b/src/tests/validity.rs
@@ -548,8 +548,7 @@ fn failable_function_with_typed_return_without_mark() {
 }
 
 #[test]
-#[should_panic(expected = "ERROR: Non-failable functions cannot be marked as failable")]
-fn non_failable_function_cannot_be_marked() {
+fn non_failable_function_marked_as_failable() {
     let code = r#"
         pub fun test(): Null ? {
             echo "Hello, World!"
@@ -1260,4 +1259,3 @@ fn optional_argument_generic(){
     ";
     test_amber!(code, "hello")
 }
-