From 23f1c1ca9f1adfc705672557895c2391fc39782b Mon Sep 17 00:00:00 2001
From: Casey Rodarmor <casey@rodarmor.com>
Date: Thu, 27 Jun 2024 11:47:33 -0700
Subject: [PATCH] Allow comments after `mod` statements (#2201)

---
 src/parser.rs    |  5 +++--
 tests/modules.rs | 17 +++++++++++++++--
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/parser.rs b/src/parser.rs
index 9ea372c..9066e27 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -365,10 +365,11 @@ impl<'run, 'src> Parser<'run, 'src> {
             });
           }
           Some(Keyword::Mod)
-            if self.next_are(&[Identifier, Identifier, StringToken])
-              || self.next_are(&[Identifier, Identifier, Identifier, StringToken])
+            if self.next_are(&[Identifier, Identifier, Comment])
               || self.next_are(&[Identifier, Identifier, Eof])
               || self.next_are(&[Identifier, Identifier, Eol])
+              || self.next_are(&[Identifier, Identifier, Identifier, StringToken])
+              || self.next_are(&[Identifier, Identifier, StringToken])
               || self.next_are(&[Identifier, QuestionMark]) =>
           {
             self.presume_keyword(Keyword::Mod)?;
diff --git a/tests/modules.rs b/tests/modules.rs
index 43a1254..3692858 100644
--- a/tests/modules.rs
+++ b/tests/modules.rs
@@ -8,8 +8,6 @@ fn modules_are_unstable() {
         mod foo
       ",
     )
-    .arg("foo")
-    .arg("foo")
     .stderr(
       "error: Modules are currently unstable. \
       Invoke `just` with the `--unstable` flag to enable unstable features.\n",
@@ -781,3 +779,18 @@ fn colon_separated_path_components_are_not_used_as_arguments() {
     .status(1)
     .run();
 }
+
+#[test]
+fn comments_can_follow_modules() {
+  Test::new()
+    .write("foo.just", "foo:\n @echo FOO")
+    .justfile(
+      "
+        mod foo # this is foo
+      ",
+    )
+    .test_round_trip(false)
+    .args(["--unstable", "foo", "foo"])
+    .stdout("FOO\n")
+    .run();
+}