Implement @unboxed annotation exemption for reach capabilities of parameters

Author: odersky

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

@@ -40,6 +40,8 @@ object ccConfig:
    */
   inline val handleEtaExpansionsSpecially = false
 
+  val useUnboxedParams = true
+
   /** If true, use existential capture set variables */
   def useExistentials(using Context) =
     Feature.sourceVersion.stable.isAtLeast(SourceVersion.`3.5`)

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

@@ -62,6 +62,9 @@ object CheckCaptures:
         val res = cur
         cur = cur.outer
         res
+
+    def ownerString(using Context): String =
+      if owner.isAnonymousFunction then "enclosing function" else owner.show
   end Env
 
   /** Similar normal substParams, but this is an approximating type map that
@@ -386,21 +389,68 @@ class CheckCaptures extends Recheck, SymTransformer:
           val included = cs.filter: c =>
             c.stripReach match
               case ref: TermRef =>
-                val isVisible = isVisibleFromEnv(ref.symbol.owner)
-                if !isVisible && c.isReach then
+                //if c.isReach then println(i"REACH $c in ${env.owner}")
+                //assert(!env.owner.isAnonymousFunction)
+                val refSym = ref.symbol
+                val refOwner = refSym.owner
+                val isVisible = isVisibleFromEnv(refOwner)
+                if !isVisible && c.isReach && refSym.is(Param) && refOwner == env.owner then
+                if refSym.hasAnnotation(defn.UnboxedAnnot) then
+                  capt.println(i"exempt: $ref in $refOwner")
+                else
                   // Reach capabilities that go out of scope have to be approximated
-                  // by their underlyiong capture set. See i20503.scala.
-                  checkSubset(CaptureSet.ofInfo(c), env.captured, pos, provenance(env))
+                  // by their underlying capture set, which cannot be universal.
+                  // Reach capabilities of @unboxed parameters are exempted.
+                  val cs = CaptureSet.ofInfo(c)
+                  if ccConfig.useUnboxedParams then
+                    cs.disallowRootCapability: () =>
+                      report.error(em"Local reach capability $c leaks into capture scope of ${env.ownerString}", pos)
+                  checkSubset(cs, env.captured, pos, provenance(env))
                 isVisible
               case ref: ThisType => isVisibleFromEnv(ref.cls)
               case _ => false
-          capt.println(i"Include call or box capture $included from $cs in ${env.owner}")
           checkSubset(included, env.captured, pos, provenance(env))
+          capt.println(i"Include call or box capture $included from $cs in ${env.owner} --> ${env.captured}")
+    end markFree
 
     /** Include references captured by the called method in the current environment stack */
     def includeCallCaptures(sym: Symbol, pos: SrcPos)(using Context): Unit =
       if sym.exists && curEnv.isOpen then markFree(capturedVars(sym), pos)
 
+    private val prefixCalls = util.EqHashSet[GenericApply]()
+    private val unboxedArgs = util.EqHashSet[Tree]()
+
+    def handleCall(meth: Symbol, call: GenericApply, eval: () => Type)(using Context): Type =
+      if prefixCalls.remove(call) then return eval()
+
+      val unboxedParamNames =
+        meth.rawParamss.flatMap: params =>
+          params.collect:
+            case param if param.hasAnnotation(defn.UnboxedAnnot) =>
+              param.name
+        .toSet
+
+      def markUnboxedArgs(call: GenericApply): Unit = call.fun.tpe.widen match
+        case MethodType(pnames) =>
+          for (pname, arg) <- pnames.lazyZip(call.args) do
+            if unboxedParamNames.contains(pname) then
+              unboxedArgs.add(arg)
+        case _ =>
+
+      def markPrefixCalls(tree: Tree): Unit = tree match
+        case tree: GenericApply =>
+          prefixCalls.add(tree)
+          markUnboxedArgs(tree)
+          markPrefixCalls(tree.fun)
+        case _ =>
+
+      markUnboxedArgs(call)
+      markPrefixCalls(call.fun)
+      val res = eval()
+      includeCallCaptures(meth, call.srcPos)
+      res
+    end handleCall
+
     override def recheckIdent(tree: Ident, pt: Type)(using Context): Type =
       if tree.symbol.is(Method) then
         if tree.symbol.info.isParameterless then
@@ -470,7 +520,6 @@ class CheckCaptures extends Recheck, SymTransformer:
      */
     override def recheckApply(tree: Apply, pt: Type)(using Context): Type =
       val meth = tree.fun.symbol
-      includeCallCaptures(meth, tree.srcPos)
 
       // Unsafe box/unbox handlng, only for versions < 3.3
       def mapArgUsing(f: Type => Type) =
@@ -503,7 +552,7 @@ class CheckCaptures extends Recheck, SymTransformer:
             tp.derivedCapturingType(forceBox(parent), refs)
         mapArgUsing(forceBox)
       else
