Add capture checking attributes to TASTy

Author: nicolasstucki

compiler/src/dotty/tools/dotc/core/CompilationUnitInfo.scala

@@ -14,18 +14,22 @@ import dotty.tools.tasty.TastyVersion
 class CompilationUnitInfo(
   val associatedFile: AbstractFile,
   val tastyVersion: Option[TastyVersion],
-  val explicitNulls: Boolean
+  val explicitNulls: Boolean,
+  val captureChecked: Boolean,
+  val withPureFuns: Boolean,
 ) {
 
   override def toString(): String =
-    s"CompilationUnitInfo($associatedFile, $tastyVersion)"
+    s"CompilationUnitInfo($associatedFile, $tastyVersion, explicitNulls = $explicitNulls, captureChecked = $captureChecked, withPureFuns = $withPureFuns)"
 }
 
 object CompilationUnitInfo:
-  def apply(assocFile: AbstractFile | Null, explicitNulls: Boolean = false): CompilationUnitInfo | Null =
+  def apply(assocFile: AbstractFile | Null, explicitNulls: Boolean = false, captureChecked: Boolean = false, withPureFuns: Boolean = false): CompilationUnitInfo | Null =
     if assocFile == null then null
     else new CompilationUnitInfo(
       assocFile,
       tastyVersion = None,
       explicitNulls = explicitNulls,
+      captureChecked = captureChecked,
+      withPureFuns = withPureFuns,
     )

compiler/src/dotty/tools/dotc/core/Definitions.scala

