Fix inferredTypeEdits for symbols

Co-authored-by: Jędrzej Rochala <[email protected]>
scala · Jan 30, 2025 · f6234e7 · f6234e7
1 parent 6d9b0f1
commit f6234e7
Show file tree

Hide file tree

Showing 2 changed files with 101 additions and 8 deletions.
diff --git a/presentation-compiler/src/main/dotty/tools/pc/InferredTypeProvider.scala b/presentation-compiler/src/main/dotty/tools/pc/InferredTypeProvider.scala
@@ -112,8 +112,8 @@ final class InferredTypeProvider(
     def imports: List[TextEdit] =
       printer.imports(autoImportsGen)
 
-    def printType(tpe: Type): String =
-      printer.tpe(tpe)
+    def printTypeAscription(tpe: Type, spaceBefore: Boolean = false): String =
+      (if spaceBefore then " : " else ": ") + printer.tpe(tpe)
 
     path.headOption match
       /* `val a = 1` or `var b = 2`
@@ -124,7 +124,7 @@ final class InferredTypeProvider(
        *     turns into
        * `.map((a: Int) => a + a)`
        */
-      case Some(vl @ ValDef(sym, tpt, rhs)) =>
+      case Some(vl @ ValDef(name, tpt, rhs)) =>
         val isParam = path match
           case head :: next :: _ if next.symbol.isAnonymousFunction => true
           case head :: (b @ Block(stats, expr)) :: next :: _
@@ -136,9 +136,11 @@ final class InferredTypeProvider(
           val endPos =
             findNamePos(sourceText, vl, keywordOffset).endPos.toLsp
           adjustOpt.foreach(adjust => endPos.setEnd(adjust.adjustedEndPos))
+          val spaceBefore = name.isOperatorName
+
           new TextEdit(
             endPos,
-            ": " + printType(optDealias(tpt.typeOpt)) + {
+            printTypeAscription(optDealias(tpt.typeOpt), spaceBefore) + {
               if withParens then ")" else ""
             }
           )
@@ -197,7 +199,7 @@ final class InferredTypeProvider(
        *     turns into
        * `def a[T](param : Int): Int = param`
        */
-      case Some(df @ DefDef(name, _, tpt, rhs)) =>
+      case Some(df @ DefDef(name, paramss, tpt, rhs)) =>
         def typeNameEdit =
           /* NOTE: In Scala 3.1.3, `List((1,2)).map((<<a>>,b) => ...)`
            * turns into `List((1,2)).map((:Inta,b) => ...)`,
@@ -208,10 +210,12 @@ final class InferredTypeProvider(
             if tpt.endPos.end > df.namePos.end then tpt.endPos.toLsp
             else df.namePos.endPos.toLsp
 
+          val spaceBefore = name.isOperatorName && paramss.isEmpty
+
           adjustOpt.foreach(adjust => end.setEnd(adjust.adjustedEndPos))
           new TextEdit(
             end,
-            ": " + printType(optDealias(tpt.typeOpt))
+            printTypeAscription(optDealias(tpt.typeOpt), spaceBefore)
           )
         end typeNameEdit
 
@@ -239,9 +243,10 @@ final class InferredTypeProvider(
        */
       case Some(bind @ Bind(name, body)) =>
         def baseEdit(withParens: Boolean) =
+          val spaceBefore = name.isOperatorName
           new TextEdit(
             bind.endPos.toLsp,
-            ": " + printType(optDealias(body.typeOpt)) + {
+            printTypeAscription(optDealias(body.typeOpt), spaceBefore) + {
               if withParens then ")" else ""
             }
           )
@@ -272,9 +277,10 @@ final class InferredTypeProvider(
        * `for(t: Int <- 0 to 10)`
        */
       case Some(i @ Ident(name)) =>
+        val spaceBefore = name.isOperatorName
         val typeNameEdit = new TextEdit(
           i.endPos.toLsp,
-          ": " + printType(optDealias(i.typeOpt.widen))
+          printTypeAscription(optDealias(i.typeOpt.widen), spaceBefore)
         )
         typeNameEdit :: imports
 

diff --git a/presentation-compiler/test/dotty/tools/pc/tests/edit/InsertInferredTypeSuite.scala b/presentation-compiler/test/dotty/tools/pc/tests/edit/InsertInferredTypeSuite.scala
@@ -883,6 +883,93 @@ class InsertInferredTypeSuite extends BaseCodeActionSuite:
         |""".stripMargin
     )
 
+  @Test def `operator-val` =
+    checkEdit(
+      """|object A {
+         |  val <<!>> = 1
+         |}
+         |""".stripMargin,
+      """|object A {
+         |  val ! : Int = 1
+         |}
+         |""".stripMargin
+    )
+
+  @Test def `operator-def` =
+    checkEdit(
+      """|object A {
+         |  def <<!>> = 1
+         |}
+         |""".stripMargin,
+      """|object A {
+         |  def ! : Int = 1
+         |}
+         |""".stripMargin
+    )
+
+  @Test def `operator-def-param` =
+    checkEdit(
+      """|object A {
+         |  def <<!>>[T] = 1
+         |}
+         |""".stripMargin,
+      """|object A {
+         |  def ![T]: Int = 1
+         |}
+         |""".stripMargin
+    )
+
+  @Test def `operator-def-type-param` =
+    checkEdit(
+      """|object A {
+         |  def <<!>>(x: Int) = 1
+         |}
+         |""".stripMargin,
+      """|object A {
+         |  def !(x: Int): Int = 1
+         |}
+         |""".stripMargin
+    )
+
+  @Test def `operator-for` =
+    checkEdit(
+      """|object A {
+         |  def foo = for(<<!>> <- List(1)) yield !
+         |}
+         |""".stripMargin,
+      """|object A {
+         |  def foo = for(! : Int <- List(1)) yield !
+         |}
+         |""".stripMargin
+    )
+  @Test def `operator-lambda` =
+    checkEdit(
+      """|object A {
+         |  val foo: Int => Int = (<<!>>) => ! + 1
+         |}
+         |""".stripMargin,
+      """|object A {
+         |  val foo: Int => Int = (! : Int) => ! + 1
+         |}
+         |""".stripMargin
+    )
+
+  @Test def `operator-ident` =
+    checkEdit(
+      """|object A {
+         |  def foo =
+         |    val ! = 1
+         |    <<!>>
+         |}
+         |""".stripMargin,
+      """|object A {
+         |  def foo =
+         |    val ! = 1
+         |    ! : Int
+         |}
+         |""".stripMargin
+    )
+
   def checkEdit(
       original: String,
       expected: String