Adapt call-by-name parameters under pureFunctions

If the source of pickled information is not compiled with -Ypure-functions,
but the current compilation is, map call-by-name parameters in that source
to impure by-name parameters.

Author: odersky

compiler/src/dotty/tools/dotc/CompilationUnit.scala

@@ -9,11 +9,11 @@ import util.{FreshNameCreator, SourceFile, NoSource}
 import util.Spans.Span
 import ast.{tpd, untpd}
 import tpd.{Tree, TreeTraverser}
-import ast.Trees.Import
+import ast.Trees.{Import, Ident}
 import typer.Nullables
 import transform.SymUtils._
 import core.Decorators._
-import config.SourceVersion
+import config.{SourceVersion, Feature}
 import StdNames.nme
 import scala.annotation.internal.sharable
 
@@ -56,6 +56,9 @@ class CompilationUnit protected (val source: SourceFile) {
   /** Will be set to true if the unit contains a captureChecking language import */
   var needsCaptureChecking: Boolean = false
 
+  /** Will be set to true if the unit contains a pureFunctions language import */
+  var knowsPureFuns: Boolean = false
+
   var suspended: Boolean = false
   var suspendedAtInliningPhase: Boolean = false
 
@@ -151,9 +154,12 @@ object CompilationUnit {
       if tree.symbol.is(Flags.Inline) then
         containsInline = true
       tree match
-        case Import(qual, selectors)
-        if tpd.languageImport(qual).isDefined && selectors.contains(nme.captureChecking) =>
-          containsCaptureChecking = true
+        case Import(qual, selectors) =>
+          tpd.languageImport(qual) match
+            case Some(prefix) =>
+              for case untpd.ImportSelector(untpd.Ident(imported), untpd.EmptyTree, _) <- selectors do
+                Feature.handleGlobalLanguageImport(prefix, imported)
+            case _ =>
         case _ =>
       traverseChildren(tree)
     }

compiler/src/dotty/tools/dotc/Run.scala

@@ -163,6 +163,11 @@ class Run(comp: Compiler, ictx: Context) extends ImplicitRunInfo with Constraint
   /** Actions that need to be performed at the end of the current compilation run */
   private var finalizeActions = mutable.ListBuffer[() => Unit]()
 
+  /** Will be set to true if any of the compiled compilation units contains
+   *  a pureFunctions or captureChecking language import.
+   */
+  var pureFunsImportEncountered = false
+
   def compile(files: List[AbstractFile]): Unit =
     try
       val sources = files.map(runContext.getSource(_))

compiler/src/dotty/tools/dotc/cc/CaptureOps.scala

@@ -41,6 +41,22 @@ extension (tree: Tree)
         tree.putAttachment(Captures, refs)
         refs
 
+  /** Under pureFunctions, add a @retainsByName(*)` annotation to the argument of
+   *  a by name parameter type, turning the latter into an impure by name parameter type.
+   */
+  def adaptByNameArgUnderPureFuns(using Context): Tree =
+    if Feature.pureFunsEnabledSomewhere then
+      val rbn = defn.RetainsByNameAnnot
+      Annotated(tree,
+        New(rbn.typeRef).select(rbn.primaryConstructor).appliedTo(
+          Typed(
+            SeqLiteral(ref(defn.captureRoot) :: Nil, TypeTree(defn.AnyType)),
+            TypeTree(defn.RepeatedParamType.appliedTo(defn.AnyType))
+          )
+        )
+      )
+    else tree
+
 extension (tp: Type)
 
   /** @pre `tp` is a CapturingType */
@@ -125,7 +141,7 @@ extension (tp: Type)
    */
   def adaptFunctionTypeUnderPureFuns(using Context): Type = tp match
     case AppliedType(fn, args)
-    if Feature.pureFunsEnabled && defn.isFunctionClass(fn.typeSymbol) =>
+    if Feature.pureFunsEnabledSomewhere && defn.isFunctionClass(fn.typeSymbol) =>
       val fname = fn.typeSymbol.name
       defn.FunctionType(
         fname.functionArity,
@@ -135,6 +151,16 @@ extension (tp: Type)
     case _ =>
       tp
 
+  /** Under pureFunctions, add a @retainsByName(*)` annotation to the argument of
+   *  a by name parameter type, turning the latter into an impure by name parameter type.
+   */
+  def adaptByNameArgUnderPureFuns(using Context): Type =
+    if Feature.pureFunsEnabledSomewhere then
+      AnnotatedType(tp,
+        CaptureAnnotation(CaptureSet.universal, boxed = false)(defn.RetainsByNameAnnot))
+    else
+      tp
+
   def isCapturingType(using Context): Boolean =
     tp match
       case CapturingType(_, _) => true

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

@@ -80,10 +80,18 @@ object Feature:
   def scala2ExperimentalMacroEnabled(using Context) = enabled(scala2macros)
 
   def pureFunsEnabled(using Context) =
-    enabled(pureFunctions) || ccEnabled
+    enabledBySetting(pureFunctions)
+    || ctx.compilationUnit.knowsPureFuns
+    || ccEnabled
 
   def ccEnabled(using Context) =
-    enabledBySetting(captureChecking) || ctx.compilationUnit.needsCaptureChecking
+    enabledBySetting(captureChecking)
+    || ctx.compilationUnit.needsCaptureChecking
+
+  def pureFunsEnabledSomewhere(using Context) =
+    enabledBySetting(pureFunctions)
+    || enabledBySetting(captureChecking)
+    || ctx.run != null && ctx.run.nn.pureFunsImportEncountered
 
   def sourceVersionSetting(using Context): SourceVersion =
     SourceVersion.valueOf(ctx.settings.source.value)
@@ -131,4 +139,16 @@ object Feature:
   def isExperimentalEnabled(using Context): Boolean =
     Properties.experimental && !ctx.settings.YnoExperimental.value
 
+  def handleGlobalLanguageImport(prefix: TermName, imported: Name)(using Context): Boolean =
+    val fullFeatureName = QualifiedName(prefix, imported.asTermName)
+    if fullFeatureName == pureFunctions then
+      ctx.compilationUnit.knowsPureFuns = true
+      if ctx.run != null then ctx.run.nn.pureFunsImportEncountered = true
+      true
+    else if fullFeatureName == captureChecking then
+      ctx.compilationUnit.needsCaptureChecking = true
+      if ctx.run != null then ctx.run.nn.pureFunsImportEncountered = true
+      true
+    else
+      false
 end Feature

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

@@ -32,7 +32,7 @@ import ast.{Trees, tpd, untpd}
 import Trees._
 import Decorators._
 import transform.SymUtils._
-import cc.adaptFunctionTypeUnderPureFuns
+import cc.{adaptFunctionTypeUnderPureFuns, adaptByNameArgUnderPureFuns}
 
 import dotty.tools.tasty.{TastyBuffer, TastyReader}
 import TastyBuffer._
@@ -455,7 +455,8 @@ class TreeUnpickler(reader: TastyReader,
           val ref = readAddr()
           typeAtAddr.getOrElseUpdate(ref, forkAt(ref).readType())
         case BYNAMEtype =>
-          ExprType(readType())
+          val arg = readType()
+          ExprType(if knowsPureFuns then arg else arg.adaptByNameArgUnderPureFuns)
         case _ =>
           ConstantType(readConstant(tag))
       }
@@ -1178,7 +1179,8 @@ class TreeUnpickler(reader: TastyReader,
         case SINGLETONtpt =>
           SingletonTypeTree(readTerm())
         case BYNAMEtpt =>
-          ByNameTypeTree(readTpt())
+          val arg = readTpt()
+          ByNameTypeTree(if knowsPureFuns then arg else arg.adaptByNameArgUnderPureFuns)
         case NAMEDARG =>
           NamedArg(readName(), readTerm())
         case _ =>

compiler/src/dotty/tools/dotc/core/unpickleScala2/Scala2Unpickler.scala

@@ -33,7 +33,7 @@ import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
 import scala.annotation.switch
 import reporting._
-import cc.adaptFunctionTypeUnderPureFuns
+import cc.{adaptFunctionTypeUnderPureFuns, adaptByNameArgUnderPureFuns}
 
 object Scala2Unpickler {
 
@@ -817,7 +817,7 @@ class Scala2Unpickler(bytes: Array[Byte], classRoot: ClassDenotation, moduleClas
         }
         val tycon = select(pre, sym)
         val args = until(end, () => readTypeRef())
-        if (sym == defn.ByNameParamClass2x) ExprType(args.head)
+        if (sym == defn.ByNameParamClass2x) ExprType(args.head.adaptByNameArgUnderPureFuns)
         else if (ctx.settings.scalajs.value && args.length == 2 &&
             sym.owner == JSDefinitions.jsdefn.ScalaJSJSPackageClass && sym == JSDefinitions.jsdefn.PseudoUnionClass) {
           // Treat Scala.js pseudo-unions as real unions, this requires a

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

@@ -196,7 +196,7 @@ object Parsers {
 
     def isIdent = in.isIdent
     def isIdent(name: Name) = in.isIdent(name)
-    def isPureArrow(name: Name): Boolean = in.pureFunsEnabled && isIdent(name)
+    def isPureArrow(name: Name): Boolean = isIdent(name) && Feature.pureFunsEnabled
     def isPureArrow: Boolean = isPureArrow(nme.PUREARROW) || isPureArrow(nme.PURECTXARROW)
     def isErased = isIdent(nme.erased) && in.erasedEnabled
     def isSimpleLiteral =
@@ -972,7 +972,7 @@ object Parsers {
      *  capture set `{ref1, ..., refN}` followed by a token that can start a type?
      */
     def followingIsCaptureSet(): Boolean =
-      in.featureEnabled(Feature.captureChecking) && {
+      Feature.ccEnabled && {
         val lookahead = in.LookaheadScanner()
         def followingIsTypeStart() =
           lookahead.nextToken()
@@ -1485,7 +1485,7 @@ object Parsers {
             if !imods.flags.isEmpty || params.isEmpty then
               syntaxError(em"illegal parameter list for type lambda", start)
               token = ARROW
-          else if in.pureFunsEnabled then
+          else if Feature.pureFunsEnabled then
             // `=>` means impure function under pureFunctions or captureChecking
             // language imports, whereas `->` is then a regular function.
             imods |= Impure
@@ -1891,7 +1891,7 @@ object Parsers {
       if in.token == ARROW || isPureArrow(nme.PUREARROW) then
         val isImpure = in.token == ARROW
         val tp = atSpan(in.skipToken()) { ByNameTypeTree(core()) }
-        if isImpure && in.pureFunsEnabled then ImpureByNameTypeTree(tp) else tp
+        if isImpure && Feature.pureFunsEnabled then ImpureByNameTypeTree(tp) else tp
       else if in.token == LBRACE && followingIsCaptureSet() then
         val start = in.offset
         val cs = captureSet()
@@ -3308,14 +3308,9 @@ object Parsers {
       languageImport(tree) match
         case Some(prefix) =>
           in.languageImportContext = in.languageImportContext.importContext(imp, NoSymbol)
-          for
-            case ImportSelector(id @ Ident(imported), EmptyTree, _) <- selectors
-          do
-            val fullFeatureName = QualifiedName(prefix, imported.asTermName)
-            if globalOnlyImports.contains(fullFeatureName) && !outermost then
+          for case ImportSelector(id @ Ident(imported), EmptyTree, _) <- selectors do
+            if Feature.handleGlobalLanguageImport(prefix, imported) && !outermost then
               syntaxError(i"this language import is only allowed at the toplevel", id.span)
-            if fullFeatureName == Feature.captureChecking then
-              ctx.compilationUnit.needsCaptureChecking = true
             if allSourceVersionNames.contains(imported) && prefix.isEmpty then
               if !outermost then
                 syntaxError(i"source version import is only allowed at the toplevel", id.span)

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

@@ -230,15 +230,6 @@ object Scanners {
         postfixOpsEnabledCtx = myLanguageImportContext
       postfixOpsEnabledCache
 
-    private var pureFunsEnabledCache = false
-    private var pureFunsEnabledCtx: Context = NoContext
-
-    def pureFunsEnabled =
-      if pureFunsEnabledCtx ne myLanguageImportContext then
-        pureFunsEnabledCache = featureEnabled(Feature.pureFunctions) || featureEnabled(Feature.captureChecking)
-        pureFunsEnabledCtx = myLanguageImportContext
-      pureFunsEnabledCache
-
     /** All doc comments kept by their end position in a `Map`.
       *
       * Note: the map is necessary since the comments are looked up after an

compiler/test/dotc/pos-test-pickling.blacklist

@@ -81,6 +81,9 @@ i13842.scala
 # Position change under captureChecking
 boxmap-paper.scala
 
+# Function types print differnt after unpickling since test mispredicts Feature.preFundsEnabled
+caps.universal.scala
+
 # GADT cast applied to singleton type difference
 i4176-gadt.scala