Check user defined PolyFunction refinements (#18457)

Fixes #18302

Author: nicolasstucki

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

@@ -1165,10 +1165,12 @@ class Definitions {
           Some(mt)
         case _ => None
 
-    private def isValidPolyFunctionInfo(info: Type)(using Context): Boolean =
+    def isValidPolyFunctionInfo(info: Type)(using Context): Boolean =
       def isValidMethodType(info: Type) = info match
         case info: MethodType =>
-          !info.resType.isInstanceOf[MethodOrPoly] // Has only one parameter list
+          !info.resType.isInstanceOf[MethodOrPoly] && // Has only one parameter list
+          !info.isVarArgsMethod &&
+          !info.isMethodWithByNameArgs // No by-name parameters
         case _ => false
       info match
         case info: PolyType => isValidMethodType(info.resType)

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

@@ -429,6 +429,12 @@ object Types {
       case _ => false
     }
 
+    /** Is this the type of a method that has a by-name parameters? */
+    def isMethodWithByNameArgs(using Context): Boolean = stripPoly match {
+      case mt: MethodType => mt.paramInfos.exists(_.isInstanceOf[ExprType])
+      case _ => false
+    }
+
     /** Is this the type of a method with a leading empty parameter list?
      */
     def isNullaryMethod(using Context): Boolean = stripPoly match {

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

@@ -381,13 +381,15 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
         case tree: ValDef =>
           registerIfHasMacroAnnotations(tree)
           checkErasedDef(tree)
+          Checking.checkPolyFunctionType(tree.tpt)
           val tree1 = cpy.ValDef(tree)(rhs = normalizeErasedRhs(tree.rhs, tree.symbol))
           if tree1.removeAttachment(desugar.UntupledParam).isDefined then
             checkStableSelection(tree.rhs)
           processValOrDefDef(super.transform(tree1))
         case tree: DefDef =>
           registerIfHasMacroAnnotations(tree)
           checkErasedDef(tree)
+          Checking.checkPolyFunctionType(tree.tpt)
           annotateContextResults(tree)
           val tree1 = cpy.DefDef(tree)(rhs = normalizeErasedRhs(tree.rhs, tree.symbol))
           processValOrDefDef(superAcc.wrapDefDef(tree1)(super.transform(tree1).asInstanceOf[DefDef]))
@@ -492,6 +494,9 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
           )
         case Block(_, Closure(_, _, tpt)) if ExpandSAMs.needsWrapperClass(tpt.tpe) =>
           superAcc.withInvalidCurrentClass(super.transform(tree))
+        case tree: RefinedTypeTree =>
+          Checking.checkPolyFunctionType(tree)
+          super.transform(tree)
         case _: Quote | _: QuotePattern =>
           ctx.compilationUnit.needsStaging = true
           super.transform(tree)

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

@@ -816,6 +816,31 @@ object Checking {
           else Feature.checkExperimentalFeature("features", imp.srcPos)
         case _ =>
   end checkExperimentalImports
+
+  /** Checks that PolyFunction only have valid refinements.
+   *
+   *  It only supports `apply` methods with one parameter list and optional type arguments.
+   */
+  def checkPolyFunctionType(tree: Tree)(using Context): Unit = new TreeTraverser {
+    def traverse(tree: Tree)(using Context): Unit = tree match
+      case tree: RefinedTypeTree if tree.tpe.derivesFrom(defn.PolyFunctionClass) =>
+        if tree.refinements.isEmpty then
+          reportNoRefinements(tree.srcPos)
+        tree.refinements.foreach {
+          case refinement: DefDef if refinement.name != nme.apply =>
+            report.error("PolyFunction only supports apply method refinements", refinement.srcPos)
+          case refinement: DefDef if !defn.PolyFunctionOf.isValidPolyFunctionInfo(refinement.tpe.widen) =>
+            report.error("Implementation restriction: PolyFunction apply must have exactly one parameter list and optionally type arguments. No by-name nor varags are allowed.", refinement.srcPos)
+          case _ =>
+        }
+      case _: RefTree if tree.symbol == defn.PolyFunctionClass =>
+        reportNoRefinements(tree.srcPos)
+      case _ =>
+        traverseChildren(tree)
+
+    def reportNoRefinements(pos: SrcPos) =
+      report.error("PolyFunction subtypes must refine the apply method", pos)
+  }.traverse(tree)
 }
 
 trait Checking {

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

@@ -528,7 +528,7 @@ object Nullables:
   def postProcessByNameArgs(fn: TermRef, app: Tree)(using Context): Tree =
     fn.widen match
       case mt: MethodType
-      if mt.paramInfos.exists(_.isInstanceOf[ExprType]) && !fn.symbol.is(Inline) =>
+      if mt.isMethodWithByNameArgs && !fn.symbol.is(Inline) =>
         app match
           case Apply(fn, args) =>
             object dropNotNull extends TreeMap:

tests/neg/i18302b.check

@@ -0,0 +1,4 @@
+-- Error: tests/neg/i18302b.scala:3:32 ---------------------------------------------------------------------------------
+3 |def polyFun: PolyFunction { def apply(x: Int)(y: Int): Int } = // error
+  |                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  |Implementation restriction: PolyFunction apply must have exactly one parameter list and optionally type arguments. No by-name nor varags are allowed.

tests/neg/i18302b.scala

@@ -0,0 +1,5 @@
+def test = polyFun(1)(2)
+
+def polyFun: PolyFunction { def apply(x: Int)(y: Int): Int } = // error
+  new PolyFunction:
+    def apply(x: Int)(y: Int): Int = x + y

tests/neg/i18302c.check

@@ -0,0 +1,4 @@
+-- Error: tests/neg/i18302c.scala:4:32 ---------------------------------------------------------------------------------
+4 |def polyFun: PolyFunction { def foo(x: Int): Int } = // error
+  |                            ^^^^^^^^^^^^^^^^^^^^
+  |                            PolyFunction only supports apply method refinements

tests/neg/i18302c.scala

@@ -0,0 +1,5 @@
+import scala.reflect.Selectable.reflectiveSelectable
+
+def test = polyFun.foo(1)
+def polyFun: PolyFunction { def foo(x: Int): Int } = // error
+  new PolyFunction { def foo(x: Int): Int = x + 1 }

tests/neg/i18302d.check

@@ -0,0 +1,4 @@
+-- Error: tests/neg/i18302d.scala:1:32 ---------------------------------------------------------------------------------
+1 |def polyFun: PolyFunction { def apply: Int } = // error
+  |                            ^^^^^^^^^^^^^^
+  |Implementation restriction: PolyFunction apply must have exactly one parameter list and optionally type arguments. No by-name nor varags are allowed.