Handle captures in by-name parameters

 1. Infrastructure to deal with capturesets in byname parameters

 2. Handle retainsByName annotations in ElimByName

    Convert them to regular annotations on the generated function types.
    This enables capture checking on by-name parameters.

 3. Add a style warning for misleading by-name parameter type formatting.

    By-name types should be formatted `{...}-> T`. `{...} -> T` looks too much
    like a function type.

Author: odersky

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

@@ -468,7 +468,7 @@ object desugar {
 
     if mods.is(Trait) then
       for vparams <- originalVparamss; vparam <- vparams do
-        if vparam.tpt.isInstanceOf[ByNameTypeTree] then
+        if isByNameType(vparam.tpt) then
           report.error(em"implementation restriction: traits cannot have by name parameters", vparam.srcPos)
 
     // Annotations on class _type_ parameters are set on the derived parameters
@@ -576,9 +576,8 @@ object desugar {
       appliedTypeTree(tycon, targs)
     }
 
-    def isRepeated(tree: Tree): Boolean = tree match {
+    def isRepeated(tree: Tree): Boolean = stripByNameType(tree) match {
       case PostfixOp(_, Ident(tpnme.raw.STAR)) => true
-      case ByNameTypeTree(tree1) => isRepeated(tree1)
       case _ => false
     }
 
@@ -1811,8 +1810,13 @@ object desugar {
       case ext: ExtMethods =>
         Block(List(ext), Literal(Constant(())).withSpan(ext.span))
       case CapturingTypeTree(refs, parent) =>
-        val annot = New(scalaDot(tpnme.retains), List(refs))
-        Annotated(parent, annot)
+        def annotate(annotName: TypeName, tp: Tree) =
+          Annotated(tp, New(scalaDot(annotName), List(refs)))
+        parent match
+          case ByNameTypeTree(restpt) =>
+            cpy.ByNameTypeTree(parent)(annotate(tpnme.retainsByName, restpt))
+          case _ =>
+            annotate(tpnme.retains, parent)
     }
     desugared.withSpan(tree.span)
   }

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

@@ -178,8 +178,7 @@ trait TreeInfo[T >: Untyped <: Type] { self: Trees.Instance[T] =>
   }
 
   /** Is tpt a vararg type of the form T* or => T*? */
-  def isRepeatedParamType(tpt: Tree)(using Context): Boolean = tpt match {
-    case ByNameTypeTree(tpt1) => isRepeatedParamType(tpt1)
+  def isRepeatedParamType(tpt: Tree)(using Context): Boolean = stripByNameType(tpt) match {
     case tpt: TypeTree => tpt.typeOpt.isRepeatedParam
     case AppliedTypeTree(Select(_, tpnme.REPEATED_PARAM_CLASS), _) => true
     case _ => false
@@ -196,6 +195,16 @@ trait TreeInfo[T >: Untyped <: Type] { self: Trees.Instance[T] =>
     case arg => arg.typeOpt.widen.isRepeatedParam
   }
 
+  def isByNameType(tree: Tree)(using Context): Boolean =
+    stripByNameType(tree) ne tree
+
+  def stripByNameType(tree: Tree)(using Context): Tree = unsplice(tree) match
+    case ByNameTypeTree(t1) => t1
+    case untpd.CapturingTypeTree(_, parent) =>
+      val parent1 = stripByNameType(parent)
+      if parent1 eq parent then tree else parent1
+    case _ => tree
+
   /** All type and value parameter symbols of this DefDef */
   def allParamSyms(ddef: DefDef)(using Context): List[Symbol] =
     ddef.paramss.flatten.map(_.symbol)
@@ -388,6 +397,16 @@ trait UntypedTreeInfo extends TreeInfo[Untyped] { self: Trees.Instance[Untyped]
       case _ => None
     }
   }
+
+  object ImpureByNameTypeTree:
+    def apply(tp: ByNameTypeTree)(using Context): untpd.CapturingTypeTree =
+      untpd.CapturingTypeTree(
+        Ident(nme.CAPTURE_ROOT).withSpan(tp.span.startPos) :: Nil, tp)
+    def unapply(tp: Tree)(using Context): Option[ByNameTypeTree] = tp match
+      case untpd.CapturingTypeTree(id @ Ident(nme.CAPTURE_ROOT) :: Nil, bntp: ByNameTypeTree)
+      if id.span == bntp.span.startPos => Some(bntp)
+      case _ => None
+  end ImpureByNameTypeTree
 }
 
 trait TypedTreeInfo extends TreeInfo[Type] { self: Trees.Instance[Type] =>

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

@@ -12,7 +12,7 @@ import printing.Printer
 import printing.Texts.Text
 
 
-case class CaptureAnnotation(refs: CaptureSet, boxed: Boolean) extends Annotation:
+case class CaptureAnnotation(refs: CaptureSet, kind: CapturingKind) extends Annotation:
   import CaptureAnnotation.*
   import tpd.*
 
@@ -25,25 +25,26 @@ case class CaptureAnnotation(refs: CaptureSet, boxed: Boolean) extends Annotatio
     val arg = repeated(elems, TypeTree(defn.AnyType))
     New(symbol.typeRef, arg :: Nil)
 
-  override def symbol(using Context) = defn.RetainsAnnot
+  override def symbol(using Context) =
+    if kind == CapturingKind.ByName then defn.RetainsByNameAnnot else defn.RetainsAnnot
 
   override def derivedAnnotation(tree: Tree)(using Context): Annotation =
     unsupported("derivedAnnotation(Tree)")
 
-  def derivedAnnotation(refs: CaptureSet, boxed: Boolean)(using Context): Annotation =
-    if (this.refs eq refs) && (this.boxed == boxed) then this
-    else CaptureAnnotation(refs, boxed)
+  def derivedAnnotation(refs: CaptureSet, kind: CapturingKind)(using Context): Annotation =
+    if (this.refs eq refs) && (this.kind == kind) then this
+    else CaptureAnnotation(refs, kind)
 
   override def sameAnnotation(that: Annotation)(using Context): Boolean = that match
-    case CaptureAnnotation(refs2, boxed2) => refs == refs2 && boxed == boxed2
+    case CaptureAnnotation(refs2, kind2) => refs == refs2 && kind == kind2
     case _ => false
 
   override def mapWith(tp: TypeMap)(using Context) =
     val elems = refs.elems.toList
     val elems1 = elems.mapConserve(tp)
     if elems1 eq elems then this
     else if elems1.forall(_.isInstanceOf[CaptureRef])
-    then derivedAnnotation(CaptureSet(elems1.asInstanceOf[List[CaptureRef]]*), boxed)
+    then derivedAnnotation(CaptureSet(elems1.asInstanceOf[List[CaptureRef]]*), kind)
     else EmptyAnnotation
 
   override def refersToParamOf(tl: TermLambda)(using Context): Boolean =
@@ -54,10 +55,11 @@ case class CaptureAnnotation(refs: CaptureSet, boxed: Boolean) extends Annotatio
 
   override def toText(printer: Printer): Text = refs.toText(printer)
 
-  override def hash: Int = (refs.hashCode << 1) | (if boxed then 1 else 0)
+  override def hash: Int =
+    (refs.hashCode << 1) | (if kind == CapturingKind.Regular then 0 else 1)
 
   override def eql(that: Annotation) = that match
-    case that: CaptureAnnotation => (this.refs eq that.refs) && (this.boxed == boxed)
+    case that: CaptureAnnotation => (this.refs eq that.refs) && (this.kind == kind)
     case _ => false
 
 end CaptureAnnotation

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

@@ -43,9 +43,9 @@ extension (tree: Tree)
 extension (tp: Type)
 
   def derivedCapturingType(parent: Type, refs: CaptureSet)(using Context): Type = tp match
-    case CapturingType(p, r, b) =>
+    case CapturingType(p, r, k) =>
       if (parent eq p) && (refs eq r) then tp
-      else CapturingType(parent, refs, b)
+      else CapturingType(parent, refs, k)
 
   /** If this is  type variable instantiated or upper bounded with a capturing type,
    *  the capture set associated with that type. Extended to and-or types and
@@ -54,7 +54,8 @@ extension (tp: Type)
    */
   def boxedCaptured(using Context): CaptureSet =
     def getBoxed(tp: Type): CaptureSet = tp match
-      case CapturingType(_, refs, boxed) => if boxed then refs else CaptureSet.empty
+      case CapturingType(_, refs, CapturingKind.Boxed) => refs
+      case CapturingType(_, _, _) => CaptureSet.empty
       case tp: TypeProxy => getBoxed(tp.superType)
       case tp: AndType => getBoxed(tp.tp1) ++ getBoxed(tp.tp2)
       case tp: OrType => getBoxed(tp.tp1) ** getBoxed(tp.tp2)

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

@@ -209,8 +209,9 @@ sealed abstract class CaptureSet extends Showable:
     ((NoType: Type) /: elems) ((tp, ref) =>
       if tp.exists then OrType(tp, ref, soft = false) else ref)
 
-  def toRegularAnnotation(using Context): Annotation  =
-    Annotation(CaptureAnnotation(this, boxed = false).tree)
+  def toRegularAnnotation(byName: Boolean)(using Context): Annotation =
+    val kind = if byName then CapturingKind.ByName else CapturingKind.Regular
+    Annotation(CaptureAnnotation(this, kind).tree)
 
   override def toText(printer: Printer): Text =
     Str("{") ~ Text(elems.toList.map(printer.toTextCaptureRef), ", ") ~ Str("}")

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

@@ -0,0 +1,9 @@
+package dotty.tools
+package dotc
+package cc
+
+/** Possible kinds of captures */
+enum CapturingKind:
+  case Regular     // normal capture
+  case Boxed       // capture under box
+  case ByName      // capture applies to enclosing by-name type (only possible before ElimByName)

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

@@ -5,26 +5,46 @@ package cc
 import core.*
 import Types.*, Symbols.*, Contexts.*
 
+/** A capturing type. This is internally represented as an annotated type with a `retains`
+ *  annotation, but the extractor will succeed only at phase CheckCaptures.
+ *  Annotated types with `@retainsByName` annotation can also be created that way, by
+ *  giving a `CapturingKind.ByName` as `kind` argument, but they are never extracted,
+ *  since they have already been converted to regular capturing types before CheckCaptures.
+ */
 object CapturingType:
 
-  def apply(parent: Type, refs: CaptureSet, boxed: Boolean)(using Context): Type =
+  def apply(parent: Type, refs: CaptureSet, kind: CapturingKind)(using Context): Type =
     if refs.isAlwaysEmpty then parent
-    else AnnotatedType(parent, CaptureAnnotation(refs, boxed))
-
-  def unapply(tp: AnnotatedType)(using Context): Option[(Type, CaptureSet, Boolean)] =
-    if ctx.phase == Phases.checkCapturesPhase then EventuallyCapturingType.unapply(tp)
+    else AnnotatedType(parent, CaptureAnnotation(refs, kind))
+
+  def unapply(tp: AnnotatedType)(using Context): Option[(Type, CaptureSet, CapturingKind)] =
+    if ctx.phase == Phases.checkCapturesPhase then
+      val r = EventuallyCapturingType.unapply(tp)
+      r match
+        case Some((_, _, CapturingKind.ByName)) => None
+        case _ => r
     else None
 
 end CapturingType
 
+/** An extractor for types that will be capturing types at phase CheckCaptures. Also
+ *  included are types that indicate captures on enclosing call-by-name parameters
+ *  before phase ElimByName
+ */
 object EventuallyCapturingType:
 
-  def unapply(tp: AnnotatedType)(using Context): Option[(Type, CaptureSet, Boolean)] =
-    if tp.annot.symbol == defn.RetainsAnnot then
+  def unapply(tp: AnnotatedType)(using Context): Option[(Type, CaptureSet, CapturingKind)] =
+    val sym = tp.annot.symbol
+    if sym == defn.RetainsAnnot || sym == defn.RetainsByNameAnnot then
       tp.annot match
-        case ann: CaptureAnnotation => Some((tp.parent, ann.refs, ann.boxed))
+        case ann: CaptureAnnotation =>
+          Some((tp.parent, ann.refs, ann.kind))
         case ann =>
-          try Some((tp.parent, ann.tree.toCaptureSet, ann.tree.isBoxedCapturing))
+          val kind =
+            if ann.tree.isBoxedCapturing then CapturingKind.Boxed
+            else if sym == defn.RetainsByNameAnnot then CapturingKind.ByName
+            else CapturingKind.Regular
+          try Some((tp.parent, ann.tree.toCaptureSet, kind))
           catch case ex: IllegalCaptureRef => None
     else None
 

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

@@ -25,7 +25,8 @@ extends tpd.TreeTraverser:
       .toFunctionType(isJava = false, alwaysDependent = true)
 
   private def box(tp: Type)(using Context): Type = tp match
-    case CapturingType(parent, refs, false) => CapturingType(parent, refs, true)
+    case CapturingType(parent, refs, CapturingKind.Regular) =>
+      CapturingType(parent, refs, CapturingKind.Boxed)
     case _ => tp
 
   private def setBoxed(tp: Type)(using Context) = tp match
@@ -77,7 +78,7 @@ extends tpd.TreeTraverser:
             cls.paramGetters.foldLeft(tp) { (core, getter) =>
               if getter.termRef.isTracked then
                 val getterType = tp.memberInfo(getter).strippedDealias
-                RefinedType(core, getter.name, CapturingType(getterType, CaptureSet.Var(), boxed = false))
+                RefinedType(core, getter.name, CapturingType(getterType, CaptureSet.Var(), CapturingKind.Regular))
                   .showing(i"add capture refinement $tp --> $result", capt)
               else
                 core
@@ -130,7 +131,7 @@ extends tpd.TreeTraverser:
       case tp @ OrType(tp1, CapturingType(parent2, refs2, boxed2)) =>
         CapturingType(OrType(tp1, parent2, tp.isSoft), refs2, boxed2)
       case _ if canHaveInferredCapture(tp) =>
-        CapturingType(tp, CaptureSet.Var(), boxed = false)
+        CapturingType(tp, CaptureSet.Var(), CapturingKind.Regular)
       case _ =>
         tp
 

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

@@ -15,7 +15,7 @@ import Comments.CommentsContext
 import Comments.Comment
 import util.Spans.NoSpan
 import Symbols.requiredModuleRef
-import cc.{CapturingType, CaptureSet}
+import cc.{CapturingType, CaptureSet, CapturingKind, EventuallyCapturingType}
 
 import scala.annotation.tailrec
 
@@ -118,9 +118,9 @@ class Definitions {
    *
    *  ErasedFunctionN and ErasedContextFunctionN erase to Function0.
    *
-   *  EffXYZFunctionN afollow this template:
+   *  ImpureXYZFunctionN follow this template:
    *
-   *      type EffXYZFunctionN[-T0,...,-T{N-1}, +R] = {*} XYZFunctionN[T0,...,T{N-1}, R]
+   *      type ImpureXYZFunctionN[-T0,...,-T{N-1}, +R] = {*} XYZFunctionN[T0,...,T{N-1}, R]
    */
   private def newFunctionNType(name: TypeName): Symbol = {
     val impure = name.startsWith("Impure")
@@ -136,7 +136,7 @@ class Definitions {
             HKTypeLambda(argParamNames :+ "R".toTypeName, argVariances :+ Covariant)(
               tl => List.fill(arity + 1)(TypeBounds.empty),
               tl => CapturingType(underlyingClass.typeRef.appliedTo(tl.paramRefs),
-                CaptureSet.universal, boxed = false)
+                CaptureSet.universal, CapturingKind.Regular)
             ))
         else
           val cls = denot.asClass.classSymbol
@@ -1016,6 +1016,7 @@ class Definitions {
   @tu lazy val VarargsAnnot: ClassSymbol = requiredClass("scala.annotation.varargs")
   @tu lazy val SinceAnnot: ClassSymbol = requiredClass("scala.annotation.since")
   @tu lazy val RetainsAnnot: ClassSymbol = requiredClass("scala.retains")
+  @tu lazy val RetainsByNameAnnot: ClassSymbol = requiredClass("scala.retainsByName")
 
   @tu lazy val JavaRepeatableAnnot: ClassSymbol = requiredClass("java.lang.annotation.Repeatable")
 
@@ -1149,9 +1150,16 @@ class Definitions {
     }
   }
 
+  /** Extractor for function types representing by-name parameters, of the form
+   *  `() ?=> T`.
+   *  Under -Ycc, this becomes `() ?-> T` or `{r1, ..., rN} () ?-> T`.
+   */
   object ByNameFunction:
-    def apply(tp: Type)(using Context): Type =
-      defn.ContextFunction0.typeRef.appliedTo(tp :: Nil)
+    def apply(tp: Type)(using Context): Type = tp match
+      case EventuallyCapturingType(tp1, refs, CapturingKind.ByName) =>
+        CapturingType(apply(tp1), refs, CapturingKind.Regular)
+      case _ =>
+        defn.ContextFunction0.typeRef.appliedTo(tp :: Nil)
     def unapply(tp: Type)(using Context): Option[Type] = tp match
       case tp @ AppliedType(tycon, arg :: Nil) if defn.isByNameFunctionClass(tycon.typeSymbol) =>
         Some(arg)

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

@@ -566,6 +566,7 @@ object StdNames {
     val reify : N               = "reify"
     val releaseFence : N        = "releaseFence"
     val retains: N              = "retains"
+    val retainsByName: N        = "retainsByName"
     val rootMirror : N          = "rootMirror"
     val run: N                  = "run"
     val runOrElse: N            = "runOrElse"

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

@@ -23,7 +23,7 @@ import typer.ProtoTypes.constrained
 import typer.Applications.productSelectorTypes
 import reporting.trace
 import annotation.constructorOnly
-import cc.{CapturingType, derivedCapturingType, CaptureSet, stripCapturing}
+import cc.{CapturingType, derivedCapturingType, CaptureSet, CapturingKind, stripCapturing}
 
 /** Provides methods to compare types.
  */
@@ -876,7 +876,7 @@ class TypeComparer(@constructorOnly initctx: Context) extends ConstraintHandling
           tp1 match
             case tp1: CaptureRef if tp1.isTracked =>
               val stripped = tp1w.stripCapturing
-              tp1w = CapturingType(stripped, tp1.singletonCaptureSet, boxed = false)
+              tp1w = CapturingType(stripped, tp1.singletonCaptureSet, CapturingKind.Regular)
             case _ =>
           isSubType(tp1w, tp2, approx.addLow)
         }

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

@@ -36,7 +36,7 @@ import config.Printers.{core, typr, matchTypes}
 import reporting.{trace, Message}
 import java.lang.ref.WeakReference
 import compiletime.uninitialized
-import cc.{CapturingType, CaptureSet, derivedCapturingType, retainedElems, isBoxedCapturing}
+import cc.{CapturingType, CaptureSet, derivedCapturingType, retainedElems, isBoxedCapturing, CapturingKind}
 import CaptureSet.CompareResult
 
 import scala.annotation.internal.sharable
@@ -1880,13 +1880,15 @@ object Types {
 
     def capturing(ref: CaptureRef)(using Context): Type =
       if captureSet.accountsFor(ref) then this
-      else CapturingType(this, ref.singletonCaptureSet, this.isBoxedCapturing)
+      else CapturingType(this, ref.singletonCaptureSet,
+        if this.isBoxedCapturing then CapturingKind.Boxed else CapturingKind.Regular)
 
     def capturing(cs: CaptureSet)(using Context): Type =
       if cs.isConst && cs.subCaptures(captureSet, frozen = true).isOK then this
       else this match
         case CapturingType(parent, cs1, boxed) => parent.capturing(cs1 ++ cs)
-        case _ => CapturingType(this, cs, this.isBoxedCapturing)
+        case _ => CapturingType(this, cs,
+          if this.isBoxedCapturing then CapturingKind.Boxed else CapturingKind.Regular)
 
     /** The set of distinct symbols referred to by this type, after all aliases are expanded */
     def coveringSet(using Context): Set[Symbol] =
@@ -3840,10 +3842,11 @@ object Types {
                   CapturingType(parent1, CaptureSet.universal, boxed))
             case AnnotatedType(parent, ann) if ann.refersToParamOf(thisLambdaType) =>
               val parent1 = mapOver(parent)
-              if ann.symbol == defn.RetainsAnnot then
+              if ann.symbol == defn.RetainsAnnot || ann.symbol == defn.RetainsByNameAnnot then
+                val byName = ann.symbol == defn.RetainsByNameAnnot
                 range(
-                  AnnotatedType(parent1, CaptureSet.empty.toRegularAnnotation),
-                  AnnotatedType(parent1, CaptureSet.universal.toRegularAnnotation))
+                  AnnotatedType(parent1, CaptureSet.empty.toRegularAnnotation(byName)),
+                  AnnotatedType(parent1, CaptureSet.universal.toRegularAnnotation(byName)))
               else
                 parent1
             case _ => mapOver(tp)