Check user defined PolyFunction refinements (#18457)

Fixes #18302

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.

tests/neg/i18302d.scala

@@ -0,0 +1,2 @@
+def polyFun: PolyFunction { def apply: Int } = // error
+  new PolyFunction { def apply: Int = 1 }

tests/neg/i18302e.check

@@ -0,0 +1,8 @@
+-- Error: tests/neg/i18302e.scala:1:13 ---------------------------------------------------------------------------------
+1 |def polyFun: PolyFunction {  } = // error
+  |             ^^^^^^^^^^^^^^^^^
+  |             PolyFunction subtypes must refine the apply method
+-- Error: tests/neg/i18302e.scala:4:15 ---------------------------------------------------------------------------------
+4 |def polyFun(f: PolyFunction {  }) = () // error
+  |               ^^^^^^^^^^^^^^^^^
+  |               PolyFunction subtypes must refine the apply method

tests/neg/i18302e.scala

@@ -0,0 +1,4 @@
+def polyFun: PolyFunction {  } = // error
+  new PolyFunction {  }
+
+def polyFun(f: PolyFunction {  }) = () // error

tests/neg/i18302f.check

@@ -0,0 +1,12 @@
+-- Error: tests/neg/i18302f.scala:1:13 ---------------------------------------------------------------------------------
+1 |def polyFun: PolyFunction = // error
+  |             ^^^^^^^^^^^^
+  |             PolyFunction subtypes must refine the apply method
+-- Error: tests/neg/i18302f.scala:4:16 ---------------------------------------------------------------------------------
+4 |def polyFun2(a: PolyFunction) = () // error
+  |                ^^^^^^^^^^^^
+  |                PolyFunction subtypes must refine the apply method
+-- Error: tests/neg/i18302f.scala:6:14 ---------------------------------------------------------------------------------
+6 |val polyFun3: PolyFunction = // error
+  |              ^^^^^^^^^^^^
+  |              PolyFunction subtypes must refine the apply method

tests/neg/i18302f.scala

@@ -0,0 +1,7 @@
+def polyFun: PolyFunction = // error
+  new PolyFunction {  }
+
+def polyFun2(a: PolyFunction) = () // error
+
+val polyFun3: PolyFunction = // error
+  new PolyFunction {  }

tests/neg/i18302i.scala

@@ -0,0 +1,6 @@
+def polyFun1: Option[PolyFunction] = ??? // error
+def polyFun2: PolyFunction & Any = ??? // error
+def polyFun3: Any & PolyFunction = ??? // error
+def polyFun4: PolyFunction | Any = ??? // error
+def polyFun5: Any | PolyFunction = ??? // error
+def polyFun6(a: Any | PolyFunction) = ??? // error

tests/neg/i18302j.scala

@@ -0,0 +1,5 @@
+def polyFunByName: PolyFunction { def apply(thunk: => Int): Int } = // error
+  new PolyFunction { def apply(thunk: => Int): Int = 1 }
+
+def polyFunVarArgs: PolyFunction { def apply(args: Int*): Int } = // error
+  new PolyFunction { def apply(thunk: Int*): Int = 1 }

tests/neg/i8299.scala

@@ -0,0 +1,8 @@
+package example
+
+object Main {
+  def main(a: Array[String]): Unit = {
+    val p: PolyFunction = // error: PolyFunction subtypes must refine the apply method
+      [A] => (xs: List[A]) => xs.headOption
+  }
+}

tests/pos/i18302a.scala

@@ -0,0 +1,4 @@
+def test = polyFun(1)
+
+def polyFun: PolyFunction { def apply(x: Int): Int } =
+  new PolyFunction { def apply(x: Int): Int = x + 1 }