Copy @consume annotations to the type

Author: Linyxus

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

@@ -395,6 +395,9 @@ extension (tp: Type)
     RefinedType(tp, name,
       AnnotatedType(rinfo, Annotation(defn.RefineOverrideAnnot, util.Spans.NoSpan)))
 
+  def dropUseAndConsumeAnnots(using Context): Type =
+    tp.dropAnnot(defn.UseAnnot).dropAnnot(defn.ConsumeAnnot)
+
 extension (tp: MethodType)
   /** A method marks an existential scope unless it is the prefix of a curried method */
   def marksExistentialScope(using Context): Boolean =
@@ -490,7 +493,7 @@ extension (sym: Symbol)
   def hasTrackedParts(using Context): Boolean =
     !CaptureSet.ofTypeDeeply(sym.info).isAlwaysEmpty
 
-  /** `sym` is annotated @use or it is a type parameter with a matching
+  /** `sym` itself or its info is annotated @use or it is a type parameter with a matching
    *  @use-annotated term parameter that contains `sym` in its deep capture set.
    */
   def isUseParam(using Context): Boolean =
@@ -503,6 +506,11 @@ extension (sym: Symbol)
                 case c: TypeRef => c.symbol == sym
                 case _ => false
 
+  /** `sym` or its info is annotated with `@consume`. */
+  def isConsumeParam(using Context): Boolean =
+    sym.hasAnnotation(defn.ConsumeAnnot)
+    || sym.info.hasAnnotation(defn.ConsumeAnnot)
+
   def isUpdateMethod(using Context): Boolean =
     sym.isAllOf(Mutable | Method, butNot = Accessor)
 

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

@@ -716,7 +716,7 @@ class CheckCaptures extends Recheck, SymTransformer:
         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 Some(ann) if !tp.hasAnnotation(cls) => AnnotatedType(tp, ann)
             case _ => tp
           copyAnnot(copyAnnot(formal, defn.UseAnnot), defn.ConsumeAnnot)
       funtpe.derivedLambdaType(paramInfos = paramInfosWithUses)
@@ -1616,7 +1616,10 @@ class CheckCaptures extends Recheck, SymTransformer:
       if noWiden(actual, expected) then
         actual
       else
-        val improvedVAR = improveCaptures(actual.widen.dealiasKeepAnnots, actual)
+        // Compute the widened type. Drop `@use` and `@consume` annotations from the type,
+        // since they obscures the capturing type.
+        val widened = actual.widen.dealiasKeepAnnots.dropUseAndConsumeAnnots
+        val improvedVAR = improveCaptures(widened, actual)
         val improved = improveReadOnly(improvedVAR, expected)
         val adapted = adaptBoxed(
             improved.withReachCaptures(actual), expected, tree,

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

@@ -620,7 +620,7 @@ class SepCheck(checker: CheckCaptures.CheckerAPI) extends tpd.TreeTraverser:
             if currentOwner.enclosingMethodOrClass.isProperlyContainedIn(refSym.maybeOwner.enclosingMethodOrClass) then
               report.error(em"""Separation failure: $descr non-local $refSym""", pos)
             else if refSym.is(TermParam)
-              && !refSym.hasAnnotation(defn.ConsumeAnnot)
+              && !refSym.isConsumeParam
               && currentOwner.isContainedIn(refSym.owner)
             then
               badParams += refSym
@@ -899,7 +899,7 @@ class SepCheck(checker: CheckCaptures.CheckerAPI) extends tpd.TreeTraverser:
     if !isUnsafeAssumeSeparate(tree) then trace(i"checking separate $tree"):
       checkUse(tree)
       tree match
-        case tree @ Select(qual, _) if tree.symbol.is(Method) && tree.symbol.hasAnnotation(defn.ConsumeAnnot) =>
+        case tree @ Select(qual, _) if tree.symbol.is(Method) && tree.symbol.isConsumeParam =>
           traverseChildren(tree)
           checkConsumedRefs(
               captures(qual).footprint(), qual.nuType,
@@ -962,4 +962,4 @@ class SepCheck(checker: CheckCaptures.CheckerAPI) extends tpd.TreeTraverser:
             consumeInLoopError(ref, pos)
         case _ =>
           traverseChildren(tree)
-end SepCheck
+end SepCheck

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

@@ -4234,8 +4234,11 @@ object Types extends TypeUtils {
         paramType = addAnnotation(paramType, defn.InlineParamAnnot, param)
       if param.is(Erased) then
         paramType = addAnnotation(paramType, defn.ErasedParamAnnot, param)
-      if param.isUseParam then
+      // Copy `@use` and `@consume` annotations from parameter symbols to the type.
+      if param.hasAnnotation(defn.UseAnnot) then
         paramType = addAnnotation(paramType, defn.UseAnnot, param)
+      if param.hasAnnotation(defn.ConsumeAnnot) then
+        paramType = addAnnotation(paramType, defn.ConsumeAnnot, param)
       paramType
 
     def adaptParamInfo(param: Symbol)(using Context): Type =

tests/neg-custom-args/captures/cc-annot-value-classes2.scala

@@ -0,0 +1,16 @@
+import language.experimental.captureChecking
+import caps.*
+trait Ref extends Mutable
+def kill(@consume x: Ref^): Unit = ()
+
+class C1:
+  def myKill(@consume x: Ref^): Unit = kill(x)  // ok
+
+class C2(val dummy: Int) extends AnyVal:
+  def myKill(@consume x: Ref^): Unit = kill(x)  // ok, too
+
+class C3:
+  def myKill(x: Ref^): Unit = kill(x)  // error
+
+class C4(val dummy: Int) extends AnyVal:
+  def myKill(x: Ref^): Unit = kill(x)  // error, too