Attempt to pass and check capability from parents correctly

Author: noti0na1

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

@@ -207,6 +207,17 @@ extension (tp: Type)
     case _: TypeRef | _: AppliedType => tp.typeSymbol.hasAnnotation(defn.CapabilityAnnot)
     case _ => false
 
+  /** Check if the class has universal capability, which means:
+   *  1. the class has a capability annotation,
+   *  2. the class is an impure function type,
+   *  3. or one of its base classes has universal capability.
+   */
+  def hasUniversalCapability(using Context): Boolean = tp match
+    case CapturingType(parent, ref) =>
+      ref.isUniversal || parent.hasUniversalCapability
+    case tp =>
+      tp.isCapabilityClassRef || tp.parents.exists(_.hasUniversalCapability)
+
   /** Drop @retains annotations everywhere */
   def dropAllRetains(using Context): Type = // TODO we should drop retains from inferred types before unpickling
     val tm = new TypeMap:

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

@@ -519,16 +519,6 @@ class CheckCaptures extends Recheck, SymTransformer:
       if sym.isConstructor then
         val cls = sym.owner.asClass
 
-        /** Check if the class or one of its parents has a root capability,
-         *  which means that the class has a capability annotation or an impure
-         *  function type.
-         */
-        def hasUniversalCapability(tp: Type): Boolean = tp match
-          case CapturingType(parent, ref) =>
-            ref.isUniversal || hasUniversalCapability(parent)
-          case tp =>
-            tp.isCapabilityClassRef || tp.parents.exists(hasUniversalCapability)
-
         /** First half of result pair:
          *  Refine the type of a constructor call `new C(t_1, ..., t_n)`
          *  to C{val x_1: T_1, ..., x_m: T_m} where x_1, ..., x_m are the tracked
@@ -538,7 +528,7 @@ class CheckCaptures extends Recheck, SymTransformer:
          */
         def addParamArgRefinements(core: Type, initCs: CaptureSet): (Type, CaptureSet) =
           var refined: Type = core
-          var allCaptures: CaptureSet = if hasUniversalCapability(core)
+          var allCaptures: CaptureSet = if core.hasUniversalCapability
             then CaptureSet.universal else initCs
           for (getterName, argType) <- mt.paramNames.lazyZip(argTypes) do
             val getter = cls.info.member(getterName).suchThat(_.is(ParamAccessor)).symbol

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

@@ -269,12 +269,6 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
           CapturingType(fntpe, cs, boxed = false)
         else fntpe
 
-      /** Map references to capability classes C to C^ */
-      private def expandCapabilityClass(tp: Type): Type =
-        if tp.isCapabilityClassRef
-        then CapturingType(tp, defn.expandedUniversalSet, boxed = false)
-        else tp
-
       private def recur(t: Type): Type = normalizeCaptures(mapOver(t))
 
       def apply(t: Type) =
@@ -297,7 +291,8 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
           case t: TypeVar =>
             this(t.underlying)
           case t =>
-            if t.isCapabilityClassRef
+            // Map references to capability classes C to C^
+            if t.hasUniversalCapability
             then CapturingType(t, defn.expandedUniversalSet, boxed = false)
             else recur(t)
     end expandAliases

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

@@ -895,13 +895,20 @@ class TypeComparer(@constructorOnly initctx: Context) extends ConstraintHandling
       canWidenAbstract && acc(true, tp)
 
     def tryBaseType(cls2: Symbol) =
-      val base = nonExprBaseType(tp1, cls2)
+      var base = nonExprBaseType(tp1, cls2)
       if base.exists && (base ne tp1)
           && (!caseLambda.exists
               || widenAbstractOKFor(tp2)
               || tp1.widen.underlyingClassRef(refinementOK = true).exists)
       then
         def checkBase =
+          // Strip existing capturing set from base type
+          base = base.stripCapturing
+          // Pass capture set of tp1 to base type
+          tp1 match
+            case tp1 @ CapturingType(_, refs1) =>
+              base = CapturingType(base, refs1, tp1.isBoxed)
+            case _ =>
           isSubType(base, tp2, if tp1.isRef(cls2) then approx else approx.addLow)
           && recordGadtUsageIf { MatchType.thatReducesUsingGadt(tp1) }
         if tp1.widenDealias.isInstanceOf[AndType] || base.isInstanceOf[OrType] then

tests/neg-custom-args/captures/extending-cap-classes.scala

@@ -0,0 +1,15 @@
+import annotation.capability
+
+class C1
+@capability class C2 extends C1
+class C3 extends C2
+
+def test =
+  val x1: C1 = new C1
+  val x2: C1 = new C2 // error
+  val x3: C1 = new C3 // error
+
+  val y1: C2 = new C2
+  val y2: C2 = new C3
+  
+  val z1: C3 = new C3

tests/neg-custom-args/captures/extending-impure-function.scala.scala

@@ -0,0 +1,18 @@
+class F extends (Int => Unit) {
+  def apply(x: Int): Unit = ()
+}
+
+def test =
+  val x1 = new (Int => Unit) {
+    def apply(x: Int): Unit = ()
+  }
+
+  val x2: Int -> Unit = new (Int => Unit) { // error
+    def apply(x: Int): Unit = ()
+  }
+
+  val y1: Int => Unit = new F
+  val y2: Int -> Unit = new F // error
+
+  val z1 = () => ()
+  val z2: () -> Unit = () => ()