Cleanup transformation of inferred types

Author: odersky

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

@@ -63,20 +63,6 @@ extension (tp: Type)
 
  def isBoxedCapturing(using Context) = !tp.boxedCaptured.isAlwaysEmpty
 
- def canHaveInferredCapture(using Context): Boolean = tp match
- case tp: TypeRef if tp.symbol.isClass =>
- !tp.symbol.isValueClass && tp.symbol != defn.AnyClass
- case _: TypeVar | _: TypeParamRef =>
- false
- case tp: TypeProxy =>
- tp.superType.canHaveInferredCapture
- case tp: AndType =>
- tp.tp1.canHaveInferredCapture && tp.tp2.canHaveInferredCapture
- case tp: OrType =>
- tp.tp1.canHaveInferredCapture || tp.tp2.canHaveInferredCapture
- case _ =>
- false
-
  def stripCapturing(using Context): Type = tp.dealiasKeepAnnots match
  case CapturingType(parent, _, _) =>
  parent.stripCapturing

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

@@ -117,85 +117,150 @@ class CheckCaptures extends Recheck:
 
  override def transformType(tp: Type, inferred: Boolean, boxed: Boolean)(using Context): Type =
 
- def addInnerVars(tp: Type): Type = tp match
- case tp @ AppliedType(tycon, args) =>
- tp.derivedAppliedType(tycon, args.map(addVars(_, boxed = true)))
- case tp @ RefinedType(core, rname, rinfo) =>
- val rinfo1 = addVars(rinfo)
- if defn.isFunctionType(tp) then
- rinfo1.toFunctionType(isJava = false, alwaysDependent = true)
-  else
-  tp.derivedRefinedType(addInnerVars(core), rname, rinfo1)
- case tp: MethodType =>
- tp.derivedLambdaType(
-  paramInfos = tp.paramInfos.mapConserve(addVars(_)),
-  resType = addVars(tp.resType))
-  case tp: PolyType =>
- tp.derivedLambdaType(
- resType = addVars(tp.resType))
-  case tp: ExprType =>
- tp.derivedExprType(resType = addVars(tp.resType))
-  case _ =>
- tp
-
- /** Turn plain function types into dependent function types, so that
- * we can refer to the parameter in capture sets
+ def depFun(tycon: Type, argTypes: List[Type], resType: Type): Type =
+ MethodType.companion(
+  isContextual = defn.isContextFunctionClass(tycon.classSymbol),
+  isErased = defn.isErasedFunctionClass(tycon.classSymbol)
+ )(argTypes, resType)
+ .toFunctionType(isJava = false, alwaysDependent = true)
+
+ def box(tp: Type): Type = tp match
+ case CapturingType(parent, refs, false) => CapturingType(parent, refs, true)
+ case _ => tp
+
+ /** Perform the following transformation steps everywhere in a type:
+ * 1. Drop retains annotations
+ * 2. Turn plain function types into dependent function types, so that
+ *  we can refer to their parameter in capture sets. Currently this is
+ * only done at the toplevel, i.e. for function types that are not
+ * themselves argument types of other function types. Without this restriction
+ *  boxmap-paper.scala fails. Need to figure out why.
+ * 3. Refine other class types C by adding capture set variables to their parameter getters
+ *  (see addCaptureRefinements)
+ * 4. Add capture set variables to all types that can be tracked
+  *
+ * Polytype bounds are only cleaned using step 1, but not otherwise transformed.
  */
- def addFunctionRefinements(tp: Type): Type = tp match
- case tp @ AppliedType(tycon, args) =>
- if defn.isNonRefinedFunction(tp) then
- MethodType.companion(
- isContextual = defn.isContextFunctionClass(tycon.classSymbol),
- isErased = defn.isErasedFunctionClass(tycon.classSymbol)
- )(args.init, addFunctionRefinements(args.last))
- .toFunctionType(isJava = false, alwaysDependent = true)
- .showing(i"add function refinement $tp --> $result", capt)
- else
- tp.derivedAppliedType(tycon, args.map(addFunctionRefinements(_)))
- case tp @ RefinedType(core, rname, rinfo) if !defn.isFunctionType(tp) =>
- tp.derivedRefinedType(
- addFunctionRefinements(core), rname, addFunctionRefinements(rinfo))
- case tp: MethodOrPoly =>
- tp.derivedLambdaType(resType = addFunctionRefinements(tp.resType))
- case tp: ExprType =>
- tp.derivedExprType(resType = addFunctionRefinements(tp.resType))
- case _ =>
- tp
+ def mapInferred = new TypeMap:
 
