line number adjustment for more cases

core/src/main/scala/replpp/DottyReplDriver.scala

@@ -24,7 +24,7 @@ import dotty.tools.dotc.core.Symbols.{Symbol, defn}
 import dotty.tools.dotc.interfaces
 import dotty.tools.dotc.interactive.Completion
 import dotty.tools.dotc.printing.SyntaxHighlighting
-import dotty.tools.dotc.reporting.{ConsoleReporter, StoreReporter}
+import dotty.tools.dotc.reporting.{ConsoleReporter, Diagnostic, HideNonSensicalMessages, Reporter, StoreReporter, UniqueMessagePositions}
 import dotty.tools.dotc.reporting.Diagnostic
 import dotty.tools.dotc.util.Spans.Span
 import dotty.tools.dotc.util.{SourceFile, SourcePosition}
@@ -51,6 +51,7 @@ class DottyReplDriver(settings: Array[String],
                       out: PrintStream,
                       maxHeight: Option[Int],
                       classLoader: Option[ClassLoader])(using Colors) extends Driver:
+  private var invocationCount = 0
 
   /** Overridden to `false` in order to not have to give sources on the
    *  commandline
@@ -79,7 +80,8 @@ class DottyReplDriver(settings: Array[String],
   }
 
   /** the initial, empty state of the REPL session */
-  final def initialState: State = State(0, 0, Map.empty, Set.empty, rootCtx)
+  final def initialState: State =
+    State(0, 0, Map.empty, Set.empty, rootCtx)
 
   /** Reset state of repl to the initial state
    *

core/src/main/scala/replpp/InteractiveShell.scala

@@ -2,7 +2,6 @@ package replpp
 
 import dotty.tools.repl.State
 
-import java.lang.System.lineSeparator
 import scala.util.control.NoStackTrace
 
 object InteractiveShell {
@@ -19,24 +18,31 @@ object InteractiveShell {
       prompt = config0.prompt.getOrElse("scala"),
       maxHeight = config0.maxHeight,
       runAfter = config0.runAfter,
-      verbose = verbose,
+      verbose = verbose
     )
 
-    val initialState: State = replDriver.initialState
-    val runBeforeLines = (DefaultRunBeforeLines ++ globalRunBeforeLines ++ config.runBefore).mkString(lineSeparator)
-    val state: State = {
-      if (verbose) {
-        println(s"compiler arguments: ${compilerArgs.mkString(",")}")
-        println(runBeforeLines)
-        replDriver.run(runBeforeLines)(using initialState)
-      } else {
-        replDriver.runQuietly(runBeforeLines)(using initialState)
+    if (verbose) println(s"compiler arguments: ${compilerArgs.mkString(",")}")
+
+    var state: State = replDriver.initialState
+    var expectedStateObjectIndex = 0
+    Seq(DefaultRunBeforeLines, globalRunBeforeLines, config.runBefore).foreach { runBeforeLines =>
+      val runBeforeCode = runBeforeLines.mkString("\n").trim
+      if (runBeforeCode.nonEmpty) {
+        expectedStateObjectIndex += 1
+        state =
+          if (verbose) {
+            println(s"executing runBeforeCode: $runBeforeCode")
+            replDriver.run(runBeforeCode)(using state)
+          } else {
+            replDriver.runQuietly(runBeforeCode)(using state)
+          }
       }
     }
 
-    if (runBeforeLines.nonEmpty && state.objectIndex != 1) {
-      throw new AssertionError(s"compilation error for predef code - error should have been reported above ^") with NoStackTrace
-    }
+    assert(
+      state.objectIndex == expectedStateObjectIndex,
+      s"compilation error(s) for predef code - see error above ^^^"
+    )
 
     replDriver.runUntilQuit(using state)()
   }

core/src/main/scala/replpp/ReplDriver.scala

@@ -24,7 +24,7 @@ class ReplDriver(compilerArgs: Array[String],
   /** Run REPL with `state` until `:quit` command found
     * Main difference to the 'original': different greeting, trap Ctrl-c
    */
