Greatly simplify `CustomCodeWrapper`

3 files changed, 15 insertions, 36 deletions

diff --git a/core/src/main/scala/mill/Main.scala b/core/src/main/scala/mill/Main.scala
index c6d05577..0d5d282e 100644
--- a/core/src/main/scala/mill/Main.scala
+++ b/core/src/main/scala/mill/Main.scala
@@ -47,9 +47,8 @@ object Main {
         System.exit(1)
       case Right((cliConfig, leftoverArgs)) =>
         val config =
-          if(!repl) cliConfig.copy(defaultPredef = false)
+          if(!repl) cliConfig
           else cliConfig.copy(
-            defaultPredef = false,
             predefCode = "implicit val replApplyHandler = mill.main.ReplApplyHandler(mapping)",
             predefFile = Some(pwd/"build.sc"),
             welcomeBanner = None
diff --git a/core/src/main/scala/mill/main/CustomCodeWrapper.scala b/core/src/main/scala/mill/main/CustomCodeWrapper.scala
index 3d840128..f92850f4 100644
--- a/core/src/main/scala/mill/main/CustomCodeWrapper.scala
+++ b/core/src/main/scala/mill/main/CustomCodeWrapper.scala
@@ -5,49 +5,26 @@ import ammonite.util.{Imports, Name, Util}
 
 object CustomCodeWrapper extends Preprocessor.CodeWrapper {
   def top(pkgName: Seq[Name], imports: Imports, indexedWrapperName: Name) = {
+    val wrapName = indexedWrapperName.backticked
     s"""
        |package ${pkgName.head.encoded}
        |package ${Util.encodeScalaSourcePath(pkgName.tail)}
        |$imports
        |import mill._
-       |sealed abstract class ${indexedWrapperName.backticked} extends mill.Module{\n
+       |
+       |object $wrapName extends $wrapName with mill.main.MainWrapper[$wrapName]{
+       |  lazy val discovered = mill.discover.Discovered.make[$wrapName]
+       |}
+       |
+       |sealed abstract class $wrapName extends mill.Module{
        |""".stripMargin
   }
 
 
   def bottom(printCode: String, indexedWrapperName: Name, extraCode: String) = {
-    val wrapName = indexedWrapperName.backticked
-    val tmpName = ammonite.util.Name(indexedWrapperName.raw + "-Temp").backticked
-
-    // Define `discovered` in the `tmpName` trait, before mixing in `MainWrapper`,
-    // to ensure that `$tempName#discovered` is initialized before `MainWrapper` is.
-    //
-    // `import $wrapName._` is necessary too let Ammonite pick up all the
-    // members of class wrapper, which are inherited but otherwise not visible
-    // in the AST of the `$wrapName` object
-    //
-    // We need to duplicate the Ammonite predef as part of the wrapper because
-    // imports within the body of the class wrapper are not brought into scope
-    // by the `import $wrapName._`. Other non-Ammonite-predef imports are not
-    // made available, and that's just too bad
-    s"""
-       |}
-       |trait $tmpName{
-       |  val discovered = mill.discover.Discovered.make[$wrapName]
-       |  val interpApi = ammonite.interp.InterpBridge.value
-       |}
-       |
-       |object $wrapName
-       |extends $wrapName
-       |with $tmpName
-       |with mill.main.MainWrapper[$wrapName] {
-       |  ${ammonite.main.Defaults.replPredef}
-       |  ${ammonite.main.Defaults.predefString}
-       |  ${ammonite.Main.extraPredefString}
-       |  import ammonite.repl.ReplBridge.{value => repl}
-       |  import ammonite.interp.InterpBridge.{value => interp}
-       |  import $wrapName._
-      """.stripMargin +
-      Preprocessor.CodeWrapper.bottom(printCode, indexedWrapperName, extraCode)
+    // We need to disable the `$main` method definition inside the wrapper
+    // class, because otherwise it might get picked up by Ammonite and run as
+    // a static class, which naturally blows up
+    "\n}"
   }
 }
diff --git a/core/src/main/scala/mill/main/MainWrapper.scala b/core/src/main/scala/mill/main/MainWrapper.scala
index da8fc55a..16fd58c8 100644
--- a/core/src/main/scala/mill/main/MainWrapper.scala
+++ b/core/src/main/scala/mill/main/MainWrapper.scala
@@ -5,5 +5,8 @@ package mill.main
   */
 trait MainWrapper[T]{
   val discovered: mill.discover.Discovered[T]
+  // Stub to make sure Ammonite has something to call after it evaluates a script,
+  // even if it does nothing...
+  def $main() = Iterator[String]()
   lazy val mapping = discovered.mapping(this.asInstanceOf[T])
 }