Check capture sets for well-formedness

Author: odersky

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

@@ -12,6 +12,10 @@ import config.Printers.capt
 object CaptureOps:
   import tpd.*
 
+  def retainedElems(tree: Tree)(using Context): List[Tree] = tree match
+    case Apply(_, Typed(SeqLiteral(elems, _), _) :: Nil) => elems
+    case _ => Nil
+
   extension (cs: CaptureSet)
     def toAnnotation(using Context): Annotation =
       val refs = cs.elems.toList.map {
@@ -29,10 +33,8 @@ object CaptureOps:
   extension (annot: Annotation)
     def toCaptureSet(using Context): CaptureSet =
       assert(annot.symbol == defn.RetainsAnnot)
-      annot.tree match
-        case Apply(_, Typed(SeqLiteral(elems, _), _) :: Nil) =>
-          CaptureSet(elems.map(_.tpe.asInstanceOf[CaptureRef])*)
-            .showing(i"toCaptureSet $annot --> $result", capt)
+      CaptureSet(retainedElems(annot.tree).map(_.tpe.asInstanceOf[CaptureRef])*)
+        .showing(i"toCaptureSet $annot --> $result", capt)
 
   extension (tp: AnnotatedType)
     def toCapturingType(using Context): Type =

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

@@ -37,11 +37,13 @@ sealed abstract class CaptureSet extends Showable:
    */
   def isConst: Boolean
 
-  /** Is this capture set (always) empty? For capture veraiables, returns
+  /** Is this capture set always empty? For capture veraiables, returns
    *  always false
    */
-  def isEmpty: Boolean
-  def nonEmpty: Boolean = !isEmpty
+  def isAlwaysEmpty: Boolean
+
+  /** Is this capture set definitely non-empty? */
+  final def isNotEmpty: Boolean = !elems.isEmpty
 
   /** Add new elements to this capture set if allowed.
    *  @pre `newElems` is not empty and does not overlap with `this.elems`.
@@ -74,8 +76,10 @@ sealed abstract class CaptureSet extends Showable:
    *                 as frozen.
    */
   def accountsFor(x: CaptureRef)(using Context): Boolean =
-    elems.contains(x)
-    || !x.isRootCapability && (x.captureSetOfInfo frozen_<:< this) == CompareResult.OK
+    reporting.trace(i"$this accountsFor $x, ${x.captureSetOfInfo}?", show = true) {
+      elems.contains(x)
+      || !x.isRootCapability && (x.captureSetOfInfo frozen_<:< this) == CompareResult.OK
+   }
 
   /** The subcapturing test */
   def <:< (that: CaptureSet)(using Context): CompareResult =
@@ -172,7 +176,7 @@ object CaptureSet:
   class Const private[CaptureSet] (val elems: Refs) extends CaptureSet:
     assert(elems != null)
     def isConst = true
-    def isEmpty: Boolean = elems.isEmpty
+    def isAlwaysEmpty = elems.isEmpty
 
     def addNewElems(elems: Refs)(using Context, VarState): CompareResult =
       CompareResult.fail(this)
@@ -190,7 +194,7 @@ object CaptureSet:
     var elems: Refs = initialElems
     var deps: Deps = emptySet
     def isConst = false
-    def isEmpty = false
+    def isAlwaysEmpty = false
 
     private def recordElemsState()(using VarState): Boolean =
       varState.getElems(this) match

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

@@ -2041,7 +2041,7 @@ object Types {
     private var mySingletonCaptureSet: CaptureSet = null
 
     def canBeTracked(using Context): Boolean
-    final def isTracked(using Context): Boolean = canBeTracked && captureSetOfInfo.nonEmpty
+    final def isTracked(using Context): Boolean = canBeTracked && !captureSetOfInfo.isAlwaysEmpty
     def isRootCapability(using Context): Boolean = false
     def normalizedRef(using Context): CaptureRef = this
 
@@ -2067,7 +2067,7 @@ object Types {
 
     override def captureSet(using Context): CaptureSet =
       val cs = captureSetOfInfo
-      if canBeTracked && cs.nonEmpty then singletonCaptureSet else cs
+      if canBeTracked && !cs.isAlwaysEmpty then singletonCaptureSet else cs
   end CaptureRef
 
   /** A trait for types that bind other types that refer to them.
@@ -5224,7 +5224,7 @@ object Types {
 
   object CapturingType:
     def apply(parent: Type, refs: CaptureSet)(using Context): Type =
-      if refs.isEmpty then parent
+      if refs.isAlwaysEmpty then parent
       else unique(CachedCapturingType(parent, refs))
   end CapturingType
 

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

@@ -31,12 +31,39 @@ import CaptureSet.CompareResult
 import cc.CaptureOps.*
 
 object CheckCaptures:
+  import ast.tpd.*
+
   case class Env(owner: Symbol, captured: CaptureSet, isBoxed: Boolean, outer: Env):
-    def isOpen = !captured.isEmpty && !isBoxed
+    def isOpen = !captured.isAlwaysEmpty && !isBoxed
 
   /** Attachment key for printing trees with rechecked types */
   val RecheckedType = Property.Key[Type]
 
+  /** Check that a @retains annotation only mentions references that can be tracked
+   *  This check is performed at Typer.
+   */
+  def checkWellformed(ann: Tree)(using Context): Unit =
+    for elem <- retainedElems(ann) do
+      elem.tpe match
+        case ref: CaptureRef =>
+          if !ref.canBeTracked then
+            report.error(em"$elem cannot be tracked since it is not a parameter or a local variable", elem.srcPos)
+        case tpe =>
+          report.error(em"$tpe is not a legal type for a capture set", elem.srcPos)
+
+  /** If `tp` is a capturing type, check that all references it mentions have non-empty
+   *  capture sets.
+   *  This check is performed after capture sets are computed in phase cc.
+   */
+  def checkWellformedPost(tp: Type, pos: SrcPos)(using Context): Unit = tp match
+    case tp: CapturingType =>
+      for ref <- tp.refs.elems do
+        if (ref.captureSet frozen_<:< CaptureSet.empty) == CompareResult.OK then
+          report.error(em"$ref cannot be tracked since its capture set is empty", pos)
+    case _ =>
+
+  private inline val disallowGlobal = true
+
 class CheckCaptures extends Recheck:
   thisPhase =>
 
@@ -71,6 +98,7 @@ class CheckCaptures extends Recheck:
               parent
             case _ =>
               mapOver(t)
+
         def addVars(tp: Type): Type =
           if tp.canHaveInferredCapture then
             val tp1 = tp match
@@ -80,6 +108,7 @@ class CheckCaptures extends Recheck:
                 tp
             CapturingType(tp1, CaptureSet.Var())
           else tp
+
         tp match
           case tp: MethodOrPoly =>
             tp.derivedLambdaType(resType = reinfer(tp.resType))
@@ -88,6 +117,7 @@ class CheckCaptures extends Recheck:
           case _ =>
             val tp1 = cleanType(tp)
             addVars(tp1)
+      end reinfer
 
       val tp1 = if inferred then reinfer(tp) else tp
       mapType(tp1)
@@ -115,7 +145,7 @@ class CheckCaptures extends Recheck:
       if curEnv.isOpen then
         val ownEnclosure = ctx.owner.enclosingMethodOrClass
         var targetSet = capturedVars(sym)
-        if !targetSet.isEmpty && sym.enclosure == ownEnclosure then
+        if !targetSet.isAlwaysEmpty && sym.enclosure == ownEnclosure then
           targetSet = targetSet.filter {
             case ref: TermRef => ref.symbol.enclosure != ownEnclosure
             case _ => true
@@ -166,14 +196,14 @@ class CheckCaptures extends Recheck:
       remember(tree, sym.localReturnType)
       val saved = curEnv
       val localSet = capturedVars(sym)
-      if !localSet.isEmpty then curEnv = Env(sym, localSet, false, curEnv)
+      if !localSet.isAlwaysEmpty then curEnv = Env(sym, localSet, false, curEnv)
       try super.recheckDefDef(tree, sym)
       finally curEnv = saved
 
     override def recheckClassDef(tree: TypeDef, impl: Template, sym: ClassSymbol)(using Context): Type =
       val saved = curEnv
       val localSet = capturedVars(sym)
-      if !localSet.isEmpty then curEnv = Env(sym, localSet, false, curEnv)
+      if !localSet.isAlwaysEmpty then curEnv = Env(sym, localSet, false, curEnv)
       try super.recheckClassDef(tree, impl, sym)
       finally curEnv = saved
 
@@ -198,44 +228,47 @@ class CheckCaptures extends Recheck:
       super.checkUnit(unit)
       PostRefinerCheck.traverse(unit.tpdTree)
 
-  end CaptureChecker
+    def checkNotGlobal(tree: Tree, allArgs: Tree*)(using Context): Unit =
+      if disallowGlobal then
+        tree match
+          case LambdaTypeTree(_, restpt) =>
+            checkNotGlobal(restpt, allArgs*)
+          case _ =>
+            for ref <- tree.tpe.captureSet.elems do
+              val isGlobal = ref match
+                case ref: TermRef =>
+                  ref.isRootCapability || ref.prefix != NoPrefix && ref.symbol.hasAnnotation(defn.AbilityAnnot)
+                case _ => false
+              val what = if ref.isRootCapability then "universal" else "global"
+              if isGlobal then
+                val notAllowed = i" is not allowed to capture the $what capability $ref"
+                def msg = tree match
+                  case tree: InferredTypeTree =>
+                    i"""inferred type argument ${tree.tpe}$notAllowed
+                        |
+                        |The inferred arguments are: [$allArgs%, %]"""
+                  case _ => s"type argument$notAllowed"
+                report.error(msg, tree.srcPos)
+
+    object PostRefinerCheck extends TreeTraverser:
+      def traverse(tree: Tree)(using Context) =
+        tree match
+          case _: InferredTypeTree =>
+          case tree: TypeTree =>
+            transformType(tree.tpe, inferred = false).foreachPart(
+              checkWellformedPost(_, tree.srcPos))
+
+          case tree1 @ TypeApply(fn, args) if disallowGlobal =>
+            for arg <- args do
+              //println(i"checking $arg in $tree: ${arg.tpe.captureSet}")
+              checkNotGlobal(arg, args*)
+          case _ =>
+        traverseChildren(tree)
 
-  inline val disallowGlobal = true
-
-  def checkNotGlobal(tree: Tree, allArgs: Tree*)(using Context): Unit =
-    if disallowGlobal then
-      tree match
-        case LambdaTypeTree(_, restpt) =>
-          checkNotGlobal(restpt, allArgs*)
-        case _ =>
-          for ref <- tree.tpe.captureSet.elems do
-            val isGlobal = ref match
-              case ref: TermRef =>
-                ref.isRootCapability || ref.prefix != NoPrefix && ref.symbol.hasAnnotation(defn.AbilityAnnot)
-              case _ => false
-            val what = if ref.isRootCapability then "universal" else "global"
-            if isGlobal then
-              val notAllowed = i" is not allowed to capture the $what capability $ref"
-              def msg = tree match
-                case tree: InferredTypeTree =>
-                  i"""inferred type argument ${tree.tpe}$notAllowed
-                      |
-                      |The inferred arguments are: [$allArgs%, %]"""
-                case _ => s"type argument$notAllowed"
-              report.error(msg, tree.srcPos)
-
-  object PostRefinerCheck extends TreeTraverser:
-    def traverse(tree: Tree)(using Context) =
-      tree match
-        case tree1 @ TypeApply(fn, args) if disallowGlobal =>
-          for arg <- args do
-            //println(i"checking $arg in $tree: ${arg.tpe.captureSet}")
-            checkNotGlobal(arg, args*)
-        case _ =>
-      traverseChildren(tree)
-
-  def postRefinerCheck(tree: tpd.Tree)(using Context): Unit =
-    PostRefinerCheck.traverse(tree)
+    def postRefinerCheck(tree: tpd.Tree)(using Context): Unit =
+      PostRefinerCheck.traverse(tree)
+
+  end CaptureChecker
 
   override def show(tree: untpd.Tree)(using Context): String =
     val addRecheckedTypes = new TreeMap:

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

@@ -2490,6 +2490,8 @@ class Typer extends Namer
     val annot1 = typedExpr(tree.annot, defn.AnnotationClass.typeRef)
     val arg1 = typed(tree.arg, pt)
     if (ctx.mode is Mode.Type) {
+      if annot1.symbol.maybeOwner == defn.RetainsAnnot then
+        CheckCaptures.checkWellformed(annot1)
       if arg1.isType then
         assignType(cpy.Annotated(tree)(arg1, annot1), arg1, annot1)
       else

tests/neg-custom-args/captures/capt-wf-typer.scala

@@ -0,0 +1,10 @@
+class C
+type Cap = {*} C
+
+object foo
+
+def test(c: Cap, other: String): Unit =
+  val x7: {c} String = ??? // OK
+  val x8: String @retains(x7 + x7) = ??? // error
+  val x9: String @retains(foo) = ??? // error
+  ()

tests/neg-custom-args/captures/capt-wf.scala

@@ -0,0 +1,33 @@
+class C
+type Cap = {*} C
+
+object foo
+
+def test(c: Cap, other: String): Unit =
+  val x1: {*} C = ??? // OK
+  val x2: {other} C = ??? // error: cs is empty
+  val s1 = () => "abc"
+  val x3: {s1} C = ??? // error: cs is empty
+  val x3a: () => String = s1
+  val s2 = () => if x1 == null then "" else "abc"
+  val x4: {s2} C = ??? // OK
+  val x5: {c, c} C = ??? // warning: redundant
+  val x6: {c} {c} C = ??? // warning: redundant
+  val x7: {c} {*} C = ??? // warning: redundant
+  val x8: {*} {c} C = ??? // OK
+
+  def even(n: Int): Boolean = if n == 0 then true else odd(n - 1)
+  def odd(n: Int): Boolean = if n == 1 then true else even(n - 1)
+  val e1 = even
+  val o1 = odd
+
+  val y1: {e1} String = ??? // error cs is empty
+  val y2: {o1} String = ??? // error cs is empty
+
+  lazy val ev: (Int => Boolean) = (n: Int) =>
+    lazy val od: (Int => Boolean) = (n: Int) =>
+      if n == 1 then true else ev(n - 1)
+    if n == 0 then true else od(n - 1)
+  val y3: {ev} String = ??? // error cs is empty
+
+  ()

tests/new/test.scala

@@ -1,7 +1,3 @@
 object Test:
 
-  val x = ""
-
-  class C extends AnyRef @scala.retains(x)
-
   def test = ???

tests/pos-custom-args/captures/capt0.scala

@@ -1,7 +1,7 @@
 object Test:
 
   def test() =
-    val x = "abc"
+    val x: {*} Any = "abc"
     val y: Object @scala.retains(x) = ???
     val z: Object @scala.retains(x, *) = y: Object @scala.retains(x)
 

tests/pos-custom-args/captures/capt2.scala

@@ -3,11 +3,18 @@ class C
 type Cap = C @retains(*)
 
 def test1() =
-  val y = ""
+  val y: {*} String = ""
   def x: Object @retains(y) = y
 
 def test2() =
   val x: Cap = C()
   val y = () => { x; () }
   def z: (() => Unit) @retains(x) = y
-  z: (() => Unit) @retains(x)  // TODO: replace x with y
+  z: (() => Unit) @retains(x)
+  def z2: (() => Unit) @retains(y) = y
+  z2: (() => Unit) @retains(y)
+  val p: {*} () => String = () => "abc"
+  val q: {p} C = ???
+  p: ({p} () => String)
+
+