Check that only @consume parameters flow to @consume parameters

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

@@ -427,10 +427,6 @@ extension (tp: Type)
           mapOver(t)
     tm(tp)
 
-  def hasUseAnnot(using Context): Boolean = tp match
-    case AnnotatedType(_, ann) => ann.symbol == defn.UseAnnot
-    case _ => false
-
   /** If `x` is a capture ref, its maybe capability `x?`, represented internally
    *  as `x @maybeCapability`. `x?` stands for a capability `x` that might or might
    *  not be part of a capture set. We have `{} <: {x?} <: {x}`. Maybe capabilities

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

@@ -655,11 +655,13 @@ class CheckCaptures extends Recheck, SymTransformer:
      *  on method parameter symbols to the corresponding paramInfo types.
      */
     override def prepareFunction(funtpe: MethodType, meth: Symbol)(using Context): MethodType =
-      val paramInfosWithUses = funtpe.paramInfos.zipWithConserve(funtpe.paramNames): (formal, pname) =>
-        val param = meth.paramNamed(pname)
-        param.getAnnotation(defn.UseAnnot) match
-          case Some(ann) => AnnotatedType(formal, ann)
-          case _ => formal
+      val paramInfosWithUses =
+        funtpe.paramInfos.zipWithConserve(funtpe.paramNames): (formal, pname) =>
+          val param = meth.paramNamed(pname)
+          def copyAnnot(tp: Type, cls: ClassSymbol) = param.getAnnotation(cls) match
+            case Some(ann) => AnnotatedType(tp, ann)
+            case _ => tp
+          copyAnnot(copyAnnot(formal, defn.UseAnnot), defn.ConsumeAnnot)
       funtpe.derivedLambdaType(paramInfos = paramInfosWithUses)
 
     /** Recheck applications, with special handling of unsafeAssumePure.
@@ -687,7 +689,7 @@ class CheckCaptures extends Recheck, SymTransformer:
       val freshenedFormal = Fresh.fromCap(formal)
       val argType = recheck(arg, freshenedFormal)
         .showing(i"recheck arg $arg vs $freshenedFormal", capt)
-      if formal.hasUseAnnot then
+      if formal.hasAnnotation(defn.UseAnnot) then
         // The @use annotation is added to `formal` by `prepareFunction`
         capt.println(i"charging deep capture set of $arg: ${argType} = ${argType.deepCaptureSet}")
         markFree(argType.deepCaptureSet, arg)
@@ -722,7 +724,7 @@ class CheckCaptures extends Recheck, SymTransformer:
       val qualCaptures = qualType.captureSet
       val argCaptures =
         for (argType, formal) <- argTypes.lazyZip(funType.paramInfos) yield
-          if formal.hasUseAnnot then argType.deepCaptureSet else argType.captureSet
+          if formal.hasAnnotation(defn.UseAnnot) then argType.deepCaptureSet else argType.captureSet
       appType match
         case appType @ CapturingType(appType1, refs)
         if qualType.exists
@@ -1569,10 +1571,11 @@ class CheckCaptures extends Recheck, SymTransformer:
             (params1, params2) <- member.rawParamss.lazyZip(other.rawParamss)
             (param1, param2) <- params1.lazyZip(params2)
           do
-            if param1.hasAnnotation(defn.UseAnnot) != param2.hasAnnotation(defn.UseAnnot) then
-              fail(i"has a parameter ${param1.name} with different @use status than the corresponding parameter in the overridden definition")
-            if param1.hasAnnotation(defn.ConsumeAnnot) != param2.hasAnnotation(defn.ConsumeAnnot) then
-              fail(i"has a parameter ${param1.name} with different @consume status than the corresponding parameter in the overridden definition")
+            def checkAnnot(cls: ClassSymbol) =
+              if param1.hasAnnotation(cls) != param2.hasAnnotation(cls) then
+                fail(i"has a parameter ${param1.name} with different @${cls.name} status than the corresponding parameter in the overridden definition")
+            checkAnnot(defn.UseAnnot)
+            checkAnnot(defn.ConsumeAnnot)
       end OverridingPairsCheckerCC
 
       def traverse(t: Tree)(using Context) =

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

@@ -11,6 +11,7 @@ import CaptureSet.{Refs, emptySet, HiddenSet}
 import config.Printers.capt
 import StdNames.nme
 import util.{SimpleIdentitySet, EqHashMap, SrcPos}
+import tpd.*
 
 object SepChecker:
 
@@ -31,15 +32,14 @@ object SepChecker:
   /** The kind of checked type, used for composing error messages */
   enum TypeKind:
     case Result(sym: Symbol, inferred: Boolean)
