Reposition span to original file for wrapper scripts

Author: Linyxus

compiler/src/dotty/tools/dotc/ast/Positioned.scala

@@ -3,7 +3,8 @@ package dotc
 package ast
 
 import util.Spans.*
-import util.{SourceFile, SourcePosition, SrcPos}
+import util.{SourceFile, SourcePosition, SrcPos, WrappedSourceFile}
+import WrappedSourceFile.MagicHeaderInfo, MagicHeaderInfo.*
 import core.Contexts.*
 import core.Decorators.*
 import core.NameOps.*
@@ -51,7 +52,12 @@ abstract class Positioned(implicit @constructorOnly src: SourceFile) extends Src
 
   def source: SourceFile = mySource
 
-  def sourcePos(using Context): SourcePosition = source.atSpan(span)
+  def sourcePos(using Context): SourcePosition =
+    val info = WrappedSourceFile.locateMagicHeader(source)
+    info match
+      case HasHeader(offset, originalFile) =>
+        originalFile.atSpan(span `shift` -offset)
+      case _ => source.atSpan(span)
 
   /** This positioned item, widened to `SrcPos`. Used to make clear we only need the
    *  position, typically for error reporting.

compiler/src/dotty/tools/dotc/config/ScalaSettings.scala

@@ -443,7 +443,7 @@ private sealed trait YSettings:
   val YbestEffort: Setting[Boolean] = BooleanSetting(ForkSetting, "Ybest-effort", "Enable best-effort compilation attempting to produce betasty to the META-INF/best-effort directory, regardless of errors, as part of the pickler phase.")
   val YwithBestEffortTasty: Setting[Boolean] = BooleanSetting(ForkSetting, "Ywith-best-effort-tasty", "Allow to compile using best-effort tasty files. If such file is used, the compiler will stop after the pickler phase.")
 
-  val YmagicOffsetHeader: Setting[String] = StringSetting(AdvancedSetting, "Ymagic-offset-header", "header", "Specify the magic header comment that marks the start of the actual code in generated wrapper scripts. Example: -Ymagic-offset-header:SOURCE_CODE_START. Then, in the source, the magic comment `///SOURCE_CODE_START` marks the start of user code. The comment can be optionally suffixed by `:<ORIGINAL_FILE_PATH>` to indicate the original file.", "")
+  val YmagicOffsetHeader: Setting[String] = StringSetting(ForkSetting, "Ymagic-offset-header", "header", "Specify the magic header comment that marks the start of the actual code in generated wrapper scripts. Example: -Ymagic-offset-header:SOURCE_CODE_START. Then, in the source, the magic comment `///SOURCE_CODE_START:<ORIGINAL_FILE_PATH>` marks the start of user code. The comment should be suffixed by `:<ORIGINAL_FILE_PATH>` to indicate the original file.", "")
 
   // Experimental language features
   @deprecated(message = "This flag has no effect and will be removed in a future version.", since = "3.7.0")

compiler/src/dotty/tools/dotc/util/SourceFile.scala

@@ -7,6 +7,7 @@ import scala.language.unsafeNulls
 import dotty.tools.io.*
 import Spans.*
 import core.Contexts.*
+import core.Decorators.*
 
 import scala.io.Codec
 import Chars.*
@@ -62,19 +63,34 @@ object ScriptSourceFile {
 }
 
 object WrappedSourceFile:
-  private val cache: mutable.HashMap[SourceFile, Int] = mutable.HashMap.empty
-  def locateMagicHeader(sourceFile: SourceFile)(using Context): Option[Int] =
-    def findOffset: Int =
-      val magicHeader = ctx.settings.XmagicOffsetHeader.value
-      if magicHeader.isEmpty then
-        -1
+  enum MagicHeaderInfo:
+    case HasHeader(offset: Int, originalFile: SourceFile)
+    case NoHeader
+  import MagicHeaderInfo.*
+
+  private val cache: mutable.HashMap[SourceFile, MagicHeaderInfo] = mutable.HashMap.empty
+
+  def locateMagicHeader(sourceFile: SourceFile)(using Context): MagicHeaderInfo =
+    def findOffset: MagicHeaderInfo =
+      val magicHeader = ctx.settings.YmagicOffsetHeader.value
+      if magicHeader.isEmpty then NoHeader
       else
-        val s = new String(sourceFile.content)
-        val regex = ("(?m)^" + java.util.regex.Pattern.quote(magicHeader) + "$").r
-        val pos = regex.findFirstMatchIn(s).map(_.start).map(sourceFile.offsetToLine(_))
-        pos.getOrElse(-1)
+        val text = new String(sourceFile.content)
+        val headerQuoted = java.util.regex.Pattern.quote("///" + magicHeader)
+        val regex = s"(?m)^$headerQuoted:(.+)$$".r
+        regex.findFirstMatchIn(text) match
+          case Some(m) =>
+            val markerOffset = m.start
+            val sourceStartOffset = sourceFile.nextLine(markerOffset)
+            val file = ctx.getFile(m.group(1))
+            if file.exists then
+              HasHeader(sourceStartOffset, ctx.getSource(file))
+            else
+              report.warning(em"original source file not found: ${file.path}")
+              NoHeader
+          case None => NoHeader
     val result = cache.getOrElseUpdate(sourceFile, findOffset)
-    if result >= 0 then Some(result + 1) else None
+    result
 
 class SourceFile(val file: AbstractFile, computeContent: => Array[Char]) extends interfaces.SourceFile {
   import SourceFile.*