Better handling of vars with universal capture sets

 - add unsafeBox operation so that a variable with universal capture set
   can be initialized and assigned to
 - infer unsafebox

Author: odersky

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

@@ -96,16 +96,16 @@ extension (tp: Type)
   /** Is the boxedCaptureSet of this type nonempty? */
   def isBoxedCapturing(using Context) = !tp.boxedCaptureSet.isAlwaysEmpty
 
-  /** If this type is a boxed capturing type, its unboxed version
-   *  If it is a TermRef of boxed capturing type, an unboxed capturing
-   *  type capturing the TermRef.
+  /** If this type is a capturing type, the version with boxed statues as given by `boxed`.
+   *  If it is a TermRef of a capturing type, and the box status flips, widen to a capturing
+   *  type that captures the TermRef.
    */
-  def unbox(using Context): Type = tp.widenDealias match
-    case tp @ CapturingType(parent, refs) if tp.isBoxed =>
+  def forceBoxStatus(boxed: Boolean)(using Context): Type = tp.widenDealias match
+    case tp @ CapturingType(parent, refs) if tp.isBoxed != boxed =>
       val refs1 = tp match
         case ref: CaptureRef if ref.isTracked => ref.singletonCaptureSet
         case _ => refs
-      CapturingType(parent, refs1, boxed = false)
+      CapturingType(parent, refs1, boxed)
     case _ =>
       tp
 
@@ -168,6 +168,7 @@ extension (sym: Symbol)
       case _ => false
     containsEnclTypeParam(sym.info.finalResultType)
     && !sym.allowsRootCapture
+    && sym != defn.Caps_unsafeBox
     && sym != defn.Caps_unsafeUnbox
 
 extension (tp: AnnotatedType)

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

@@ -332,29 +332,31 @@ class CheckCaptures extends Recheck, SymTransformer:
      *  and Cr otherwise.
      */
     override def recheckApply(tree: Apply, pt: Type)(using Context): Type =
-      includeCallCaptures(tree.symbol, tree.srcPos)
-      tree match
-        case Apply(fn, arg :: Nil) if fn.symbol == defn.Caps_unsafeUnbox =>
-          val argType0 = recheckStart(arg, pt).unbox
-          val argType = super.recheckFinish(argType0, arg, pt)
-          super.recheckFinish(argType, tree, pt)
-        case _ =>
-          super.recheckApply(tree, pt) match
-            case appType @ CapturingType(appType1, refs) =>
-              tree.fun match
-                case Select(qual, _)
-                if !tree.fun.symbol.isConstructor
-                    && !qual.tpe.isBoxedCapturing
-                    && !tree.args.exists(_.tpe.isBoxedCapturing)
-                    && qual.tpe.captureSet.mightSubcapture(refs)
-                    && tree.args.forall(_.tpe.captureSet.mightSubcapture(refs))
-                =>
-                  val callCaptures = tree.args.foldLeft(qual.tpe.captureSet)((cs, arg) =>
-                    cs ++ arg.tpe.captureSet)
-                  appType.derivedCapturingType(appType1, callCaptures)
-                    .showing(i"narrow $tree: $appType, refs = $refs, qual = ${qual.tpe.captureSet} --> $result", capt)
-                case _ => appType
-            case appType => appType
+      val meth = tree.fun.symbol
+      includeCallCaptures(meth, tree.srcPos)
+      if meth == defn.Caps_unsafeBox || meth == defn.Caps_unsafeUnbox then
+        val arg :: Nil = tree.args: @unchecked
+        val argType0 = recheckStart(arg, pt)
+          .forceBoxStatus(boxed = meth == defn.Caps_unsafeBox)
+        val argType = super.recheckFinish(argType0, arg, pt)
+        super.recheckFinish(argType, tree, pt)
+      else
+        super.recheckApply(tree, pt) match
+          case appType @ CapturingType(appType1, refs) =>
+            tree.fun match
+              case Select(qual, _)
+              if !tree.fun.symbol.isConstructor
+                  && !qual.tpe.isBoxedCapturing
+                  && !tree.args.exists(_.tpe.isBoxedCapturing)
+                  && qual.tpe.captureSet.mightSubcapture(refs)
+                  && tree.args.forall(_.tpe.captureSet.mightSubcapture(refs))
+              =>
+                val callCaptures = tree.args.foldLeft(qual.tpe.captureSet)((cs, arg) =>
+                  cs ++ arg.tpe.captureSet)
+                appType.derivedCapturingType(appType1, callCaptures)
+                  .showing(i"narrow $tree: $appType, refs = $refs, qual = ${qual.tpe.captureSet} --> $result", capt)
+              case _ => appType
+          case appType => appType
     end recheckApply
 
     /** Handle an application of method `sym` with type `mt` to arguments of types `argTypes`.
@@ -460,10 +462,25 @@ class CheckCaptures extends Recheck, SymTransformer:
         case _ =>
       super.recheckBlock(block, pt)
 
+    /** If `rhsProto` has `*` as its capture set, wrap `rhs` in a `unsafeBox`.
+     *  Used to infer `unsafeBox` for expressions that get assigned to variables
+     *  that have universal capture set.
+     */
+    def maybeBox(rhs: Tree, rhsProto: Type)(using Context): Tree =
+      if rhsProto.captureSet.isUniversal then
+        ref(defn.Caps_unsafeBox).appliedToType(rhsProto).appliedTo(rhs)
+      else rhs
+
+    override def recheckAssign(tree: Assign)(using Context): Type =
+      val rhsProto = recheck(tree.lhs).widen
+      recheck(maybeBox(tree.rhs, rhsProto), rhsProto)
+      defn.UnitType
+
     override def recheckValDef(tree: ValDef, sym: Symbol)(using Context): Unit =
       try
         if !sym.is(Module) then // Modules are checked by checking the module class
