From 26a607eb00261df2ccafc8fcaeb96d94c4ca1772 Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Tue, 30 May 2023 11:07:39 -0400
Subject: [PATCH] fix: bad parsing of comments at end of file closes #551

---
 CHANGELOG.md                          |  1 +
 crates/aiken-lang/src/parser/lexer.rs |  4 ++--
 crates/aiken-lang/src/tests/parser.rs | 20 ++++++++++++++++++++
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 97fdda7ba..cc12ffd0c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@
 - **uplc**: Fix pair formatting
 - **aiken-lang**: forced new line in formatter for assignments
 - **aiken-lang**: Incorrect parsing of generic type annotation prefixed with module
+- **aiken-lang**: Incorrect handling of comments at end of a file when newline not present
 
 ## v1.0.6-alpha - 2023-05-17
 
diff --git a/crates/aiken-lang/src/parser/lexer.rs b/crates/aiken-lang/src/parser/lexer.rs
index 3587c37fb..28202721b 100644
--- a/crates/aiken-lang/src/parser/lexer.rs
+++ b/crates/aiken-lang/src/parser/lexer.rs
@@ -148,13 +148,13 @@ pub fn lexer() -> impl Parser<char, Vec<(Token, Span)>, Error = ParseError> {
             // NOTE: The first case here work around a bug introduced with chumsky=0.9.0 which
             // miscalculate the offset for empty comments.
             just("/".repeat(n))
-                .ignore_then(text::newline().rewind())
+                .ignore_then(choice((text::newline().rewind(), end())))
                 .to(token.clone())
                 .map_with_span(move |token, span: Span| {
                     (token, Span::new((), span.start + n..span.end))
                 }),
             just("/".repeat(n)).ignore_then(
-                take_until(text::newline().rewind())
+                take_until(choice((text::newline().rewind(), end())))
                     .to(token)
                     .map_with_span(|token, span| (token, span)),
             ),
diff --git a/crates/aiken-lang/src/tests/parser.rs b/crates/aiken-lang/src/tests/parser.rs
index 605c61b06..2c2ff67b6 100644
--- a/crates/aiken-lang/src/tests/parser.rs
+++ b/crates/aiken-lang/src/tests/parser.rs
@@ -38,6 +38,26 @@ fn windows_newline() {
     )
 }
 
+#[test]
+fn can_handle_comments_at_end_of_file() {
+    let code = indoc! {r#"
+       use aiken
+
+       // some comment
+       // more comments"#};
+
+    assert_definitions(
+        code,
+        vec![ast::UntypedDefinition::Use(Use {
+            location: Span::new((), 0..9),
+            module: vec!["aiken".to_string()],
+            as_name: None,
+            unqualified: vec![],
+            package: (),
+        })],
+    )
+}
+
 #[test]
 fn type_annotation_with_module_prefix() {
     let code = indoc! {r#"