From 90d99c36e8e9260b9f81dbe3e4eec1289e17160a Mon Sep 17 00:00:00 2001
From: Matt Bovel <matthieu@bovel.net>
Date: Fri, 19 Apr 2024 17:17:08 +0200
Subject: [PATCH] Add AnnotationsMappingBenchmark

---
 .../AnnotationsMappingBenchmark.scala         | 71 +++++++++++++++++++
 bench-micro/tests/someAnnotatedTypes.scala    | 28 ++++++++
 2 files changed, 99 insertions(+)
 create mode 100644 bench-micro/src/main/scala/dotty/tools/benchmarks/AnnotationsMappingBenchmark.scala
 create mode 100644 bench-micro/tests/someAnnotatedTypes.scala

diff --git a/bench-micro/src/main/scala/dotty/tools/benchmarks/AnnotationsMappingBenchmark.scala b/bench-micro/src/main/scala/dotty/tools/benchmarks/AnnotationsMappingBenchmark.scala
new file mode 100644
index 000000000000..310a1745171f
--- /dev/null
+++ b/bench-micro/src/main/scala/dotty/tools/benchmarks/AnnotationsMappingBenchmark.scala
@@ -0,0 +1,71 @@
+package dotty.tools.benchmarks
+
+import org.openjdk.jmh.annotations.{Benchmark, BenchmarkMode, Fork, Level, Measurement, Mode as JMHMode, Param, Scope, Setup, State, Warmup}
+import java.util.concurrent.TimeUnit.SECONDS
+
+import dotty.tools.dotc.{Driver, Run, Compiler}
+import dotty.tools.dotc.ast.{tpd, TreeTypeMap}, tpd.{Apply, Block, Tree, TreeAccumulator, TypeApply}
+import dotty.tools.dotc.core.Annotations.{Annotation, ConcreteAnnotation, EmptyAnnotation}
+import dotty.tools.dotc.core.Contexts.{ContextBase, Context, ctx, withMode}
+import dotty.tools.dotc.core.Mode
+import dotty.tools.dotc.core.Phases.Phase
+import dotty.tools.dotc.core.Symbols.{defn, mapSymbols, Symbol}
+import dotty.tools.dotc.core.Types.{AnnotatedType, NoType, SkolemType, TermRef, Type, TypeMap}
+import dotty.tools.dotc.parsing.Parser
+import dotty.tools.dotc.typer.TyperPhase
+
+/** Measures the performance of mapping over annotated types.
+  *
+  * Run with: scala3-bench-micro / Jmh / run AnnotationsMappingBenchmark
+  */
+@Fork(value = 4)
+@Warmup(iterations = 4, time = 1, timeUnit = SECONDS)
+@Measurement(iterations = 4, time = 1, timeUnit = SECONDS)
+@BenchmarkMode(Array(JMHMode.Throughput))
+@State(Scope.Thread)
+class AnnotationsMappingBenchmark:
+  var tp: Type = null
+  var specialIntTp: Type = null
+  var context: Context = null
+  var typeFunction: Context ?=> Type => Type = null
+  var typeMap: TypeMap = null
+
+  @Param(Array("v1", "v2", "v3", "v4"))
+  var valName: String = null
+
+  @Param(Array("id", "mapInts"))
+  var typeFunctionName: String = null
+
+  @Setup(Level.Iteration)
+  def setup(): Unit =
+    val testPhase =
+      new Phase:
+        final override def phaseName = "testPhase"
+        final override def run(using ctx: Context): Unit =
+          val pkg = ctx.compilationUnit.tpdTree.symbol
+          tp = pkg.requiredClass("Test").requiredValueRef(valName).underlying
+          specialIntTp = pkg.requiredClass("Test").requiredType("SpecialInt").typeRef
+          context = ctx
+
+    val compiler =
+      new Compiler:
+        private final val baseCompiler = new Compiler()
+        final override def phases = List(List(Parser()), List(TyperPhase()), List(testPhase))
+
+    val driver =
+      new Driver:
+        final override def newCompiler(using Context): Compiler = compiler
+
+    driver.process(Array("-classpath", System.getProperty("BENCH_CLASS_PATH"), "tests/someAnnotatedTypes.scala"))
+
+    typeFunction =
+      typeFunctionName match
+        case "id"      => tp => tp
+        case "mapInts" => tp => (if tp frozen_=:= defn.IntType then specialIntTp else tp)
+        case _         => throw new IllegalArgumentException(s"Unknown type function: $typeFunctionName")
+
+    typeMap =
+      new TypeMap(using context):
+        final override def apply(tp: Type): Type = typeFunction(mapOver(tp))
+
+  @Benchmark def applyTypeMap() = typeMap.apply(tp)
diff --git a/bench-micro/tests/someAnnotatedTypes.scala b/bench-micro/tests/someAnnotatedTypes.scala
new file mode 100644
index 000000000000..8b12d4f7c2c6
--- /dev/null
+++ b/bench-micro/tests/someAnnotatedTypes.scala
@@ -0,0 +1,28 @@
+class Test:
+  class FlagAnnot extends annotation.StaticAnnotation
+  class StringAnnot(val s: String) extends annotation.StaticAnnotation
+  class LambdaAnnot(val f: Int => Boolean) extends annotation.StaticAnnotation
+
+  type SpecialInt <: Int
+
+  val v1: Int @FlagAnnot = 42
+
+  val v2: Int @StringAnnot("hello") = 42
+
+  val v3: Int @LambdaAnnot(it => it == 42) = 42
+
+  val v4: Int @LambdaAnnot(it => {
+    def g(x: Int, y: Int) = x - y + 5
+    g(it, 7) * 2 == 80
+  }) = 42
+
+  /*val v5: Int @LambdaAnnot(it => {
+    class Foo(x: Int):
+      def xPlus10 = x + 10
+      def xPlus20 = x + 20
+      def xPlus(y: Int) = x + y
+    val foo = Foo(it)
+    foo.xPlus10 - foo.xPlus20 + foo.xPlus(30) == 62
+  }) = 42*/
+
+  def main(args: Array[String]): Unit = ???