Add defn.RefinedFunctionOf extractor

Author: nicolasstucki

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

@@ -876,9 +876,8 @@ object CaptureSet:
  empty
  case CapturingType(parent, refs) =>
  recur(parent) ++ refs
- case tpd @ RefinedType(parent, _, rinfo: MethodType)
- if followResult && defn.isFunctionNType(tpd) =>
- ofType(parent, followResult = false) // pick up capture set from parent type
+ case tpd @ defn.RefinedFunctionOf(rinfo: MethodType) if followResult =>
+ ofType(tpd.parent, followResult = false) // pick up capture set from parent type
  ++ (recur(rinfo.resType) // add capture set of result
  -- CaptureSet(rinfo.paramRefs.filter(_.isTracked)*)) // but disregard bound parameters
  case tpd @ AppliedType(tycon, args) =>

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

@@ -195,7 +195,7 @@ class CheckCaptures extends Recheck, SymTransformer:
  capt.println(i"solving $t")
  refs.solve()
  traverse(parent)
- case t @ RefinedType(_, nme.apply, rinfo) if defn.isFunctionType(t) =>
+ case t @ defn.RefinedFunctionOf(rinfo) =>
  traverse(rinfo)
  case tp: TypeVar =>
  case tp: TypeRef =>
@@ -302,8 +302,8 @@ class CheckCaptures extends Recheck, SymTransformer:
  t
  case _ =>
  val t1 = t match
- case t @ RefinedType(parent, rname, rinfo: MethodType) if defn.isFunctionType(t) =>
- t.derivedRefinedType(parent, rname, this(rinfo))
+ case t @ defn.RefinedFunctionOf(rinfo: MethodType) =>
+ t.derivedRefinedType(t.parent, t.refinedName, this(rinfo))
  case _ =>
  mapOver(t)
  if variance > 0 then t1
@@ -845,7 +845,7 @@ class CheckCaptures extends Recheck, SymTransformer:
  case actual @ AppliedType(tycon, args) if defn.isNonRefinedFunction(actual) =>
  adaptFun(actual, args.init, args.last, expected, covariant, insertBox,
  (aargs1, ares1) => actual.derivedAppliedType(tycon, aargs1 :+ ares1))
- case actual @ RefinedType(_, _, rinfo: MethodType) if defn.isFunctionType(actual) =>
+ case actual @ defn.RefinedFunctionOf(rinfo: MethodType) =>
  // TODO Find a way to combine handling of generic and dependent function types (here and elsewhere)
  adaptFun(actual, rinfo.paramInfos, rinfo.resType, expected, covariant, insertBox,
  (aargs1, ares1) =>
@@ -855,11 +855,11 @@ class CheckCaptures extends Recheck, SymTransformer:
  adaptFun(actual, actual.paramInfos, actual.resType, expected, covariant, insertBox,
  (aargs1, ares1) =>
  actual.derivedLambdaType(paramInfos = aargs1, resType = ares1))
- case actual @ RefinedType(p, nme, rinfo: PolyType) if defn.isFunctionType(actual) =>
+ case actual @ defn.RefinedFunctionOf(rinfo: PolyType) =>
  adaptTypeFun(actual, rinfo.resType, expected, covariant, insertBox,
  ares1 =>
  val rinfo1 = rinfo.derivedLambdaType(rinfo.paramNames, rinfo.paramInfos, ares1)
- val actual1 = actual.derivedRefinedType(p, nme, rinfo1)
+ val actual1 = actual.derivedRefinedType(actual.parent, actual.refinedName, rinfo1)
  actual1
  )
  case _ =>
@@ -1080,7 +1080,7 @@ class CheckCaptures extends Recheck, SymTransformer:
  case CapturingType(parent, refs) =>
  healCaptureSet(refs)
  traverse(parent)
- case tp @ RefinedType(parent, rname, rinfo: MethodType) if defn.isFunctionType(tp) =>
+ case defn.RefinedFunctionOf(rinfo: MethodType) =>
  traverse(rinfo)
  case tp: TermLambda =>
  val saved = allowed

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

@@ -54,7 +54,7 @@ extends tpd.TreeTraverser:
  val boxedRes = recur(res)
  if boxedRes eq res then tp
  else tp1.derivedAppliedType(tycon, args.init :+ boxedRes)
- case tp1 @ RefinedType(_, _, rinfo: MethodType) if defn.isFunctionType(tp1) =>
+ case tp1 @ defn.RefinedFunctionOf(rinfo: MethodType) =>
  val boxedRinfo = recur(rinfo)
  if boxedRinfo eq rinfo then tp
  else boxedRinfo.toFunctionType(isJava = false, alwaysDependent = true)
@@ -149,7 +149,7 @@ extends tpd.TreeTraverser:
  tp.derivedAppliedType(tycon1, args1 :+ res1)
  else
  tp.derivedAppliedType(tycon1, args.mapConserve(arg => this(arg)))
- case tp @ RefinedType(core, rname, rinfo: MethodType) if defn.isFunctionType(tp) =>
+ case defn.RefinedFunctionOf(rinfo: MethodType) =>
  val rinfo1 = apply(rinfo)
  if rinfo1 ne rinfo then rinfo1.toFunctionType(isJava = false, alwaysDependent = true)
  else tp

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

@@ -1129,6 +1129,20 @@ class Definitions {
  }
  }
 
+ object RefinedFunctionOf {
+ /** Matches a refined `PolyFunction`/`FunctionN[...]`/`ContextFunctionN[...]`.
+ * Extracts the method type type and apply info.
+ */
+ def unapply(tpe: RefinedType)(using Context): Option[MethodOrPoly] = {
+ tpe.refinedInfo match
+ case mt: MethodOrPoly
+ if tpe.refinedName == nme.apply
+ && (tpe.parent.derivesFrom(defn.PolyFunctionClass) || isFunctionNType(tpe.parent)) =>
+ Some(mt)
+ case _ => None
+ }
+ }
+
  object PolyFunctionOf {
 
  /** Creates a refined `PolyFunction` with an `apply` method with the given info. */

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

@@ -4097,8 +4097,8 @@ object Types {
  tp.derivedAppliedType(tycon, addInto(args.head) :: Nil)
  case tp @ AppliedType(tycon, args) if defn.isFunctionNType(tp) =>
  wrapConvertible(tp.derivedAppliedType(tycon, args.init :+ addInto(args.last)))
- case tp @ RefinedType(parent, rname, rinfo) if defn.isFunctionType(tp) =>
- wrapConvertible(tp.derivedRefinedType(parent, rname, addInto(rinfo)))
+ case tp @ defn.RefinedFunctionOf(rinfo) =>
+ wrapConvertible(tp.derivedRefinedType(tp.parent, tp.refinedName, addInto(rinfo)))
  case tp: MethodOrPoly =>
  tp.derivedLambdaType(resType = addInto(tp.resType))
  case ExprType(resType) =>