- /** Refine a possibly applied class type C where the class has tracked parameters
- * x_1: T_1, ..., x_n: T_n to C { val x_1: CV_1 T_1, ..., val x_n: CV_n T_n }
- * where CV_1, ..., CV_n are fresh capture sets.
- */
- def addCaptureRefinements(tp: Type): Type = tp.stripped match
- case _: TypeRef | _: AppliedType if tp.typeSymbol.isClass =>
- val cls = tp.typeSymbol.asClass
- cls.paramGetters.foldLeft(tp) { (core, getter) =>
- if getter.termRef.isTracked then
- val getterType = tp.memberInfo(getter).strippedDealias
- RefinedType(core, getter.name, CapturingType(getterType, CaptureSet.Var(), boxed = false))
- .showing(i"add capture refinement $tp --> $result", capt)
- else
- core
- }
- case _ =>
- tp
+ /** Drop @retains annotations everywhere */
+ object cleanup extends TypeMap:
+ def apply(t: Type) = t match
+ case AnnotatedType(parent, annot) if annot.symbol == defn.RetainsAnnot =>
+ apply(parent)
+ case _ =>
+ mapOver(t)
 
- def addVars(tp: Type, boxed: Boolean = false): Type =
- var tp1 = addInnerVars(tp)
- val tp2 = addCaptureRefinements(tp1)
- if tp1.canHaveInferredCapture
- then CapturingType(tp2, CaptureSet.Var(), boxed)
- else tp2
+ /** Refine a possibly applied class type C where the class has tracked parameters
+ * x_1: T_1, ..., x_n: T_n to C { val x_1: CV_1 T_1, ..., val x_n: CV_n T_n }
+ * where CV_1, ..., CV_n are fresh capture sets.
+ */
+ def addCaptureRefinements(tp: Type): Type = tp match
+ case _: TypeRef | _: AppliedType if tp.typeParams.isEmpty =>
+ tp.typeSymbol match
+ case cls: ClassSymbol if !defn.isFunctionClass(cls) =>
+ cls.paramGetters.foldLeft(tp) { (core, getter) =>
+ if getter.termRef.isTracked then
+ val getterType = tp.memberInfo(getter).strippedDealias
+ RefinedType(core, getter.name, CapturingType(getterType, CaptureSet.Var(), boxed = false))
+ .showing(i"add capture refinement $tp --> $result", capt)
+ else
+ core
+ }
+ case _ => tp
+ case _ => tp
+
+ /** Should a capture set variable be added on type `tp`? */
+ def canHaveInferredCapture(tp: Type): Boolean =
+ tp.typeParams.isEmpty && tp.match
+ case tp: (TypeRef | AppliedType) =>
+ val sym = tp.typeSymbol
+ if sym.isClass then !sym.isValueClass && sym != defn.AnyClass
+ else canHaveInferredCapture(tp.superType.dealias)
+ case tp: (RefinedOrRecType | MatchType) =>
+ canHaveInferredCapture(tp.underlying)
+ case tp: AndType =>
+ canHaveInferredCapture(tp.tp1) && canHaveInferredCapture(tp.tp2)
+ case tp: OrType =>
+ canHaveInferredCapture(tp.tp1) || canHaveInferredCapture(tp.tp2)
+ case _ =>
+ false
+
+ /** Add a capture set variable to `tp` if necessary, or maybe pull out
+ * an embedded capture set variables from a part of `tp`.
+ */
+ def addVar(tp: Type) = tp match
+ case tp @ RefinedType(parent @ CapturingType(parent1, refs, boxed), rname, rinfo) =>
+ CapturingType(tp.derivedRefinedType(parent1, rname, rinfo), refs, boxed)
+ case tp: RecType =>
+ tp.parent match
+ case CapturingType(parent1, refs, boxed) =>
+ CapturingType(tp.derivedRecType(parent1), refs, boxed)
+ case _ =>
+ tp // can return `tp` here since unlike RefinedTypes, RecTypes are never created
+ // by `mapInferred`. Hence if the underlying type admits capture variables
+ // a variable was already added, and the first case above would apply.
+ case AndType(CapturingType(parent1, refs1, boxed1), CapturingType(parent2, refs2, boxed2)) =>
+ assert(refs1.asVar.elems.isEmpty)
+ assert(refs2.asVar.elems.isEmpty)
+ assert(boxed1 == boxed2)
+ CapturingType(AndType(parent1, parent2), refs1, boxed1)
+ case tp @ OrType(CapturingType(parent1, refs1, boxed1), CapturingType(parent2, refs2, boxed2)) =>
+ assert(refs1.asVar.elems.isEmpty)
+ assert(refs2.asVar.elems.isEmpty)
+ assert(boxed1 == boxed2)
+ CapturingType(OrType(parent1, parent2, tp.isSoft), refs1, boxed1)
+ case tp @ OrType(CapturingType(parent1, refs1, boxed1), tp2) =>
+ CapturingType(OrType(parent1, tp2, tp.isSoft), refs1, boxed1)
+ case tp @ OrType(tp1, CapturingType(parent2, refs2, boxed2)) =>
+ CapturingType(OrType(tp1, parent2, tp.isSoft), refs2, boxed2)
+ case _ if canHaveInferredCapture(tp) =>
+ CapturingType(tp, CaptureSet.Var(), boxed = false)
+ case _ =>
+ tp
 
