modify spans of statements with end markers

Author: bishabosha

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

@@ -182,7 +182,7 @@ object desugar {
         tpt     = TypeTree(defn.UnitType),
         rhs     = setterRhs
       ).withMods((mods | Accessor) &~ (CaseAccessor | GivenOrImplicit | Lazy))
-       .dropEndIndex() // the end marker should only appear on the getter definition
+       .dropEndMarker() // the end marker should only appear on the getter definition
       Thicket(vdef1, setter)
     }
     else vdef1
@@ -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)
+        .withEndMarker(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,51 @@ 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 hasEndMarker then
+        val realName = srcName.stripModuleClassSuffix.lastPart
+        span.withStart(span.end - realName.length)
+      else
+        NoSpan
 
+    /** The name in source code that represents this construct,
+     *  and is the name that the user must write to create a valid
+     *  end marker.
+     *  e.g. a constructor definition is terminated in the source
+     *  code by `end this`, so it's `srcName` should return `this`.
+     */
     protected def srcName(using Context): Name
 
-    final def withEndIndex(index: Int): self.type =
-      self.withAttachment(EndIndex, index)
+    final def withEndMarker(): self.type =
+      self.withAttachment(HasEndMarker, ())
 
-    final def withEndIndex(copyFrom: WithEndMarker): self.type =
-      copyFrom.endIndex.foreach(withEndIndex)
-      this
+    final def withEndMarker(copyFrom: WithEndMarker[?]): self.type =
+      if copyFrom.hasEndMarker then
+        this.withEndMarker()
+      else
+        this
 
-    final def dropEndIndex(): self.type =
-      self.removeAttachment(EndIndex)
+    final def dropEndMarker(): self.type =
+      self.removeAttachment(HasEndMarker)
       this
 
-    protected def endIndex: Option[Int] = self.getAttachment(EndIndex)
+    protected def hasEndMarker: Boolean = self.hasAttachment(HasEndMarker)
 
   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 HasEndMarker: 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 +906,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

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

@@ -1275,7 +1275,13 @@ object Parsers {
 
     def checkEndMarker[T <: Tree](stats: ListBuffer[T]): Unit =
 
-      def matches(stat: Tree): Boolean = stat match
+      def updateSpanOfLast(last: T): Unit =
+        last match
+          case last: WithEndMarker[t] => last.withEndMarker()
+          case _ =>
+        last.span = last.span.withEnd(in.lastCharOffset)
+
+      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 +1299,10 @@ object Parsers {
         case _: (ForYield | ForDo) => in.token == FOR
         case _ => false
 
-      def matchesAndSetEnd(stat: Tree): Boolean = {
-        val didMatch = matches(stat)
+      def matchesAndSetEnd(last: T): Boolean = {
+        val didMatch = matches(last)
         if didMatch then
-          stat match
-            case stat: WithEndMarker =>
-              stat.withEndIndex(index = in.lastCharOffset)
-            case _ =>
-              ()
+          updateSpanOfLast(last)
         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)