Fix type parameters not getting properly typed

3 files changed, 107 insertions, 18 deletions

diff --git a/dottydoc/test/UsecaseTest.scala b/dottydoc/test/UsecaseTest.scala
index d5f338892..9be688297 100644
--- a/dottydoc/test/UsecaseTest.scala
+++ b/dottydoc/test/UsecaseTest.scala
@@ -52,6 +52,95 @@ class UsecaseTest extends DottyTest {
     }
   }
 
+  @Test def simpleUsecaseAddedArg = {
+    val source = new SourceFile(
+      "DefWithUseCase.scala",
+      """
+      |package scala
+      |
+      |trait Test[A] {
+      |  /** Definition with a "disturbing" signature
+      |   *
+      |   *  @usecase def foo(a: A): A
+      |   */
+      |  def foo[B]: A => B
+      |}
+      """.stripMargin
+    )
+
+    checkSources(source :: Nil) { packages =>
+      packages("scala") match {
+        case PackageImpl(_, _, List(trt: Trait), _, _) =>
+          val List(foo: Def) = trt.members
+
+          val returnValue = foo.returnValue match {
+            case ref: TypeReference => ref.title
+            case _ =>
+              assert(
+                false,
+                "Incorrect return value after usecase transformation"
+              )
+              ""
+          }
+
+          assert(
+            foo.typeParams.isEmpty,
+            "Type parameters were not stripped by usecase"
+          )
+          assert(returnValue == "A", "Incorrect return type after usecase")
+          assert(
+            foo.paramLists.head.list.head.title == "a",
+            "Incorrect parameter to function after usecase transformation"
+          )
+          assert(foo.name == "foo", s"Incorrect name after transform: ${foo.name}")
+      }
+    }
+  }
+
+  @Test def simpleTparamUsecase = {
+    val source = new SourceFile(
+      "DefWithUseCase.scala",
+      """
+      |package scala
+      |
+      |trait Test[A] {
+      |  /** Definition with a "disturbing" signature
+      |   *
+      |   *  @usecase def foo[C]: A
+      |   */
+      |  def foo[B]: A => B
+      |}
+      """.stripMargin
+    )
+
+    checkSources(source :: Nil) { packages =>
+      packages("scala") match {
+        case PackageImpl(_, _, List(trt: Trait), _, _) =>
+          val List(foo: Def) = trt.members
+
+          val returnValue = foo.returnValue match {
+            case ref: TypeReference => ref.title
+            case _ =>
+              assert(
+                false,
+                "Incorrect return value after usecase transformation"
+              )
+              ""
+          }
+
+          assert(
+            foo.typeParams.nonEmpty,
+            "Type parameters were incorrectly stripped by usecase"
+          )
+
+          assert(foo.typeParams.head == "C", "Incorrectly switched tparam")
+          assert(returnValue == "A", "Incorrect return type after usecase")
+
+          assert(foo.name == "foo", s"Incorrect name after transform: ${foo.name}")
+      }
+    }
+  }
+
   @Test def checkIterator = {
     val sources =
       "./scala-scala/src/library/scala/collection/Iterator.scala" :: Nil
diff --git a/src/dotty/tools/dotc/core/Comments.scala b/src/dotty/tools/dotc/core/Comments.scala
index 0a57d331a..1d9d53bac 100644
--- a/src/dotty/tools/dotc/core/Comments.scala
+++ b/src/dotty/tools/dotc/core/Comments.scala
@@ -60,6 +60,9 @@ object Comments {
     /** Entered by Namer */
     var symbol: Symbol = _
 
+    /** Set by typer */
+    var tpdCode: tpd.DefDef = _
+
     lazy val untpdCode: untpd.Tree = {
       val tree = new Parser(new SourceFile("<usecase>", code)).localDef(codePos.start, EmptyFlags)
 
@@ -72,19 +75,5 @@ object Comments {
           tree
       }
     }
-
-    /** Set by typer calling `typeTree` */
-    var tpdCode: tpd.DefDef = _
-
-    def typeTree()(implicit ctx: Context): Unit = untpdCode match {
-      case df: untpd.DefDef =>
-        ctx.typer.typedDefDef(df, symbol) match {
-          case tree: tpd.DefDef => tpdCode = tree
-          case _ =>
-            ctx.error("proper def could not be typed from `@usecase`", codePos)
-        }
-      case _ =>
-        ctx.error("proper def was not found in `@usecase`", codePos)
-    }
   }
 }
diff --git a/src/dotty/tools/dotc/typer/Typer.scala b/src/dotty/tools/dotc/typer/Typer.scala
index 56571116a..d633e5653 100644
--- a/src/dotty/tools/dotc/typer/Typer.scala
+++ b/src/dotty/tools/dotc/typer/Typer.scala
@@ -1182,9 +1182,6 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
       tpt1 = tpt1.withType(avoid(tpt1.tpe, vparamss1.flatMap(_.map(_.symbol))))
     }
 
-    /** Type usecases */
-    ctx.docbase.docstring(sym).map(_.usecases.map(_.typeTree()))
-
     assignType(cpy.DefDef(ddef)(name, tparams1, vparamss1, tpt1, rhs1), sym)
     //todo: make sure dependent method types do not depend on implicits or by-name params
   }
@@ -1523,7 +1520,9 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
         buf += typed(stat)(ctx.exprContext(stat, exprOwner))
         traverse(rest)
       case nil =>
-        buf.toList
+        val tpdStats = buf.toList
+        typedUsecases(tpdStats.map(_.symbol), exprOwner)
+        tpdStats
     }
     traverse(stats)
   }
@@ -1536,6 +1535,18 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
     tpd.cpy.DefDef(mdef)(rhs = Inliner.bodyToInline(mdef.symbol)) ::
         Inliner.removeInlineAccessors(mdef.symbol)
 
+  def typedUsecases(syms: List[Symbol], owner: Symbol)(implicit ctx: Context): Unit =
+    for {
+      sym <- syms
+      usecase <- ctx.docbase.docstring(sym).map(_.usecases).getOrElse(Nil)
+      List(tpdTree) = typedStats(usecase.untpdCode :: Nil, owner)
+    } yield {
+      if (tpdTree.isInstanceOf[tpd.DefDef])
+        usecase.tpdCode = tpdTree.asInstanceOf[tpd.DefDef]
+      else
+        ctx.error("Couldn't compile `@usecase`", usecase.codePos)
+    }
+
   def typedExpr(tree: untpd.Tree, pt: Type = WildcardType)(implicit ctx: Context): Tree =
     typed(tree, pt)(ctx retractMode Mode.PatternOrType)
   def typedType(tree: untpd.Tree, pt: Type = WildcardType)(implicit ctx: Context): Tree = // todo: retract mode between Type and Pattern?