-        Existential.toCap(super.recheckApply(tree, pt)) match
+        handleCall(meth, tree, () => Existential.toCap(super.recheckApply(tree, pt))) match
           case appType @ CapturingType(appType1, refs) =>
             tree.fun match
               case Select(qual, _)
@@ -521,6 +570,13 @@ class CheckCaptures extends Recheck, SymTransformer:
           case appType => appType
     end recheckApply
 
+    override def recheckArg(arg: Tree, formal: Type)(using Context): Type =
+      val argType = recheck(arg, formal)
+      if unboxedArgs.remove(arg) && ccConfig.useUnboxedParams then
+        capt.println(i"charging deep capture set of $arg: ${argType} = ${CaptureSet.deepCaptureSet(argType)}")
+        markFree(CaptureSet.deepCaptureSet(argType), arg.srcPos)
+      argType
+
     private def isDistinct(xs: List[Type]): Boolean = xs match
       case x :: xs1 => xs1.isEmpty || !xs1.contains(x) && isDistinct(xs1)
       case Nil => true
@@ -589,20 +645,21 @@ class CheckCaptures extends Recheck, SymTransformer:
     end instantiate
 
     override def recheckTypeApply(tree: TypeApply, pt: Type)(using Context): Type =
+      val meth = tree.symbol
       if ccConfig.useSealed then
         val TypeApply(fn, args) = tree
         val polyType = atPhase(thisPhase.prev):
           fn.tpe.widen.asInstanceOf[TypeLambda]
         def isExempt(sym: Symbol) =
           sym.isTypeTestOrCast || sym == defn.Compiletime_erasedValue
         for case (arg: TypeTree, formal, pname) <- args.lazyZip(polyType.paramRefs).lazyZip((polyType.paramNames)) do
-          if !isExempt(tree.symbol) then
-            def where = if fn.symbol.exists then i" in an argument of ${fn.symbol}" else ""
+          if !isExempt(meth) then
+            def where = if meth.exists then i" in an argument of $meth" else ""
             disallowRootCapabilitiesIn(arg.knownType, NoSymbol,
               i"Sealed type variable $pname", "be instantiated to",
               i"This is often caused by a local capability$where\nleaking as part of its result.",
               tree.srcPos)
-      Existential.toCap(super.recheckTypeApply(tree, pt))
+      handleCall(meth, tree, () => Existential.toCap(super.recheckTypeApply(tree, pt)))
 
     override def recheckBlock(tree: Block, pt: Type)(using Context): Type =
       inNestedLevel(super.recheckBlock(tree, pt))

compiler/src/dotty/tools/dotc/transform/Recheck.scala

@@ -289,6 +289,9 @@ abstract class Recheck extends Phase, SymTransformer:
     /** A hook to massage the type of an applied method; currently not overridden */
     protected def prepareFunction(funtpe: MethodType, meth: Symbol)(using Context): MethodType = funtpe
 
+    protected def recheckArg(arg: Tree, formal: Type)(using Context): Type =
+      recheck(arg, formal)
+
     def recheckApply(tree: Apply, pt: Type)(using Context): Type =
       val funtpe0 = recheck(tree.fun)
       // reuse the tree's type on signature polymorphic methods, instead of using the (wrong) rechecked one
@@ -303,7 +306,7 @@ abstract class Recheck extends Phase, SymTransformer:
             else fntpe.paramInfos
           def recheckArgs(args: List[Tree], formals: List[Type], prefs: List[ParamRef]): List[Type] = args match
             case arg :: args1 =>
-              val argType = recheck(arg, normalizeByName(formals.head))
+              val argType = recheckArg(arg, normalizeByName(formals.head))
               val formals1 =
                 if fntpe.isParamDependent
                 then formals.tail.map(_.substParam(prefs.head, argType))

tests/neg/leak-problem.scala

@@ -0,0 +1,31 @@
+import language.experimental.captureChecking
+
+// Some capabilities that should be used locally
+trait Async:
+  //  some method
+  def read(): Unit
+def usingAsync[X](op: Async^ => X): X = ???
+
+case class Box[+T](get: T)
+
+def useBoxedAsync(x: Box[Async^]): Unit =
+  val t0 = x
+  val t1 = t0.get // error
+  t1.read()
+
+def useBoxedAsync1(x: Box[Async^]): Unit = x.get.read() // error
+
+def test(): Unit =
+  val useBoxedAsync2 = (x: Box[Async^]) =>
+    val t0 = x
+    val t1 = x.get // error
+    t1.read()
+
+  val f: Box[Async^] => Unit = (x:  Box[Async^]) => useBoxedAsync(x)
+
+  def boom(x: Async^): () ->{f} Unit =
+    () => f(Box(x))
+
+  val leaked = usingAsync[() ->{f} Unit](boom)
+
+  leaked()  // scope violation