@@ -1039,7 +1039,6 @@ class Definitions {
   @tu lazy val UncheckedCapturesAnnot: ClassSymbol = requiredClass("scala.annotation.unchecked.uncheckedCaptures")
   @tu lazy val VolatileAnnot: ClassSymbol = requiredClass("scala.volatile")
   @tu lazy val WithPureFunsAnnot: ClassSymbol = requiredClass("scala.annotation.internal.WithPureFuns")
-  @tu lazy val CaptureCheckedAnnot: ClassSymbol = requiredClass("scala.annotation.internal.CaptureChecked")
   @tu lazy val BeanGetterMetaAnnot: ClassSymbol = requiredClass("scala.annotation.meta.beanGetter")
   @tu lazy val BeanSetterMetaAnnot: ClassSymbol = requiredClass("scala.annotation.meta.beanSetter")
   @tu lazy val FieldMetaAnnot: ClassSymbol = requiredClass("scala.annotation.meta.field")

compiler/src/dotty/tools/dotc/core/SymbolLoaders.scala

@@ -434,6 +434,8 @@ class TastyLoader(val tastyFile: AbstractFile) extends SymbolLoader {
       tastyFile,
       tastyVersion = Some(tastyVersion),
       explicitNulls = attributes.explicitNulls,
+      captureChecked = attributes.captureChecked,
+      withPureFuns = attributes.withPureFuns,
     )
 
   def description(using Context): String = "TASTy file " + tastyFile.toString

compiler/src/dotty/tools/dotc/core/Symbols.scala

@@ -293,6 +293,16 @@ object Symbols {
       val compUnitInfo = compilationUnitInfo
       compUnitInfo != null && compUnitInfo.explicitNulls
 
+    /** If this class is capture checked */
+    def captureChecked(using Context): Boolean =
+      val compUnitInfo = compilationUnitInfo
+      compUnitInfo != null && compUnitInfo.captureChecked
+
+    /** If this class is uses pure functions */
+    def withPureFuns(using Context): Boolean =
+      val compUnitInfo = compilationUnitInfo
+      compUnitInfo != null && compUnitInfo.withPureFuns
+
     /** The class file from which this class was generated, null if not applicable. */
     final def binaryFile(using Context): AbstractFile | Null = {
       val file = associatedFile

compiler/src/dotty/tools/dotc/core/tasty/AttributePickler.scala

@@ -12,7 +12,7 @@ object AttributePickler:
     pickler: TastyPickler,
     buf: TastyBuffer
   ): Unit =
-    if attributes.scala2StandardLibrary || attributes.explicitNulls then // or any other attribute is set
+    if attributes.scala2StandardLibrary || attributes.explicitNulls || attributes.captureChecked || attributes.withPureFuns then // or any other attribute is set
       pickler.newSection(AttributesSection, buf)
 
       for tag <- attributes.booleanTags do

compiler/src/dotty/tools/dotc/core/tasty/Attributes.scala

@@ -9,15 +9,23 @@ class Attributes(
     booleanTags.contains(TastyFormat.SCALA2STANDARDLIBRARYattr)
   def explicitNulls: Boolean =
     booleanTags.contains(TastyFormat.EXPLICITNULLSattr)
+  def captureChecked: Boolean =
+    booleanTags.contains(TastyFormat.CAPTURECHECKEDattr)
+  def withPureFuns: Boolean =
+    booleanTags.contains(TastyFormat.WITHPUREFUNSattr)
 }
 
 object Attributes:
   def apply(
     scala2StandardLibrary: Boolean,
     explicitNulls: Boolean,
+    captureChecked: Boolean,
+    withPureFuns: Boolean,
   ): Attributes =
     val booleanTags = List.newBuilder[Int]
     if scala2StandardLibrary then booleanTags += TastyFormat.SCALA2STANDARDLIBRARYattr
     if explicitNulls then booleanTags += TastyFormat.EXPLICITNULLSattr
+    if captureChecked then booleanTags += TastyFormat.CAPTURECHECKEDattr
+    if withPureFuns then booleanTags += TastyFormat.WITHPUREFUNSattr
     new Attributes(booleanTags.result())
   end apply

compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala

@@ -93,11 +93,9 @@ class TreeUnpickler(reader: TastyReader,
   /** The root owner tree. See `OwnerTree` class definition. Set by `enterTopLevel`. */
   private var ownerTree: OwnerTree = uninitialized
 
-  /** Was unpickled class compiled with pureFunctions? */
-  private var withPureFuns: Boolean = false
-
   /** Was unpickled class compiled with capture checks? */
-  private var withCaptureChecks: Boolean = false
+  private val withCaptureChecks: Boolean =
+    attributeUnpicklerOpt.exists(_.attributes.captureChecked)
 
   private val unpicklingScala2Library =
     attributeUnpicklerOpt.exists(_.attributes.scala2StandardLibrary)
@@ -656,13 +654,8 @@ class TreeUnpickler(reader: TastyReader,
       }
       registerSym(start, sym)
       if (isClass) {
-        if sym.owner.is(Package) then
-          if annots.exists(_.hasSymbol(defn.CaptureCheckedAnnot)) then
-            sym.setFlag(CaptureChecked)
-            withCaptureChecks = true
-            withPureFuns = true
-          else if annots.exists(_.hasSymbol(defn.WithPureFunsAnnot)) then
-            withPureFuns = true
+        if sym.owner.is(Package) && withCaptureChecks then
+          sym.setFlag(CaptureChecked)
         sym.completer.withDecls(newScope)
         forkAt(templateStart).indexTemplateParams()(using localContext(sym))
       }

compiler/src/dotty/tools/dotc/transform/Pickler.scala

@@ -7,6 +7,7 @@ import Contexts.*
 import Decorators.*
 import tasty.*
 import config.Printers.{noPrinter, pickling}
+import config.Feature
 import java.io.PrintStream
 import Periods.*
 import Phases.*
@@ -111,6 +112,8 @@ class Pickler extends Phase {
           val attributes = Attributes(
             scala2StandardLibrary = ctx.settings.YcompileScala2Library.value,
             explicitNulls = ctx.settings.YexplicitNulls.value,
+            captureChecked = Feature.ccEnabled,
+            withPureFuns = Feature.pureFunsEnabled,
           )
           AttributePickler.pickleAttributes(attributes, pickler, scratch.attributeBuffer)
 

compiler/src/dotty/tools/dotc/transform/PostTyper.scala

@@ -423,11 +423,6 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
                 val reference = ctx.settings.sourceroot.value
                 val relativePath = util.SourceFile.relativePath(ctx.compilationUnit.source, reference)
                 sym.addAnnotation(Annotation.makeSourceFile(relativePath, tree.span))
-              if sym != defn.WithPureFunsAnnot && sym != defn.CaptureCheckedAnnot then
-                if Feature.ccEnabled then
-                  sym.addAnnotation(Annotation(defn.CaptureCheckedAnnot, tree.span))
-                else if Feature.pureFunsEnabled then
-                  sym.addAnnotation(Annotation(defn.WithPureFunsAnnot, tree.span))
           else
             if !sym.is(Param) && !sym.owner.isOneOf(AbstractOrTrait) then
               Checking.checkGoodBounds(tree.symbol)

compiler/src/dotty/tools/dotc/typer/Namer.scala

@@ -247,10 +247,16 @@ class Namer { typer: Typer =>
       case tree: TypeDef if tree.isClassDef =>
         val flags = checkFlags(tree.mods.flags)
         val name = checkNoConflict(tree.name, flags.is(Private), tree.span).asTypeName
+        val compilationUnitInfo = CompilationUnitInfo(
+          ctx.source.file,
+          explicitNulls = ctx.explicitNulls,
+          captureChecked = Feature.ccEnabled,
+          withPureFuns = Feature.pureFunsEnabled,
+        )
         val cls =
           createOrRefine[ClassSymbol](tree, name, flags, ctx.owner,
             cls => adjustIfModule(new ClassCompleter(cls, tree)(ctx), tree),
-            newClassSymbol(ctx.owner, name, _, _, _, tree.nameSpan, CompilationUnitInfo(ctx.source.file, explicitNulls = ctx.explicitNulls)))
+            newClassSymbol(ctx.owner, name, _, _, _, tree.nameSpan, compilationUnitInfo))
         cls.completer.asInstanceOf[ClassCompleter].init()
         cls
       case tree: MemberDef =>