From d5ba564135692f8bcd74647f03e648d1665d044e Mon Sep 17 00:00:00 2001
From: regginator <reggie@latte.to>
Date: Wed, 4 Oct 2023 21:00:22 -0500
Subject: [PATCH] Make all standalone modules pass-through Wax's env

---
 README.md                   |  2 +-
 default.project.json        |  2 +-
 lune/lib/data/Template.luau | 17 ++++++++++++-----
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 6b9daf8..a1dfa23 100644
--- a/README.md
+++ b/README.md
@@ -171,7 +171,7 @@ SUBCOMMANDS:
 
 ## Runtime Spec Details
 
-* Like a model on Roblox's [Developer Marketplace](https://create.roblox.com/marketplace/models), a ModuleScript under the model root named (exactly) "MainModule" will automatically run at init and pass its return through the **real** script
+* If there's only 1 instance directly under the model root and it's a `ModuleScript`, that module will automatically run at init and pass its return through the **real** script, just as expected with normal module behavior. Additionally, like a model on Roblox's [Developer Marketplace](https://create.roblox.com/marketplace/models), a `ModuleScript` under the model root named (exactly) "MainModule" will also have the same functionality, even if there is more than 1 instance directly under the model root
 * The real `getfenv`/`setfenv` functions are not overriden by global flattening whatsoever, so by using for example `getfenv(0)` it will return the actual root function environment of the script. *Just know that its behavior isn't modified whatsoever for virtual closures!*
 * Instance `ClassName`s that will bundle:
   * `Folder`
diff --git a/default.project.json b/default.project.json
index 9ca99d5..b56b4d0 100644
--- a/default.project.json
+++ b/default.project.json
@@ -1,5 +1,5 @@
 {
-    "name": "MainModule",
+    "name": "Wax",
     "tree": {
         "$path": "lune/wax.luau",
 
diff --git a/lune/lib/data/Template.luau b/lune/lib/data/Template.luau
index 64a17cb..85d72b4 100644
--- a/lune/lib/data/Template.luau
+++ b/lune/lib/data/Template.luau
@@ -505,11 +505,18 @@ end
 
 -- If there's a "MainModule" top-level modulescript, we'll return it from the output's closure directly
 do
-    local MainModule 
-    for _, Ref in next, RealObjectRoot:GetChildren() do
-        if Ref.ClassName == "ModuleScript" and Ref.Name == "MainModule" then
-            MainModule = Ref
-            break
+    local MainModule
+    local RealObjectRootChildren = RealObjectRoot:GetChildren()
+
+    -- Doesn't need to be named "MainModule" if there's just 1 root obj and it's a ModuleScript
+    if #RealObjectRootChildren == 1 and RealObjectRootChildren[1].ClassName == "ModuleScript" then
+        MainModule = RealObjectRootChildren[1]
+    else
+        for _, Ref in next, RealObjectRootChildren do
+            if Ref.ClassName == "ModuleScript" and Ref.Name == "MainModule" then
+                MainModule = Ref
+                break
+            end
         end
     end