From d106e62228392b2f8d926db86126ee456da7fd22 Mon Sep 17 00:00:00 2001
From: xunilrj <xunilrj@hotmail.com>
Date: Thu, 17 Aug 2023 18:05:17 +0100
Subject: [PATCH] improvements to unfinished comments inside fns

---
 sway-parse/src/expr/mod.rs                    | 25 ++++++------
 sway-parse/src/parser.rs                      | 38 +++++++++++++++++++
 .../src/main.sw                               |  5 +++
 3 files changed, 57 insertions(+), 11 deletions(-)

diff --git a/sway-parse/src/expr/mod.rs b/sway-parse/src/expr/mod.rs
index 1869bb74625..16cb7cc951d 100644
--- a/sway-parse/src/expr/mod.rs
+++ b/sway-parse/src/expr/mod.rs
@@ -132,20 +132,29 @@ impl ParseToEnd for CodeBlockContents {
         mut parser: Parser<'a, '_>,
     ) -> ParseResult<(CodeBlockContents, ParserConsumed<'a>)> {
         let mut statements = Vec::new();
+
         let (final_expr_opt, consumed) = loop {
             if let Some(consumed) = parser.check_empty() {
                 break (None, consumed);
             }
-            match parse_stmt(&mut parser)? {
-                StmtOrTail::Stmt(s) => statements.push(s),
-                StmtOrTail::Tail(e, c) => break (Some(e), c),
+
+            match parser.parse_fn_with_recovery(parse_stmt) {
+                Ok(StmtOrTail::Stmt(s)) => statements.push(s),
+                Ok(StmtOrTail::Tail(e, c)) => break (Some(e), c),
+                Err(r) => {
+                    let (spans, error) = r
+                        .recover_at_next_line_with_fallback_error(ParseErrorKind::InvalidStatement);
+                    statements.push(Statement::Error(spans, error));
+                }
             }
         };
+
         let code_block_contents = CodeBlockContents {
             statements,
             final_expr_opt,
             span: parser.full_span().clone(),
         };
+
         Ok((code_block_contents, consumed))
     }
 }
@@ -187,14 +196,8 @@ fn parse_stmt<'a>(parser: &mut Parser<'a, '_>) -> ParseResult<StmtOrTail<'a>> {
     }
 
     // Try a `let` statement.
-    match parser.guarded_parse_with_recovery::<LetToken, StatementLet>() {
-        Ok(None) => {}
-        Ok(Some(item)) => return stmt(Statement::Let(item)),
-        Err(r) => {
-            let (spans, error) =
-                r.recover_at_next_line_with_fallback_error(ParseErrorKind::InvalidStatement);
-            return stmt(Statement::Error(spans, error));
-        }
+    if let Some(item) = parser.guarded_parse::<LetToken, StatementLet>()? {
+        return stmt(Statement::Let(item));
     }
 
     // Try an `expr;` statement.
diff --git a/sway-parse/src/parser.rs b/sway-parse/src/parser.rs
index e420c983c83..44a3b43f5e1 100644
--- a/sway-parse/src/parser.rs
+++ b/sway-parse/src/parser.rs
@@ -76,6 +76,44 @@ impl<'a, 'e> Parser<'a, 'e> {
         Peeker::with(self.token_trees).map(|(v, _)| v)
     }
 
+    pub fn parse_fn_with_recovery<
+        'original,
+        T,
+        F: FnOnce(&mut Parser<'a, '_>) -> ParseResult<T>,
+    >(
+        &'original mut self,
+        f: F,
+    ) -> Result<T, Recoverer<'original, 'a, 'e>> {
+        let handler = Handler::default();
+        let mut fork = Parser {
+            token_trees: self.token_trees,
+            full_span: self.full_span.clone(),
+            handler: &handler,
+        };
+
+        match f(&mut fork) {
+            Ok(result) => {
+                self.token_trees = fork.token_trees;
+                self.handler.append(handler);
+                Ok(result)
+            }
+            Err(error) => {
+                let Parser {
+                    token_trees,
+                    full_span,
+                    ..
+                } = fork;
+                Err(Recoverer {
+                    original: RefCell::new(self),
+                    handler,
+                    fork_token_trees: token_trees,
+                    fork_full_span: full_span,
+                    error,
+                })
+            }
+        }
+    }
+
     pub fn parse_with_recovery<'original, T: Parse>(
         &'original mut self,
     ) -> Result<T, Recoverer<'original, 'a, 'e>> {
diff --git a/test/src/e2e_vm_tests/test_programs/should_fail/recover_unclosed_multiline_comment/src/main.sw b/test/src/e2e_vm_tests/test_programs/should_fail/recover_unclosed_multiline_comment/src/main.sw
index 456f402996d..968b2586f1e 100644
--- a/test/src/e2e_vm_tests/test_programs/should_fail/recover_unclosed_multiline_comment/src/main.sw
+++ b/test/src/e2e_vm_tests/test_programs/should_fail/recover_unclosed_multiline_comment/src/main.sw
@@ -7,6 +7,11 @@ mod foo;
 mod bar;
 mod baz;
 
+fn f() {
+    /
+    f2();
+}
+
 fn a() -> bool { 0 } // recovery witness
 
 /*