-    case Argument
+    case Argument(arg: Tree)
 
     def dclSym = this match
       case Result(sym, _) => sym
       case _ => NoSymbol
   end TypeKind
 
 class SepChecker(checker: CheckCaptures.CheckerAPI) extends tpd.TreeTraverser:
-  import tpd.*
   import checker.*
   import SepChecker.*
 
@@ -214,7 +214,7 @@ class SepChecker(checker: CheckCaptures.CheckerAPI) extends tpd.TreeTraverser:
     for (arg, idx) <- indexedArgs do
       if arg.needsSepCheck then
         val ac = formalCaptures(arg)
-        checkType(arg.formalType, arg.srcPos, TypeKind.Argument)
+        checkType(arg.formalType, arg.srcPos, TypeKind.Argument(arg))
         val hiddenInArg = ac.hidden.footprint
         //println(i"check sep $arg: $ac, footprint so far = $footprint, hidden = $hiddenInArg")
         val overlap = hiddenInArg.overlapWith(footprint).deductCapturesOf(deps(arg))
@@ -252,9 +252,9 @@ class SepChecker(checker: CheckCaptures.CheckerAPI) extends tpd.TreeTraverser:
       case TypeKind.Result(sym, inferred) =>
         def inferredStr = if inferred then " inferred" else ""
         def resultStr = if sym.info.isInstanceOf[MethodicType] then " result" else ""
-        i" $sym's$inferredStr$resultStr"
-      case TypeKind.Argument =>
-        " the argument's adapted type"
+        i"$sym's$inferredStr$resultStr"
+      case TypeKind.Argument(_) =>
+        "the argument's adapted"
 
     def explicitRefs(tp: Type): Refs = tp match
       case tp: (TermRef | ThisType) => SimpleIdentitySet(tp)
@@ -292,7 +292,7 @@ class SepChecker(checker: CheckCaptures.CheckerAPI) extends tpd.TreeTraverser:
               .nextOption
               .getOrElse(("", current, globalOverlap))
             report.error(
-              em"""Separation failure in$typeDescr type $tpe.
+              em"""Separation failure in $typeDescr type $tpe.
                   |One part,  $part , $nextRel  ${CaptureSet(next)}.
                   |A previous part$prevStr $prevRel  ${CaptureSet(prevRefs)}.
                   |The two sets overlap at  ${CaptureSet(overlap)}.""",
@@ -346,10 +346,10 @@ class SepChecker(checker: CheckCaptures.CheckerAPI) extends tpd.TreeTraverser:
             case t =>
               foldOver(c, t)
 
-    def checkParameters() =
+    def checkParams(refsToCheck: Refs, descr: => String) =
       val badParams = mutable.ListBuffer[Symbol]()
       def currentOwner = kind.dclSym.orElse(ctx.owner)
-      for hiddenRef <- prune(tpe.deepCaptureSet.elems.hidden.footprint) do
+      for hiddenRef <- prune(refsToCheck.footprint) do
         val refSym = hiddenRef.termSymbol
         if refSym.is(TermParam)
           && !refSym.hasAnnotation(defn.ConsumeAnnot)
@@ -364,25 +364,29 @@ class SepChecker(checker: CheckCaptures.CheckerAPI) extends tpd.TreeTraverser:
           case p :: ps => i"${p.name}, ${paramsStr(ps)}"
         val (pluralS, singleS) = if badParams.tail.isEmpty then ("", "s") else ("s", "")
         report.error(
-          em"""Separation failure:$typeDescr type $tpe hides parameter$pluralS ${paramsStr(badParams.toList)}
+          em"""Separation failure: $descr parameter$pluralS ${paramsStr(badParams.toList)}.
               |The parameter$pluralS need$singleS to be annotated with @consume to allow this.""",
             pos)
 
