Store capability class information in a hash map during cc

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

@@ -203,21 +203,6 @@ extension (tp: Type)
     case _ =>
       false
 
-  def isCapabilityClassRef(using Context) = tp.dealiasKeepAnnots match
-    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

@@ -528,7 +528,7 @@ class CheckCaptures extends Recheck, SymTransformer:
          */
         def addParamArgRefinements(core: Type, initCs: CaptureSet): (Type, CaptureSet) =
           var refined: Type = core
-          var allCaptures: CaptureSet = if core.hasUniversalCapability
+          var allCaptures: CaptureSet = if setup.isCapabilityClassRef(core)
             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

@@ -23,6 +23,7 @@ trait SetupAPI:
   def setupUnit(tree: Tree, recheckDef: DefRecheck)(using Context): Unit
   def isPreCC(sym: Symbol)(using Context): Boolean
   def postCheck()(using Context): Unit
+  def isCapabilityClassRef(tp: Type)(using Context): Boolean
 
 object Setup:
 
@@ -67,6 +68,31 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
     && !sym.owner.is(CaptureChecked)
     && !defn.isFunctionSymbol(sym.owner)
 
+  private val capabilityClassMap = new util.HashMap[Symbol, Boolean]
+
+  /** Check if the class is capability, which means:
+   *  1. the class has a capability annotation,
+   *  2. or at least one of its parent type has universal capability.
+   */
+  def isCapabilityClassRef(tp: Type)(using Context): Boolean = tp.dealiasKeepAnnots match
+    case _: TypeRef | _: AppliedType =>
+      val sym = tp.classSymbol
+      def checkSym: Boolean =
+        sym.hasAnnotation(defn.CapabilityAnnot)
+        || sym.info.parents.exists(hasUniversalCapability)
+      sym.isClass && capabilityClassMap.getOrElseUpdate(sym, checkSym)
+    case _ => false
+
+  private def hasUniversalCapability(tp: Type)(using Context): Boolean = tp.dealiasKeepAnnots match
+    case CapturingType(parent, refs) =>
+      refs.isUniversal || hasUniversalCapability(parent)
+    case AnnotatedType(parent, ann) =>
+      if ann.symbol.isRetains then
+        try ann.tree.toCaptureSet.isUniversal || hasUniversalCapability(parent)
+        catch case ex: IllegalCaptureRef => false
+      else hasUniversalCapability(parent)
+    case tp => isCapabilityClassRef(tp)
+
   private def fluidify(using Context) = new TypeMap with IdempotentCaptRefMap:
     def apply(t: Type): Type = t match
       case t: MethodType =>
@@ -292,7 +318,7 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
             this(t.underlying)
           case t =>
             // Map references to capability classes C to C^
-            if t.hasUniversalCapability
+            if isCapabilityClassRef(t)
             then CapturingType(t, defn.expandedUniversalSet, boxed = false)
             else recur(t)
     end expandAliases

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

@@ -0,0 +1,30 @@
+class F1 extends (Int => Unit) {
+  def apply(x: Int): Unit = ()
+}
+
+class F2 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 x3: Int -> Unit = new (Int -> Unit) {
+    def apply(x: Int): Unit = ()
+  }
+
+  val y1: Int => Unit = new F1
+  val y2: Int -> Unit = new F1 // error
+  val y3: Int => Unit = new F2
+  val y4: Int -> Unit = new F2
+
+  val z1 = () => ()
+  val z2: () -> Unit = () => ()
+  val z3: () -> Unit = z1
+  val z4: () => Unit = () => ()