- if inferred then
- val cleanup = new TypeMap:
- def apply(t: Type) = t match
+ var isTopLevel = true
+
+ def mapNested(ts: List[Type]): List[Type] =
+ val saved = isTopLevel
+ isTopLevel = false
+ try ts.mapConserve(this) finally isTopLevel = saved
+
+ def apply(t: Type) =
+ val t1 = t match
  case AnnotatedType(parent, annot) if annot.symbol == defn.RetainsAnnot =>
  apply(parent)
+ case tp @ AppliedType(tycon, args) =>
+ val tycon1 = this(tycon)
+ if defn.isNonRefinedFunction(tp) then
+ val args1 = mapNested(args.init)
+ val res1 = this(args.last)
+ if isTopLevel then
+ depFun(tycon1, args1, res1)
+ .showing(i"add function refinement $tp --> $result", capt)
+ else
+ tp.derivedAppliedType(tycon1, args1 :+ res1)
+ else
+ tp.derivedAppliedType(tycon1, args.mapConserve(arg => box(this(arg))))
+ case tp @ RefinedType(core, rname, rinfo) if defn.isFunctionType(tp) =>
+ apply(rinfo).toFunctionType(isJava = false, alwaysDependent = true)
+ case tp: MethodType =>
+ tp.derivedLambdaType(
+ paramInfos = mapNested(tp.paramInfos),
+ resType = this(tp.resType))
+ case tp: TypeLambda =>
+ // Don't recurse into parameter bounds, just cleanup any stray retains annotations
+ tp.derivedLambdaType(
+ paramInfos = tp.paramInfos.mapConserve(cleanup(_).bounds),
+ resType = this(tp.resType))
  case _ =>
  mapOver(t)
- addVars(addFunctionRefinements(cleanup(tp)), boxed)
- .showing(i"reinfer $tp --> $result", capt)
+ addVar(addCaptureRefinements(t1))
+ end mapInferred
+
+ if inferred then
+ val tp1 = mapInferred(tp)
+ if boxed then box(tp1) else tp1
  else
  def setBoxed(t: Type) = t match
  case AnnotatedType(_, annot) if annot.symbol == defn.RetainsAnnot =>

tests/neg-custom-args/captures/capt1.check

@@ -40,7 +40,7 @@ longer explanation available when compiling with `-explain`
 -- [E007] Type Mismatch Error: tests/neg-custom-args/captures/capt1.scala:31:24 ----------------------------------------
 31 | val z2 = h[() -> Cap](() => x)(() => C()) // error
  | ^^^^^^^
- | Found: {x} () -> ? Cap
+ | Found: {x} () -> Cap
  | Required: () -> Cap
 
 longer explanation available when compiling with `-explain`

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

@@ -1,15 +1,14 @@
-class CC
-type Cap = {*} CC
+@annotation.capability class Cap
 
-class LazyRef[T](val elem: {*} () => T):
+class LazyRef[T](val elem: () => T):
  val get = elem
- def map[U](f: {*} T => U): {f, this} LazyRef[U] =
+ def map[U](f: T => U): {f, this} LazyRef[U] =
  new LazyRef(() => f(elem()))
 
-def map[A, B](ref: {*} LazyRef[A], f: {*} A => B): {f, ref} LazyRef[B] =
+def map[A, B](ref: {*} LazyRef[A], f: A => B): {f, ref} LazyRef[B] =
  new LazyRef(() => f(ref.elem()))
 
-def mapc[A, B]: (ref: {*} LazyRef[A], f: {*} A => B) => {f, ref} LazyRef[B] =
+def mapc[A, B]: (ref: {*} LazyRef[A], f: A => B) => {f, ref} LazyRef[B] =
  (ref1, f1) => map[A, B](ref1, f1)
 
 def test(cap1: Cap, cap2: Cap) =

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

@@ -1,6 +1,5 @@
 
-class C
-type Cap = {*} C
+@annotation.capability class Cap
 
 object Generic: