From 939d06f4e52ae881f412eda2563a57c123161efa Mon Sep 17 00:00:00 2001
From: Wojciech Mazur <wmazur@virtuslab.com>
Date: Wed, 19 Jun 2024 12:51:37 +0200
Subject: [PATCH] Improve `defn.PolyFunctionOf` extractor

* Only match `RefinedType` representing the `PolyFunction`. This will
  allow us to use `derivedRefinedType` on the function type.
* Only match the refinement if it is a `MethodOrPoly`. `ExprType` is not
  a valid `PolyFunction` refinement.
* Remove `dealias` in `PolyFunctionOf` extractor. There was only one
  case where this was necessary and it added unnecessary overhead.

[Cherry-picked 60b2d0276f4a7b36b4b00152724aa8ab5e5f0e5d][modified]
---
 compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala |  2 +-
 compiler/src/dotty/tools/dotc/core/Definitions.scala | 11 ++++++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala b/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
index a2dd3e450a50..e6287bb3db5d 100644
--- a/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
+++ b/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
@@ -732,7 +732,7 @@ class CheckCaptures extends Recheck, SymTransformer:
 
         try
           val eres = expected.dealias.stripCapturing match
-            case RefinedType(_, _, rinfo: PolyType) => rinfo.resType
+            case defn.PolyFunctionOf(rinfo: PolyType) => rinfo.resType
             case expected: PolyType => expected.resType
             case _ => WildcardType
 
diff --git a/compiler/src/dotty/tools/dotc/core/Definitions.scala b/compiler/src/dotty/tools/dotc/core/Definitions.scala
index 3813c87943a2..262317a3c120 100644
--- a/compiler/src/dotty/tools/dotc/core/Definitions.scala
+++ b/compiler/src/dotty/tools/dotc/core/Definitions.scala
@@ -1151,11 +1151,12 @@ class Definitions {
      *
      *  Pattern: `PolyFunction { def apply: $pt }`
      */
-    def unapply(ft: Type)(using Context): Option[PolyType] = ft.dealias match
-      case RefinedType(parent, nme.apply, pt: PolyType)
-      if parent.derivesFrom(defn.PolyFunctionClass) =>
-        Some(pt)
-      case _ => None
+    def unapply(tpe: RefinedType)(using Context): Option[MethodOrPoly] =
+      tpe.refinedInfo match
+        case mt: MethodOrPoly
+        if tpe.refinedName == nme.apply && tpe.parent.derivesFrom(defn.PolyFunctionClass) =>
+          Some(mt)
+        case _ => None
   }
 
   object ErasedFunctionOf {