From 660ca3fadab022c90a49b716e1e506ed087df1b0 Mon Sep 17 00:00:00 2001
From: KtorZ <matthias.benkort@gmail.com>
Date: Thu, 16 Mar 2023 23:33:02 +0100
Subject: [PATCH] Fix comment formatting wrongly assuming false invariant.

---
 crates/aiken-lang/src/format.rs       | 40 +++++++++++++++------------
 crates/aiken-lang/src/tests/check.rs  | 19 -------------
 crates/aiken-lang/src/tests/format.rs | 21 ++++++++++++--
 crates/aiken-lang/src/tipo/error.rs   |  7 -----
 4 files changed, 41 insertions(+), 46 deletions(-)

diff --git a/crates/aiken-lang/src/format.rs b/crates/aiken-lang/src/format.rs
index 56193e228..c8bdf75f4 100644
--- a/crates/aiken-lang/src/format.rs
+++ b/crates/aiken-lang/src/format.rs
@@ -1799,31 +1799,35 @@ fn printed_comments<'a, 'comments>(
 
     let mut doc = Vec::new();
     while let Some(c) = comments.next() {
-        // There will never be consecutive empty lines (None values),
-        // and whenever we peek a None, we advance past it.
-        let c = c.expect("no consecutive empty lines");
-        doc.push("//".to_doc().append(Document::String(c.to_string())));
-        match comments.peek() {
-            // Next line is a comment
-            Some(Some(_)) => doc.push(line()),
-            // Next line is empty
-            Some(None) => {
-                comments.next();
+        match c {
+            None => continue,
+            Some(c) => {
+                // There will never be consecutive empty lines (None values),
+                // and whenever we peek a None, we advance past it.
+                doc.push("//".to_doc().append(Document::String(c.to_string())));
                 match comments.peek() {
-                    Some(_) => doc.push(lines(2)),
+                    // Next line is a comment
+                    Some(Some(_)) => doc.push(line()),
+                    // Next line is empty
+                    Some(None) => {
+                        comments.next();
+                        match comments.peek() {
+                            Some(_) => doc.push(lines(2)),
+                            None => {
+                                if trailing_newline {
+                                    doc.push(lines(2));
+                                }
+                            }
+                        }
+                    }
+                    // We've reached the end, there are no more lines
                     None => {
                         if trailing_newline {
-                            doc.push(lines(2));
+                            doc.push(line());
                         }
                     }
                 }
             }
-            // We've reached the end, there are no more lines
-            None => {
-                if trailing_newline {
-                    doc.push(line());
-                }
-            }
         }
     }
     let doc = concat(doc);
diff --git a/crates/aiken-lang/src/tests/check.rs b/crates/aiken-lang/src/tests/check.rs
index 92cb10e34..0bf369a74 100644
--- a/crates/aiken-lang/src/tests/check.rs
+++ b/crates/aiken-lang/src/tests/check.rs
@@ -322,25 +322,6 @@ fn utf8_hex_literal_warning() {
     ))
 }
 
-#[test]
-fn weird_comments() {
-    let source_code = r#"
-        // A
-
-        /// B
-
-        // C
-        fn foo () {
-          todo
-        }
-    "#;
-
-    assert!(matches!(
-        check(parse(source_code)),
-        Err((_, Error::UnexpectedDocComment { .. }))
-    ))
-}
-
 #[test]
 fn discarded_let_bindings() {
     let source_code = r#"
diff --git a/crates/aiken-lang/src/tests/format.rs b/crates/aiken-lang/src/tests/format.rs
index a71581d6a..4901936ca 100644
--- a/crates/aiken-lang/src/tests/format.rs
+++ b/crates/aiken-lang/src/tests/format.rs
@@ -638,7 +638,16 @@ fn weird_comments() {
         /// B
 
         /// C
-        fn bar () {
+        fn foo() {
+          todo
+        }
+
+        // E
+
+        /// F
+
+        // G
+        fn bar() {
           todo
         }
     "#};
@@ -648,7 +657,15 @@ fn weird_comments() {
 
         /// B
         /// C
-        fn bar () {
+        fn foo() {
+          todo
+        }
+
+        // E
+
+        // G
+        /// F
+        fn bar() {
           todo
         }
     "#};
diff --git a/crates/aiken-lang/src/tipo/error.rs b/crates/aiken-lang/src/tipo/error.rs
index 02cd2a812..3486256ab 100644
--- a/crates/aiken-lang/src/tipo/error.rs
+++ b/crates/aiken-lang/src/tipo/error.rs
@@ -892,13 +892,6 @@ The best thing to do from here is to remove it."#))]
         #[label("{} arguments", if *count < 2 { "not enough" } else { "too many" })]
         location: Span,
     },
-
-    #[error("I ran into a lonely doc comment not attached to any function definition.")]
-    #[diagnostic(code("unexpected::doc_comment"))]
-    UnexpectedDocComment {
-        #[label("unexpected doc comment")]
-        location: Span,
-    },
 }
 
 impl Error {