SI-4625 Warn on first non-toplevel only

Fixed the warning when main module is accompanied by snippets.

Minor clean-up so even I can follow what is returned.

1 files changed, 46 insertions, 42 deletions

diff --git a/src/compiler/scala/tools/nsc/ast/parser/Parsers.scala b/src/compiler/scala/tools/nsc/ast/parser/Parsers.scala
index c2f2141fd3..308669256d 100644
--- a/src/compiler/scala/tools/nsc/ast/parser/Parsers.scala
+++ b/src/compiler/scala/tools/nsc/ast/parser/Parsers.scala
@@ -371,7 +371,7 @@ self =>
        * around it.  Since objects are loaded lazily the whole script would have
        * been a no-op, so we're not taking much liberty.
        */
-      def searchForMain(): Option[Tree] = {
+      def searchForMain(): Tree = {
         import PartialFunction.cond
 
         /* Have to be fairly liberal about what constitutes a main method since
@@ -387,10 +387,10 @@ self =>
           case Template(parents, _, _) => parents.exists(cond(_) { case Ident(tpnme.App) => true })
           case _ => false
         }
-        /* For now we require there only be one top level object. */
+        /* We allow only one main module. */
         var seenModule = false
         var disallowed = EmptyTree: Tree
-        val newStmts = stmts collect {
+        val newStmts = stmts.map {
           case md @ ModuleDef(mods, name, template) if !seenModule && (isApp(template) || md.exists(isMainMethod)) =>
             seenModule = true
             /* This slightly hacky situation arises because we have no way to communicate
@@ -407,54 +407,58 @@ self =>
           case t  @ Import(_, _)         => t
           case t =>
             /* If we see anything but the above, fail. */
-            disallowed = t
+            if (disallowed.isEmpty) disallowed = t
             EmptyTree
         }
-        if (disallowed.isEmpty) Some(makeEmptyPackage(0, newStmts))
+        if (disallowed.isEmpty) makeEmptyPackage(0, newStmts)
         else {
           if (seenModule)
             warning(disallowed.pos.point, "Script has a main object but statement is disallowed")
-          None
+          EmptyTree
         }
       }
 
-      if (mainModuleName == newTermName(ScriptRunner.defaultScriptMain))
-        searchForMain() foreach { return _ }
+      def mainModule: Tree =
+        if (mainModuleName == newTermName(ScriptRunner.defaultScriptMain)) searchForMain() else EmptyTree
 
-      /*  Here we are building an AST representing the following source fiction,
-       *  where `moduleName` is from -Xscript (defaults to "Main") and <stmts> are
-       *  the result of parsing the script file.
-       *
-       *  {{{
-       *  object moduleName {
-       *    def main(args: Array[String]): Unit =
-       *      new AnyRef {
-       *        stmts
-       *      }
-       *  }
-       *  }}}
-       */
-      def emptyInit   = DefDef(
-        NoMods,
-        nme.CONSTRUCTOR,
-        Nil,
-        ListOfNil,
-        TypeTree(),
-        Block(List(Apply(gen.mkSuperInitCall, Nil)), literalUnit)
-      )
-
-      // def main
-      def mainParamType = AppliedTypeTree(Ident(tpnme.Array), List(Ident(tpnme.String)))
-      def mainParameter = List(ValDef(Modifiers(Flags.PARAM), nme.args, mainParamType, EmptyTree))
-      def mainDef       = DefDef(NoMods, nme.main, Nil, List(mainParameter), scalaDot(tpnme.Unit), gen.mkAnonymousNew(stmts))
-
-      // object Main
-      def moduleName  = newTermName(ScriptRunner scriptMain settings)
-      def moduleBody  = Template(atInPos(scalaAnyRefConstr) :: Nil, noSelfType, List(emptyInit, mainDef))
-      def moduleDef   = ModuleDef(NoMods, moduleName, moduleBody)
-
-      // package <empty> { ... }
-      makeEmptyPackage(0, moduleDef :: Nil)
+      def repackaged: Tree = {
+        /*  Here we are building an AST representing the following source fiction,
+         *  where `moduleName` is from -Xscript (defaults to "Main") and <stmts> are
+         *  the result of parsing the script file.
+         *
+         *  {{{
+         *  object moduleName {
+         *    def main(args: Array[String]): Unit =
+         *      new AnyRef {
+         *        stmts
+         *      }
+         *  }
+         *  }}}
+         */
+        def emptyInit   = DefDef(
+          NoMods,
+          nme.CONSTRUCTOR,
+          Nil,
+          ListOfNil,
+          TypeTree(),
+          Block(List(Apply(gen.mkSuperInitCall, Nil)), literalUnit)
+        )
+
+        // def main
+        def mainParamType = AppliedTypeTree(Ident(tpnme.Array), List(Ident(tpnme.String)))
+        def mainParameter = List(ValDef(Modifiers(Flags.PARAM), nme.args, mainParamType, EmptyTree))
+        def mainDef       = DefDef(NoMods, nme.main, Nil, List(mainParameter), scalaDot(tpnme.Unit), gen.mkAnonymousNew(stmts))
+
+        // object Main
+        def moduleName  = newTermName(ScriptRunner scriptMain settings)
+        def moduleBody  = Template(atInPos(scalaAnyRefConstr) :: Nil, noSelfType, List(emptyInit, mainDef))
+        def moduleDef   = ModuleDef(NoMods, moduleName, moduleBody)
+
+        // package <empty> { ... }
+        makeEmptyPackage(0, moduleDef :: Nil)
+      }
+
+      mainModule orElse repackaged
     }
 
 /* --------------- PLACEHOLDERS ------------------------------------------- */