-    def flagHiddenParams =
-      kind match
-        case TypeKind.Result(sym, _) =>
-          !sym.isAnonymousFunction // we don't check return types of anonymous functions
-          && !sym.is(Case)         // We don't check so far binders in patterns since they
-                                   // have inferred universal types. TODO come back to this;
-                                   // either infer more precise types for such binders or
-                                   // "see through them" when we look at hidden sets.
-        case TypeKind.Argument =>
-          false
+    def checkParameters() = kind match
+      case TypeKind.Result(sym, _) =>
+        if !sym.isAnonymousFunction // we don't check return types of anonymous functions
+            && !sym.is(Case)        // We don't check so far binders in patterns since they
+                                    // have inferred universal types. TODO come back to this;
+                                    // either infer more precise types for such binders or
+                                    // "see through them" when we look at hidden sets.
+        then checkParams(tpe.deepCaptureSet.elems.hidden, i"$typeDescr type $tpe hides")
+      case TypeKind.Argument(arg) =>
+        if tpe.hasAnnotation(defn.ConsumeAnnot) then
+          val capts = captures(arg)
+          def descr(verb: String) = i"argument to @consume parameter with type ${arg.nuType} $verb"
+          checkParams(capts, descr("refers to"))
+          checkParams(capts.hidden, descr("hides"))
 
     if !tpe.hasAnnotation(defn.UntrackedCapturesAnnot) then
       traverse(Captures.None, tpe)
       traverse.toCheck.foreach(checkParts)
-      if flagHiddenParams then checkParameters()
+      checkParameters()
   end checkType
 
   private def collectMethodTypes(tp: Type): List[TermLambda] = tp match

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

@@ -528,7 +528,7 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
         case _ =>
           traverseChildren(tree)
       postProcess(tree)
-      checkProperUse(tree)
+      checkProperUseOrConsume(tree)
     end traverse
 
     /** Processing done on node `tree` after its children are traversed */
@@ -682,16 +682,22 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
       case _ =>
     end postProcess
 
-    /** Check that @use annotations only appear on parameters and not on anonymous function parameters */
-    def checkProperUse(tree: Tree)(using Context): Unit = tree match
+    /** Check that @use and @consume annotations only appear on parameters and not on
+     *  anonymous function parameters
+     */
+    def checkProperUseOrConsume(tree: Tree)(using Context): Unit = tree match
       case tree: MemberDef =>
-        def useAllowed(sym: Symbol) =
-          (sym.is(Param) || sym.is(ParamAccessor)) && !sym.owner.isAnonymousFunction
         for ann <- tree.symbol.annotations do
-          if ann.symbol == defn.UseAnnot && !useAllowed(tree.symbol) then
-            report.error(i"Only parameters of methods can have @use annotations", tree.srcPos)
+          def isAllowedFor(sym: Symbol) =
+            (sym.is(Param) || sym.is(ParamAccessor))
+            && (ann.symbol != defn.ConsumeAnnot || sym.isTerm)
+            && !sym.owner.isAnonymousFunction
+          def termStr =
+            if ann.symbol == defn.ConsumeAnnot then " term" else ""
+          if defn.ccParamOnlyAnnotations.contains(ann.symbol) && !isAllowedFor(tree.symbol) then
+            report.error(i"Only$termStr parameters of methods can have @${ann.symbol.name} annotations", tree.srcPos)
       case _ =>
-    end checkProperUse
+    end checkProperUseOrConsume
   end setupTraverser
 
 // --------------- Adding capture set variables ----------------------------------

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