-  override def runUntilQuit(using initialState: State = initialState)(): State = {
+  override def runUntilQuit(using initialState: State)(): State = {
     val terminal = new replpp.JLineTerminal {
       override protected def promptStr = prompt
     }

core/src/main/scala/replpp/scripting/ScriptingDriver.scala

@@ -48,8 +48,10 @@ class ScriptingDriver(compilerArgs: Array[String],
     executed = true
     val inputFiles = (wrappedScript +: predefFiles).filter(Files.exists(_))
     try {
-      new SimpleDriver(lineNumberReportingAdjustment = -wrappingResult.linesBeforeWrappedCode)
-        .compile(compilerArgs, inputFiles, verbose) { (ctx, outDir) =>
+      new SimpleDriver(
+        linesBeforeRunBeforeCode = wrappingResult.linesBeforeRunBeforeCode,
+        linesBeforeScript = wrappingResult.linesBeforeScript
+      ).compile(compilerArgs, inputFiles, verbose) { (ctx, outDir) =>
           given Context = ctx
           tempFiles += outDir
 

core/src/main/scala/replpp/scripting/WrapForMainArgs.scala

@@ -4,51 +4,55 @@ package replpp.scripting
  * https://github.com/com-lihaoyi/mainargs  */
 object WrapForMainArgs {
 
-  /** linesBeforeWrappedCode: allows us to adjust line numbers in error reporting... */
-  case class WrappingResult(fullScript: String, linesBeforeWrappedCode: Int)
+  /** linesBeforeScript: allows us to adjust line numbers in error reporting... */
+  case class WrappingResult(fullScript: String, linesBeforeRunBeforeCode: Int, linesBeforeScript: Int)
 
   def apply(scriptCode: String, runBefore: Seq[String], runAfter: Seq[String]): WrappingResult = {
-    val runBeforeCode = runBefore.mkString("\n")
-    val runAfterCode = runAfter.mkString("\n")
-
-    val wrapperCodeStart =
+    val wrapperCodeStart0 =
       s"""import replpp.shaded.mainargs
          |import mainargs.main // intentionally shadow any potentially given @main
          |
          |// ScriptingDriver expects an object with a predefined name and a main entrypoint method
          |object ${ScriptingDriver.MainClassName} {
          |// runBeforeCode START
-         |$runBeforeCode
+         |""".stripMargin
+    val linesBeforeRunBeforeCode =
+      wrapperCodeStart0.lines().count().toInt
+      + 1 // for the line break after `wrapperCodeStart0`
+
+    val wrapperCodeStart1 =
+      s"""$wrapperCodeStart0
+         |${runBefore.mkString("\n")}
          |// runBeforeCode END
          |""".stripMargin
 
-    var linesBeforeWrappedCode = 0 // to adjust line number reporting
+    var linesBeforeScript = 0 // to adjust line number reporting
 
     val mainImpl =
       if (scriptCode.contains("@main"))
         scriptCode
       else {
-        linesBeforeWrappedCode += 1 // because we added the following line _before_ the wrapped script code
+        linesBeforeScript += 1 // because we added the following line _before_ the wrapped script code
         s"""@main def _execMain(): Unit = {
            |$scriptCode
            |}""".stripMargin
       }
 
-    linesBeforeWrappedCode += wrapperCodeStart.lines().count().toInt
-    linesBeforeWrappedCode += 1 // for the line break after $wrapperCodeStart
+    linesBeforeScript += wrapperCodeStart1.lines().count().toInt
+    linesBeforeScript += 1 // for the line break after `wrapperCodeStart1`
     val fullScript =
-      s"""$wrapperCodeStart
+      s"""$wrapperCodeStart1
          |$mainImpl
          |
          |  def ${ScriptingDriver.MainMethodName}(args: Array[String]): Unit = {
          |    mainargs.ParserForMethods(this).runOrExit(args.toSeq)
          |
-         |    $runAfterCode
+         |    ${runAfter.mkString("\n")}
          |  }
          |}
          |""".stripMargin
 
-    WrappingResult(fullScript, linesBeforeWrappedCode)
+    WrappingResult(fullScript, linesBeforeRunBeforeCode, linesBeforeScript)
   }
 
 }

core/src/main/scala/replpp/util/SimpleDriver.scala

@@ -11,6 +11,7 @@ import replpp.scripting.CompilerError
 import java.nio.file.{Files, Path}
 import scala.language.unsafeNulls
 import scala.util.Try
+import scala.util.control.NoStackTrace
 
 /** Compiles input files to a temporary directory
  *
@@ -23,7 +24,7 @@ import scala.util.Try
  * i.e. store hash of all inputs?
  * that functionality must exist somewhere already, e.g. zinc incremental compiler, or even in dotty itself?
  */
-class SimpleDriver(lineNumberReportingAdjustment: Int = 0) extends Driver {
+class SimpleDriver(linesBeforeRunBeforeCode: Int = 0, linesBeforeScript: Int = 0) extends Driver {
 
   def compileAndGetOutputDir[A](compilerArgs: Array[String], inputFiles: Seq[Path], verbose: Boolean): Try[Path] =
     compile(compilerArgs, inputFiles, verbose) { (ctx, outDir) => outDir }
@@ -45,7 +46,9 @@ class SimpleDriver(lineNumberReportingAdjustment: Int = 0) extends Driver {
 
       given ctx0: Context = {
         val ctx = rootCtx.fresh.setSetting(rootCtx.settings.outputDir, new PlainDirectory(Directory(outDir)))
-        if (lineNumberReportingAdjustment != 0) ctx.setReporter(createAdjustedReporter(rootCtx.reporter))
+        if (linesBeforeRunBeforeCode != 0 || linesBeforeScript != 0) {
+          ctx.setReporter(createReporter(linesBeforeRunBeforeCode, linesBeforeScript, rootCtx.reporter))
+        }
 
         if (verbose) {
           ctx.setSetting(rootCtx.settings.help, true)
@@ -58,19 +61,27 @@ class SimpleDriver(lineNumberReportingAdjustment: Int = 0) extends Driver {
 
       if (doCompile(newCompiler, toCompile).hasErrors) {
         val msgAddonMaybe = if (verbose) "" else " - try `--verbose` for more output"
-        throw CompilerError(s"Errors encountered during compilation$msgAddonMaybe")
+        throw new CompilerError(s"Errors encountered during compilation$msgAddonMaybe") with NoStackTrace
       } else {
         fun(ctx0, outDir)
       }
     }
   }
 
-  // creates a new reporter based on the original reporter that copies Diagnostic and changes line numbers
-  private def createAdjustedReporter(originalReporter: Reporter): Reporter = {
+  private def createReporter(linesBeforeRunBeforeCode: Int, linesBeforeScript: Int, originalReporter: Reporter): Reporter = {
     new Reporter {
       override def doReport(dia: Diagnostic)(using Context): Unit = {
         val adjustedPos = new SourcePosition(source = dia.pos.source, span = dia.pos.span, outer = dia.pos.outer) {
-          override def line: Int = super.line + lineNumberReportingAdjustment
+          override def line: Int = {
+            val original = super.line
+            val adjusted = original - linesBeforeScript
+            if (adjusted >= 0) {
+              adjusted
+            } else {
+              // adjusted line number is negative, i.e. the error must be in the `runBefore` code
+              original - linesBeforeRunBeforeCode
+            }
+          }
         }
         originalReporter.doReport(new Diagnostic(dia.msg, adjustedPos, dia.level))
       }

core/src/main/scala/replpp/util/package.scala

@@ -1,8 +1,7 @@
 package replpp
 
-import replpp.shaded.fansi
-
 import java.nio.file.{FileSystems, Files, Path}
+import replpp.shaded.fansi
 import scala.collection.immutable.Seq
 import scala.io.Source
 import scala.util.{Try, Using}
@@ -51,4 +50,5 @@ package object util {
 
   def pathAsString(path: Path): String =
     path.toAbsolutePath.toString
+
 }