surface (fix): Reduce the generated code size of Surface.methodOf in …

…Scala 3 (#3148)

Reduce the macro generated code size by binding Surface to local val and reuse it. This fixes #3147

airframe-surface/src/main/scala-3/wvlet/airframe/surface/CompileTimeSurfaceFactory.scala

@@ -81,7 +81,10 @@ private[surface] class CompileTimeSurfaceFactory[Q <: Quotes](using quotes: Q) {
   private val lazySurface = scala.collection.mutable.Set[TypeRepr]()
 
   private def surfaceOf(t: TypeRepr): Expr[Surface] = {
-    if (seen.contains(t)) {
+    if (surfaceToVar.contains(t)) {
+      // println(s"==== ${t} is already cached")
+      Ref(surfaceToVar(t)).asExprOf[Surface]
+    } else if (seen.contains(t)) {
       if (memo.contains(t)) {
         memo(t)
       } else {
@@ -679,15 +682,55 @@ private[surface] class CompileTimeSurfaceFactory[Q <: Quotes](using quotes: Q) {
     methodsOf(TypeRepr.of(using t))
   }
 
-  private val methodMemo = scala.collection.mutable.Map[TypeRepr, Expr[Seq[MethodSurface]]]()
-  private val methodSeen = scala.collection.mutable.Set[TypeRepr]()
+  // To reduce the byte code size, we need to memoize the generated surface bound to a variable
+  private val surfaceToVar = scala.collection.mutable.Map[TypeRepr, Symbol]()
 
-  private def methodsOf(targetType: TypeRepr): Expr[Seq[MethodSurface]] = {
-    if (methodMemo.contains(targetType)) methodMemo(targetType)
-    else if (seen.contains(targetType)) {
-      sys.error(s"recurcive type in method: ${targetType.typeSymbol.fullName}")
+  private def methodsOf(t: TypeRepr): Expr[Seq[MethodSurface]] = {
+    // Run just for collecting known surfaces. seen variable will be updated
+    methodsOfInternal(t)
+
+    var count = 0
+    // Bind the observed surfaces to local variables __s0, __s1, ...
+    seen.foreach { s =>
+      // Update the cache so that the next call of surfaceOf method will use the local varaible reference
+      surfaceToVar += s -> Symbol.newVal(
+        Symbol.spliceOwner,
+        s"__s${count}",
+        TypeRepr.of[Surface],
+        Flags.EmptyFlags,
+        Symbol.noSymbol
+      )
+      count += 1
+    }
+    val surfaceDefs: List[ValDef] = surfaceToVar.map { x =>
+      val sym = x._2
+      ValDef(sym, Some(memo(x._1).asTerm))
+    }.toList
+
+    // Clear method observation cache
+    seenMethodParent.clear()
+
+    /**
+      * Generate a code like this: {{ val __s0 = Surface.of[A] val __s1 = Surface.of[B] ...
+      *
+      * ClassMethodSurface( .... ) }}
+      */
+    val expr = Block(
+      surfaceDefs,
+      methodsOfInternal(t).asTerm
+    ).asExprOf[Seq[MethodSurface]]
+
+    // println(s"===  methodOf: ${t.typeSymbol.fullName} => \n${expr.show}")
+    expr
+  }
+
+  private val seenMethodParent = scala.collection.mutable.Set[TypeRepr]()
+
+  private def methodsOfInternal(targetType: TypeRepr): Expr[Seq[MethodSurface]] = {
+    if (seenMethodParent.contains(targetType)) {
+      sys.error(s"recursive method found in: ${targetType.typeSymbol.fullName}")
     } else {
-      methodSeen += targetType
+      seenMethodParent += targetType
       val localMethods = localMethodsOf(targetType).distinct.sortBy(_.name)
       val methodSurfaces = localMethods.map(m => (m, m.tree)).collect { case (m, df: DefDef) =>
         val mod   = Expr(modifierBitMaskOf(m))
@@ -704,8 +747,6 @@ private[surface] class CompileTimeSurfaceFactory[Q <: Quotes](using quotes: Q) {
         }
       }
       val expr = Expr.ofSeq(methodSurfaces)
-      // println(s"methodOf: ${targetType.typeSymbol.fullName} => \n${expr.show}")
-      methodMemo += targetType -> expr
       expr
     }
   }

build.sbt

@@ -444,7 +444,7 @@ lazy val surface =
     .settings(
       name        := "airframe-surface",
       description := "A library for extracting object structure surface",
-      // TODO: This is a temporaly solution. Use AirSpec after Scala 3 support of Surface is completed
+      // TODO: This is a temporary solution. Use AirSpec after Scala 3 support of Surface is completed
       libraryDependencies += "org.scalameta" %%% "munit" % "0.7.29" % Test,
       libraryDependencies ++= surfaceDependencies(scalaVersion.value)
     )