@@ -1117,6 +1117,8 @@ class Definitions {
   @tu lazy val SilentAnnots: Set[Symbol] =
     Set(InlineParamAnnot, ErasedParamAnnot, RefineOverrideAnnot)
 
+  @tu lazy val ccParamOnlyAnnotations: Set[Symbol] = Set(UseAnnot, ConsumeAnnot)
+
   // A list of annotations that are commonly used to indicate that a field/method argument or return
   // type is not null. These annotations are used by the nullification logic in JavaNullInterop to
   // improve the precision of type nullification.

tests/neg-custom-args/captures/bad-uses-2.scala

@@ -1,7 +1,13 @@
-import caps.use
-class Test:
+import caps.{use, consume}
+class TestUse:
   @use def F = ???  // error
   @use val x = ???  // error
   @use type T       // error
   def foo[@use T](@use c: T): Unit = ??? // OK
 
+class TestConsume:
+  @consume def F = ???  // error
+  @consume val x = ???  // error
+  @consume type T       // error
+  def foo[@consume T](@use c: T): Unit = ??? // error
+

tests/neg-custom-args/captures/capt-depfun.check

@@ -8,5 +8,5 @@
 -- Error: tests/neg-custom-args/captures/capt-depfun.scala:11:24 -------------------------------------------------------
 11 |  val dc: ((Str^{y, z}) => Str^{y, z}) = ac(g()) // error // error: separatioon
    |           ^^^^^^^^^^^^^^^^^^^^^^^^^^
-   |           Separation failure: value dc's type Str^{y, z} => Str^{y, z} hides parameters y and z
+   |           Separation failure: value dc's type Str^{y, z} => Str^{y, z} hides parameters y and z.
    |           The parameters need to be annotated with @consume to allow this.

tests/neg-custom-args/captures/depfun-reach.check

@@ -15,5 +15,5 @@
 -- Error: tests/neg-custom-args/captures/depfun-reach.scala:12:17 ------------------------------------------------------
 12 |               : (xs: List[(X, () ->{io} Unit)]) => List[() ->{} Unit] = // error
    |                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-   |Separation failure: method foo's result type (xs: List[(X, box () ->{io} Unit)]) => List[() -> Unit] hides parameter op
+   |Separation failure: method foo's result type (xs: List[(X, box () ->{io} Unit)]) => List[() -> Unit] hides parameter op.
    |The parameter needs to be annotated with @consume to allow this.

tests/neg-custom-args/captures/i15772.check

@@ -42,5 +42,5 @@
 -- Error: tests/neg-custom-args/captures/i15772.scala:34:10 ------------------------------------------------------------
 34 |  def c : C^ = new C(x) // error separation
    |          ^^
-   |          Separation failure: method c's result type C^ hides parameter x
+   |          Separation failure: method c's result type C^ hides parameter x.
    |          The parameter needs to be annotated with @consume to allow this.

tests/neg-custom-args/captures/i19330.check

@@ -13,5 +13,5 @@
 -- Error: tests/neg-custom-args/captures/i19330.scala:16:14 ------------------------------------------------------------
 16 |    val t: () => Logger^ = () => l // error
    |           ^^^^^^^^^^^^^
-   |           Separation failure: value t's type () => (ex$5: caps.Exists) -> Logger^{ex$5} hides parameter l
+   |           Separation failure: value t's type () => (ex$5: caps.Exists) -> Logger^{ex$5} hides parameter l.
    |           The parameter needs to be annotated with @consume to allow this.

tests/neg-custom-args/captures/i21442.check

@@ -10,5 +10,5 @@
 -- Error: tests/neg-custom-args/captures/i21442.scala:17:10 ------------------------------------------------------------
 17 |  val x1: Boxed[IO^] = x // error
    |          ^^^^^^^^^^
-   |          Separation failure: value x1's type Boxed[box IO^] hides parameter x
+   |          Separation failure: value x1's type Boxed[box IO^] hides parameter x.
    |          The parameter needs to be annotated with @consume to allow this.

tests/neg-custom-args/captures/sepchecks2.check

@@ -36,17 +36,17 @@
 -- Error: tests/neg-custom-args/captures/sepchecks2.scala:27:6 ---------------------------------------------------------
 27 |  bar((c, c)) // error
    |      ^^^^^^
-   |      Separation failure in the argument's adapted type type (box Object^, box Object^).
+   |      Separation failure in the argument's adapted type (box Object^, box Object^).
    |      One part,  box Object^ , hides  {c}.
    |      A previous part,  box Object^ ,  also hides  {c}.
    |      The two sets overlap at  {c}.
 -- Error: tests/neg-custom-args/captures/sepchecks2.scala:30:9 ---------------------------------------------------------
 30 |  val x: (Object^, Object^{c}) = (d, c) // error
    |         ^^^^^^^^^^^^^^^^^^^^^
-   |         Separation failure: value x's type (box Object^, box Object^{c}) hides parameter d
+   |         Separation failure: value x's type (box Object^, box Object^{c}) hides parameter d.
    |         The parameter needs to be annotated with @consume to allow this.
 -- Error: tests/neg-custom-args/captures/sepchecks2.scala:33:9 ---------------------------------------------------------
 33 |  val x: (Object^, Object^) = (c, d) // error
    |         ^^^^^^^^^^^^^^^^^^
-   |         Separation failure: value x's type (box Object^, box Object^) hides parameters c and d
+   |         Separation failure: value x's type (box Object^, box Object^) hides parameters c and d.
    |         The parameters need to be annotated with @consume to allow this.

tests/neg-custom-args/captures/sepchecks4.check

@@ -1,12 +1,12 @@
 -- Error: tests/neg-custom-args/captures/sepchecks4.scala:8:12 ---------------------------------------------------------
 8 |  val x: () => Unit = () => println(io) // error
   |         ^^^^^^^^^^
-  |         Separation failure: value x's type () => Unit hides parameter io
+  |         Separation failure: value x's type () => Unit hides parameter io.
   |         The parameter needs to be annotated with @consume to allow this.
 -- Error: tests/neg-custom-args/captures/sepchecks4.scala:7:25 ---------------------------------------------------------
 7 |def bad(io: Object^): () => Unit =    // error
   |                      ^^^^^^^^^^
-  |                      Separation failure: method bad's result type () => Unit hides parameter io
+  |                      Separation failure: method bad's result type () => Unit hides parameter io.
   |                      The parameter needs to be annotated with @consume to allow this.
 -- Error: tests/neg-custom-args/captures/sepchecks4.scala:12:25 --------------------------------------------------------
 12 |  par(() => println(io))(() => println(io)) // error // (1)

tests/neg-custom-args/captures/sepchecks5.check

@@ -0,0 +1,10 @@
+-- Error: tests/neg-custom-args/captures/sepchecks5.scala:12:37 --------------------------------------------------------
+12 |def bad(io: Object^): () => Unit = f(io) // error
+   |                                     ^^
+   |             Separation failure: argument to @consume parameter with type (io : Object^) refers to parameter io.
+   |             The parameter needs to be annotated with @consume to allow this.
+-- Error: tests/neg-custom-args/captures/sepchecks5.scala:19:13 --------------------------------------------------------
+19 |  val f2 = g(io)              // error
+   |             ^^
+   |             Separation failure: argument to @consume parameter with type (io : Object^) refers to parameter io.
+   |             The parameter needs to be annotated with @consume to allow this.

tests/neg-custom-args/captures/sepchecks5.scala

@@ -0,0 +1,21 @@
+import caps.{cap, consume}
+import language.future
+import language.experimental.captureChecking
+
+def par(op1: () => Unit)(op2: () => Unit): Unit = ()
+
+def f(@consume io: Object^): () => Unit =
+  () => println(io)
+
+def g(@consume io: Object^): () => Unit = f(io) // ok
+
+def bad(io: Object^): () => Unit = f(io) // error
+
+def test(io: Object^): Unit =
+
+  val f1 = bad(io)
+  par(f1)(() => println(io))  // !!! separation failure
+
+  val f2 = g(io)              // error
+  par(f2)(() => println(io))  // !!! separation failure
+

tests/neg-custom-args/captures/unsound-reach-6.check

@@ -0,0 +1,15 @@
+-- Error: tests/neg-custom-args/captures/unsound-reach-6.scala:7:13 ----------------------------------------------------
+7 |  println(xs.head) // error
+  |          ^^^^^^^
+  |          Local reach capability xs* leaks into capture scope of method f.
+  |          To allow this, the parameter xs should be declared with a @use annotation
+-- Error: tests/neg-custom-args/captures/unsound-reach-6.scala:11:14 ---------------------------------------------------
+11 |    val z = f(ys)   // error @consume failure
+   |              ^^
+   |Separation failure: argument to @consume parameter with type (ys : List[box () ->{io} Unit]) refers to parameters ys and io.
+   |The parameters need to be annotated with @consume to allow this.
+-- Error: tests/neg-custom-args/captures/unsound-reach-6.scala:19:14 ---------------------------------------------------
+19 |    val z = f(ys)  // error @consume failure
+   |              ^^
+   |Separation failure: argument to @consume parameter with type (ys : -> List[box () ->{io} Unit]) refers to parameter io.
+   |The parameter needs to be annotated with @consume to allow this.

tests/neg-custom-args/captures/unsound-reach-6.scala

@@ -8,15 +8,15 @@ def f(@consume xs: List[() => Unit]): () => Unit = () =>
 
 def test(io: IO^)(ys: List[() ->{io} Unit]) =
   val x = () =>
-    val z = f(ys)
+    val z = f(ys)   // error @consume failure
     z()
   val _: () -> Unit = x // !!! ys* gets lost
   ()
 
 def test(io: IO^) =
   def ys: List[() ->{io} Unit] = ???
   val x = () =>
-    val z = f(ys)
+    val z = f(ys)  // error @consume failure
     z()
   val _: () -> Unit = x // !!! io gets lost
   ()