-          super.recheckValDef(tree, sym)
+          if sym.is(Mutable) then recheck(maybeBox(tree.rhs, sym.info), sym.info)
+          else super.recheckValDef(tree, sym)
       finally
         if !sym.is(Param) then
           // Parameters with inferred types belong to anonymous methods. We need to wait
@@ -765,7 +782,6 @@ class CheckCaptures extends Recheck, SymTransformer:
             recon(CapturingType(parent1, cs1, actualIsBoxed))
       }
 
-
       var actualw = actual.widenDealias
       actual match
         case ref: CaptureRef if ref.isTracked =>

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

@@ -961,6 +961,7 @@ class Definitions {
   def RuntimeTupleFunctionsModule(using Context): Symbol = requiredModule("scala.runtime.TupledFunctions")
 
   @tu lazy val CapsModule: Symbol = requiredModule("scala.caps")
+    @tu lazy val Caps_unsafeBox: Symbol = CapsModule.requiredMethod("unsafeBox")
     @tu lazy val Caps_unsafeUnbox: Symbol = CapsModule.requiredMethod("unsafeUnbox")
 
   // Annotation base classes

library/src/scala/caps.scala

@@ -5,6 +5,11 @@ import annotation.experimental
 // @experimental , suppress @experimental so we can use in compiler itself
 object caps:
 
+  /** If argument is of type `cs T`, converts to type `box cs T`. This
+   *  avoids the error that would be raised when boxing `*`.
+   */
+  extension [T](x: T) def unsafeBox: T = x
+
   /** If argument is of type `box cs T`, converts to type `cs T`. This
    *  avoids the error that would be raised when unboxing `*`.
    */

tests/pos-custom-args/captures/vars1.scala

@@ -1,8 +1,25 @@
-import caps.unsafeUnbox
-type ErrorHandler = (Int, String) => Unit
-
-var defaultIncompleteHandler: ErrorHandler = ???
-var incompleteHandler: ErrorHandler = defaultIncompleteHandler
-val x = incompleteHandler.unsafeUnbox
-val _ : ErrorHandler = x
-val _ = x(1, "a")
+import caps.*
+
+object Test:
+  type ErrorHandler = (Int, String) => Unit
+
+  var defaultIncompleteHandler: ErrorHandler = ???
+  var incompleteHandler: ErrorHandler = defaultIncompleteHandler
+  val x = incompleteHandler.unsafeUnbox
+  val _ : ErrorHandler = x
+  val _ = x(1, "a")
+
+  def defaultIncompleteHandler1(): ErrorHandler = ???
+  val defaultIncompleteHandler2: ErrorHandler = ???
+  var incompleteHandler1: ErrorHandler = defaultIncompleteHandler1()
+  var incompleteHandler2: ErrorHandler = defaultIncompleteHandler2
+  var incompleteHandler3: ErrorHandler = defaultIncompleteHandler1().unsafeBox
+  var incompleteHandler4: ErrorHandler = defaultIncompleteHandler2.unsafeBox
+  private var incompleteHandler5 = defaultIncompleteHandler1()
+  private var incompleteHandler6 = defaultIncompleteHandler2
+  private var incompleteHandler7 = defaultIncompleteHandler1().unsafeBox
+  private var incompleteHandler8 = defaultIncompleteHandler2.unsafeBox
+
+  incompleteHandler1 = defaultIncompleteHandler2
+  incompleteHandler1 = defaultIncompleteHandler2.unsafeBox
+  val saved = incompleteHandler1.unsafeUnbox