modify spans of statements with end markers

Author: bishabosha

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

@@ -875,7 +875,6 @@ object desugar {
       val modul = ValDef(moduleName, clsRef, New(clsRef, Nil))
         .withMods(mods.toTermFlags & RetainedModuleValFlags | ModuleValCreationFlags)
         .withSpan(mdef.span.startPos)
-        .withEndIndex(copyFrom = mdef) // copy over the end marker position to the module val
       val ValDef(selfName, selfTpt, _) = impl.self
       val selfMods = impl.self.mods
       if (!selfTpt.isEmpty) report.error(ObjectMayNotHaveSelfType(mdef), impl.self.srcPos)
@@ -885,6 +884,7 @@ object desugar {
       val clsTmpl = cpy.Template(impl)(self = clsSelf, body = impl.body)
       val cls = TypeDef(clsName, clsTmpl)
         .withMods(mods.toTypeFlags & RetainedModuleClassFlags | ModuleClassCreationFlags)
+        .withEndIndex(copyFrom = mdef) // copy over the end marker position to the module class def
       Thicket(modul, classDef(cls).withSpan(mdef.span))
     }
   }

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

@@ -327,42 +327,45 @@ object Trees {
 
   extension (mdef: untpd.DefTree) def mods: untpd.Modifiers = mdef.rawMods
 
-  /** PackageDef | NamedDefTree */
-  sealed trait WithEndMarker:
-    self: Attachment.Container =>
+  sealed trait WithEndMarker[-T >: Untyped]:
+    self: PackageDef[T] | NamedDefTree[T] =>
 
     import WithEndMarker.*
 
     final def endSpan(using Context): Span =
-      self.getAttachment(EndIndex) match
-        case Some(end) =>
-          val realName = srcName.stripModuleClassSuffix.lastPart
-          Span(end - realName.length, end)
-        case none => NoSpan
+      if hasEnd then
+        val realName = srcName.stripModuleClassSuffix.lastPart
+        span.withStart(span.end - realName.length)
+      else
+        NoSpan
 
     protected def srcName(using Context): Name
 
-    final def withEndIndex(index: Int): self.type =
-      self.withAttachment(EndIndex, index)
+    final def withEndIndex(): self.type =
+      self.withAttachment(EndIndex, ())
 
-    final def withEndIndex(copyFrom: WithEndMarker): self.type =
-      copyFrom.endIndex.foreach(withEndIndex)
-      this
+    final def withEndIndex(copyFrom: WithEndMarker[T @uncheckedVariance]): self.type =
+      if copyFrom.hasEnd then
+        this.withEndIndex()
+      else
+        this
 
     final def dropEndIndex(): self.type =
       self.removeAttachment(EndIndex)
       this
 
-    protected def endIndex: Option[Int] = self.getAttachment(EndIndex)
+    protected def hasEnd: Boolean = self.hasAttachment(EndIndex)
 
   object WithEndMarker:
-    /** Property key for trees with an `end` marker */
-    private val EndIndex: Property.StickyKey[Int] = Property.StickyKey()
+    /** Property key that signals the tree was terminated
+     *  with an `end` marker in the source code
+     */
+    private val EndIndex: Property.StickyKey[Unit] = Property.StickyKey()
 
   end WithEndMarker
 
   abstract class NamedDefTree[-T >: Untyped](implicit @constructorOnly src: SourceFile)
-  extends NameTree[T] with DefTree[T] with WithEndMarker {
+  extends NameTree[T] with DefTree[T] with WithEndMarker[T] {
     type ThisTree[-T >: Untyped] <: NamedDefTree[T]
 
     protected def srcName(using Context): Name =
@@ -897,7 +900,7 @@ object Trees {
 
   /** package pid { stats } */
   case class PackageDef[-T >: Untyped] private[ast] (pid: RefTree[T], stats: List[Tree[T]])(implicit @constructorOnly src: SourceFile)
-    extends ProxyTree[T] with WithEndMarker {
+    extends ProxyTree[T] with WithEndMarker[T] {
     type ThisTree[-T >: Untyped] = PackageDef[T]
     def forwardTo: RefTree[T] = pid
     protected def srcName(using Context): Name = pid.name
@@ -1092,6 +1095,8 @@ object Trees {
     type Annotated = Trees.Annotated[T]
     type Thicket = Trees.Thicket[T]
 
+    type WithEndMarker = Trees.WithEndMarker[T]
+
     @sharable val EmptyTree: Thicket = genericEmptyTree
     @sharable val EmptyValDef: ValDef = genericEmptyValDef
     @sharable val ContextualEmptyTree: Thicket = new EmptyTree() // an empty tree marking a contextual closure

compiler/src/dotty/tools/dotc/parsing/Parsers.scala

@@ -1275,7 +1275,7 @@ object Parsers {
 
     def checkEndMarker[T <: Tree](stats: ListBuffer[T]): Unit =
 
-      def matches(stat: Tree): Boolean = stat match
+      def matches(stat: T): Boolean = stat match
         case stat: MemberDef if !stat.name.isEmpty =>
           if stat.name == nme.CONSTRUCTOR then in.token == THIS
           else in.isIdent && in.name == stat.name.toTermName
@@ -1293,14 +1293,20 @@ object Parsers {
         case _: (ForYield | ForDo) => in.token == FOR
         case _ => false
 
-      def matchesAndSetEnd(stat: Tree): Boolean = {
+      def matchesAndSetEnd(stat: T): Boolean = {
         val didMatch = matches(stat)
         if didMatch then
           stat match
             case stat: WithEndMarker =>
-              stat.withEndIndex(index = in.lastCharOffset)
+              stat.withEndIndex()
             case _ =>
               ()
+
+          // VERY SNEAKY MUTATION HERE:
+          // `withSpan` can clone the tree, so we must replace the last statment
+          stats.dropRightInPlace(1)
+          stats.addOne(stat.withSpan(stat.span.withEnd(in.lastCharOffset)))
+        end if
         didMatch
       }
 

compiler/src/dotty/tools/dotc/semanticdb/ExtractSemanticDB.scala

@@ -240,7 +240,7 @@ class ExtractSemanticDB extends Phase:
           traverseChildren(tree)
 
       tree match
-        case tree: WithEndMarker =>
+        case tree: WithEndMarker[t] =>
           val endSpan = tree.endSpan
           if endSpan.exists then
             registerUseGuarded(None, tree.symbol, endSpan, tree.source)