From 94b3af6cb7a1175c8f65b7159a5998281b9219b6 Mon Sep 17 00:00:00 2001
From: Casey Rodarmor <casey@rodarmor.com>
Date: Thu, 28 Dec 2023 04:23:58 -0800
Subject: [PATCH] Allow mod statements with path to source file (#1786)

---
 src/analyzer.rs  |  2 +-
 src/compiler.rs  | 12 ++++++++++--
 src/item.rs      | 11 ++++++++++-
 src/parser.rs    | 17 +++++++++++++++--
 tests/modules.rs | 49 ++++++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 85 insertions(+), 6 deletions(-)

diff --git a/src/analyzer.rs b/src/analyzer.rs
index 0f138a7c3c..99c3db21d8 100644
--- a/src/analyzer.rs
+++ b/src/analyzer.rs
@@ -77,7 +77,7 @@ impl<'src> Analyzer<'src> {
           Item::Import { absolute, .. } => {
             stack.push(asts.get(absolute.as_ref().unwrap()).unwrap());
           }
-          Item::Mod { absolute, name } => {
+          Item::Mod { absolute, name, .. } => {
             define(*name, "module", false)?;
             modules.insert(
               name.to_string(),
diff --git a/src/compiler.rs b/src/compiler.rs
index b7b0eca242..457e889486 100644
--- a/src/compiler.rs
+++ b/src/compiler.rs
@@ -27,7 +27,11 @@ impl Compiler {
 
       for item in &mut ast.items {
         match item {
-          Item::Mod { name, absolute } => {
+          Item::Mod {
+            name,
+            absolute,
+            path,
+          } => {
             if !unstable {
               return Err(Error::Unstable {
                 message: "Modules are currently unstable.".into(),
@@ -36,7 +40,11 @@ impl Compiler {
 
             let parent = current.parent().unwrap();
 
-            let import = Self::find_module_file(parent, *name)?;
+            let import = if let Some(path) = path {
+              parent.join(&path.cooked)
+            } else {
+              Self::find_module_file(parent, *name)?
+            };
 
             if srcs.contains_key(&import) {
               return Err(Error::CircularImport { current, import });
diff --git a/src/item.rs b/src/item.rs
index a709e3606f..654b6bae34 100644
--- a/src/item.rs
+++ b/src/item.rs
@@ -13,6 +13,7 @@ pub(crate) enum Item<'src> {
   Mod {
     name: Name<'src>,
     absolute: Option<PathBuf>,
+    path: Option<StringLiteral<'src>>,
   },
   Recipe(UnresolvedRecipe<'src>),
   Set(Set<'src>),
@@ -25,7 +26,15 @@ impl<'src> Display for Item<'src> {
       Item::Assignment(assignment) => write!(f, "{assignment}"),
       Item::Comment(comment) => write!(f, "{comment}"),
       Item::Import { relative, .. } => write!(f, "import {relative}"),
-      Item::Mod { name, .. } => write!(f, "mod {name}"),
+      Item::Mod { name, path, .. } => {
+        write!(f, "mod {name}")?;
+
+        if let Some(path) = path {
+          write!(f, " {path}")?;
+        }
+
+        Ok(())
+      }
       Item::Recipe(recipe) => write!(f, "{}", recipe.color_display(Color::never())),
       Item::Set(set) => write!(f, "{set}"),
     }
diff --git a/src/parser.rs b/src/parser.rs
index 019f564e67..3758988507 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -335,11 +335,24 @@ impl<'tokens, 'src> Parser<'tokens, 'src> {
               absolute: None,
             });
           }
-          Some(Keyword::Mod) if self.next_are(&[Identifier, Identifier]) => {
+          Some(Keyword::Mod)
+            if self.next_are(&[Identifier, Identifier, StringToken])
+              || self.next_are(&[Identifier, Identifier, Eof])
+              || self.next_are(&[Identifier, Identifier, Eol]) =>
+          {
             self.presume_keyword(Keyword::Mod)?;
+            let name = self.parse_name()?;
+
+            let path = if self.next_is(StringToken) {
+              Some(self.parse_string_literal()?)
+            } else {
+              None
+            };
+
             items.push(Item::Mod {
-              name: self.parse_name()?,
+              name,
               absolute: None,
+              path,
             });
           }
           Some(Keyword::Set)
diff --git a/tests/modules.rs b/tests/modules.rs
index 59d356446a..ca69340faa 100644
--- a/tests/modules.rs
+++ b/tests/modules.rs
@@ -444,3 +444,52 @@ fn dotenv_settings_in_submodule_are_ignored() {
     .stdout("dotenv-value\n")
     .run();
 }
+
+#[test]
+fn modules_may_specify_path() {
+  Test::new()
+    .write("commands/foo.just", "foo:\n @echo FOO")
+    .justfile(
+      "
+        mod foo 'commands/foo.just'
+      ",
+    )
+    .test_round_trip(false)
+    .arg("--unstable")
+    .arg("foo")
+    .arg("foo")
+    .stdout("FOO\n")
+    .run();
+}
+
+#[test]
+fn modules_with_paths_are_dumped_correctly() {
+  Test::new()
+    .write("commands/foo.just", "foo:\n @echo FOO")
+    .justfile(
+      "
+        mod foo 'commands/foo.just'
+      ",
+    )
+    .test_round_trip(false)
+    .arg("--unstable")
+    .arg("--dump")
+    .stdout("mod foo 'commands/foo.just'\n")
+    .run();
+}
+
+#[test]
+fn recipes_may_be_named_mod() {
+  Test::new()
+    .justfile(
+      "
+        mod foo:
+          @echo FOO
+      ",
+    )
+    .test_round_trip(false)
+    .arg("mod")
+    .arg("bar")
+    .stdout("FOO\n")
+    .run();
+}