Reimplement code generation using KotlinPoet (open-toast#143)

Co-authored-by: pranjalvachaspati-toast <pranjal.vachaspati@toasttab.com>
andrewparmet · Oct 22, 2021 · 67db17d · 67db17d
1 parent 4f6afd8
commit 67db17d
Show file tree

Hide file tree

Showing 91 changed files with 2,687 additions and 4,224 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -68,6 +68,8 @@ subprojects {
                 allWarningsAsErrors = true
                 jvmTarget = "1.8"
                 freeCompilerArgs = listOf("-Xinline-classes")
+                languageVersion = "1.4"
+                apiVersion = "1.4"
             }
         }
 

diff --git a/buildSrc/build.gradle.kts b/buildSrc/build.gradle.kts
@@ -23,11 +23,11 @@ repositories {
 }
 
 dependencies {
-    implementation("com.diffplug.spotless:spotless-plugin-gradle:5.15.0")
+    implementation("com.diffplug.spotless:spotless-plugin-gradle:5.15.1")
     implementation("com.google.protobuf:protobuf-gradle-plugin:0.8.17")
     implementation("io.codearte.gradle.nexus:gradle-nexus-staging-plugin:0.30.0")
     implementation("com.google.guava:guava:30.1.1-jre")
     implementation("ru.vyarus:gradle-animalsniffer-plugin:1.5.0")
-    implementation("org.jetbrains.kotlinx:binary-compatibility-validator:0.5.0")
+    implementation("org.jetbrains.kotlinx:binary-compatibility-validator:0.7.0")
     implementation(kotlin("gradle-plugin-api"))
 }
diff --git a/buildSrc/src/main/kotlin/Dependencies.kt b/buildSrc/src/main/kotlin/Dependencies.kt
@@ -21,22 +21,22 @@ object versions {
     const val grpc = "1.41.0"
     const val grpcKotlin = "1.1.0"
     const val kollection = "0.7"
-    const val kotlin = "1.4.32"
+    const val kotlin = "1.5.21"
+    const val kotlinPoet = "1.10.1"
     const val kotlinxCoroutines = "1.3.9"
     const val protobuf = DEFAULT_PROTOBUF_VERSION
     const val protobufPlugin = "0.8.17"
-    const val stringTemplate = "4.3.1"
 
     // Test
-    const val jackson = "2.12.3"
+    const val jackson = "2.13.0"
     const val junit = "5.7.1"
     const val truth = "1.1.3"
 
     // Benchmarks
     const val datasets = "0.1.0"
     const val gradleDownload = "4.1.1"
     const val jmh = "1.26"
-    const val wire = "3.5.0"
+    const val wire = "4.0.0-alpha.12"
 
     // Third Party
     const val protoGoogleCommonProtos = "2.5.0"
@@ -57,6 +57,7 @@ object libraries {
 
     const val kollection = "com.github.andrewoma.dexx:kollection:${versions.kollection}"
 
+    const val kotlinPoet = "com.squareup:kotlinpoet:${versions.kotlinPoet}"
     const val kotlinReflect = "org.jetbrains.kotlin:kotlin-reflect:${versions.kotlin}"
     const val kotlinStdlib = "org.jetbrains.kotlin:kotlin-stdlib:${versions.kotlin}"
     const val kotlinxCoroutinesCore = "org.jetbrains.kotlinx:kotlinx-coroutines-core:${versions.kotlinxCoroutines}"
@@ -66,8 +67,6 @@ object libraries {
     const val protobufLite = "com.google.protobuf:protobuf-javalite:${versions.protobuf}"
     const val protoc = "com.google.protobuf:protoc:${versions.protobuf}"
 
-    const val stringTemplate = "org.antlr:ST4:${versions.stringTemplate}"
-
     // Test
     const val jackson = "com.fasterxml.jackson.module:jackson-module-kotlin:${versions.jackson}"
     const val junit = "org.junit.jupiter:junit-jupiter:${versions.junit}"

diff --git a/examples/protos/src/main/kotlin/io/grpc/examples/routeguide/Database.kt b/examples/protos/src/main/kotlin/io/grpc/examples/routeguide/Database.kt
@@ -22,7 +22,9 @@ import com.fasterxml.jackson.module.kotlin.readValue
 object Database {
     fun features(): List<Feature> {
         return javaClass.getResourceAsStream("route_guide_db.json").use {
-            ObjectMapper().registerModule(KotlinModule()).readValue<FeatureDatabase>(it!!.reader())
+            ObjectMapper()
+                .registerModule(KotlinModule.Builder().build())
+                .readValue<FeatureDatabase>(it!!.reader())
         }.feature
     }
 }
diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties
@@ -1,5 +1,5 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-7.1.1-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-7.2-bin.zip
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
diff --git a/protokt-codegen/build.gradle.kts b/protokt-codegen/build.gradle.kts
@@ -42,10 +42,10 @@ dependencies {
     implementation(libraries.arrow)
     implementation(libraries.grpcStub)
     implementation(libraries.kollection)
+    implementation(libraries.kotlinPoet)
     implementation(libraries.kotlinReflect)
     implementation(libraries.kotlinxCoroutinesCore)
     implementation(libraries.protobufJava)
-    implementation(libraries.stringTemplate)
 
     testImplementation(project(":testing:testing-util"))
 

diff --git a/protokt-codegen/src/main/kotlin/com/toasttab/protokt/Main.kt b/protokt-codegen/src/main/kotlin/com/toasttab/protokt/Main.kt
@@ -20,16 +20,9 @@ import com.google.protobuf.ExtensionRegistry
 import com.google.protobuf.compiler.PluginProtos.CodeGeneratorRequest
 import com.google.protobuf.compiler.PluginProtos.CodeGeneratorResponse
 import com.google.protobuf.compiler.PluginProtos.CodeGeneratorResponse.Feature
-import com.toasttab.protokt.codegen.impl.Accumulator
-import com.toasttab.protokt.codegen.impl.Annotator
-import com.toasttab.protokt.codegen.impl.FileDescriptorResolver
-import com.toasttab.protokt.codegen.impl.ImportResolver
-import com.toasttab.protokt.codegen.impl.resolvePackage
-import com.toasttab.protokt.codegen.protoc.AnnotatedType
+import com.squareup.kotlinpoet.FileSpec
+import com.toasttab.protokt.codegen.impl.FileBuilder
 import com.toasttab.protokt.codegen.protoc.ProtocolContext
-import com.toasttab.protokt.codegen.protoc.TypeDesc
-import com.toasttab.protokt.codegen.protoc.fileName
-import com.toasttab.protokt.codegen.protoc.respectJavaPackage
 import com.toasttab.protokt.codegen.protoc.toProtocol
 import com.toasttab.protokt.ext.Protokt
 import java.io.OutputStream
@@ -51,11 +44,8 @@ internal fun main(bytes: ByteArray, out: OutputStream) {
     val files = req.protoFileList
         .filter { filesToGenerate.contains(it.name) }
         .mapNotNull {
-            response(
-                it,
-                generate(it, req.protoFileList, filesToGenerate, params),
-                params
-            )
+            val fileSpec = generate(it, req.protoFileList, filesToGenerate, params)
+            fileSpec?.let(::response)
         }
 
     if (files.isNotEmpty()) {
@@ -72,9 +62,8 @@ private fun generate(
     protoFileList: List<FileDescriptorProto>,
     filesToGenerate: Set<String>,
     params: Map<String, String>
-): String {
-    val code = StringBuilder()
-    val protocol =
+): FileSpec? =
+    FileBuilder.buildFile(
         toProtocol(
             ProtocolContext(
                 fdp,
@@ -83,35 +72,31 @@ private fun generate(
                 protoFileList
             )
         )
-
-    val descs = protocol.types.map { TypeDesc(protocol.desc, AnnotatedType(it)) }
-
-    Accumulator.apply(
-        descs.map(Annotator::apply),
-        ImportResolver.resolveImports(descs),
-        FileDescriptorResolver.resolveFileDescriptor(descs),
-        code::append
     )
 
-    return code.toString()
-}
+private fun response(fileSpec: FileSpec) =
+    CodeGeneratorResponse.File
+        .newBuilder()
+        .setContent(fileSpec.toString().let(::tidy))
+        .setName(fileSpec.name)
+        .build()
 
-private fun response(
-    fdp: FileDescriptorProto,
-    code: String,
-    params: Map<String, String>
-) =
-    code.takeIf { it.isNotBlank() }?.let {
-        CodeGeneratorResponse.File
-            .newBuilder()
-            .setContent(code)
-            .setName(
-                fileName(
-                    resolvePackage(fdp, respectJavaPackage(params)),
-                    fdp.name
-                )
-            ).build()
-    }
+// strips Explicit API mode declarations
+// https://kotlinlang.org/docs/whatsnew14.html#explicit-api-mode-for-library-authors
+private fun tidy(code: String) =
+    code
+        // https://stackoverflow.com/a/64970734
+        .replace("public class ", "class ")
+        .replace("public val ", "val ")
+        .replace("public var ", "var ")
+        .replace("public fun ", "fun ")
+        .replace("public object ", "object ")
+        .replace("public companion ", "companion ")
+        .replace("public override ", "override ")
+        .replace("public sealed ", "sealed ")
+        .replace("public data ", "data ")
+        // https://github.com/square/kotlinpoet/pull/932
+        .replace("): Unit {", ") {")
 
 private fun parseParams(req: CodeGeneratorRequest) =
     if (req.parameter == null || req.parameter.isEmpty()) {

diff --git a/...oasttab/protokt/codegen/impl/Annotator.kt → ...b/protokt/codegen/annotators/Annotator.kt b/...oasttab/protokt/codegen/impl/Annotator.kt → ...b/protokt/codegen/annotators/Annotator.kt
@@ -13,27 +13,25 @@
  * limitations under the License.
  */
 
-package com.toasttab.protokt.codegen.impl
+package com.toasttab.protokt.codegen.annotators
 
-import arrow.core.Some
 import com.github.andrewoma.dexx.kollection.immutableListOf
-import com.toasttab.protokt.codegen.impl.EnumAnnotator.Companion.annotateEnum
-import com.toasttab.protokt.codegen.impl.MessageAnnotator.Companion.annotateMessage
-import com.toasttab.protokt.codegen.impl.ServiceAnnotator.annotateService
-import com.toasttab.protokt.codegen.model.PPackage
+import com.squareup.kotlinpoet.TypeSpec
+import com.toasttab.protokt.codegen.annotators.MessageAnnotator.Companion.annotateMessage
+import com.toasttab.protokt.codegen.annotators.ServiceAnnotator.annotateService
+import com.toasttab.protokt.codegen.impl.EnumBuilder
 import com.toasttab.protokt.codegen.protoc.AnnotatedType
 import com.toasttab.protokt.codegen.protoc.Enum
 import com.toasttab.protokt.codegen.protoc.FileDesc
 import com.toasttab.protokt.codegen.protoc.Message
+import com.toasttab.protokt.codegen.protoc.Protocol
 import com.toasttab.protokt.codegen.protoc.ProtocolContext
 import com.toasttab.protokt.codegen.protoc.Service
 import com.toasttab.protokt.codegen.protoc.TopLevelType
 import com.toasttab.protokt.codegen.protoc.TypeDesc
 
 /**
- * STAnnotator is an implementation of a side effect free function.
- * The input is an AST<TypeDesc> and its output is a NEW AST<TypeDesc> that has
- * a NEW fully constructed AnnotatedType attached to each AST node.
+ * Annotates an unannotated AST. This effectively converts the protobuf AST to a Kotlin AST.
  */
 object Annotator {
     const val rootGoogleProto = "google.protobuf"
@@ -45,43 +43,42 @@ object Annotator {
 
     data class Context(
         val enclosing: List<Message>,
-        val pkg: PPackage,
         val desc: FileDesc
     )
 
-    fun apply(data: TypeDesc) =
-        TypeDesc(
-            data.desc,
-            AnnotatedType(
-                data.type.rawType,
-                Some(
-                    annotate(
-                        data.type.rawType,
-                        Context(
-                            immutableListOf(),
-                            kotlinPackage(data),
-                            data.desc
-                        )
+    fun apply(protocol: Protocol) =
+        protocol.types.flatMap {
+            val annotated =
+                annotate(
+                    it,
+                    Context(
+                        immutableListOf(),
+                        protocol.desc
                     )
                 )
-            )
-        )
 
-    fun annotate(type: TopLevelType, ctx: Context): String =
+            annotated.map { type ->
+                TypeDesc(protocol.desc, AnnotatedType(it, type))
+            }
+        }
+
+    fun annotate(type: TopLevelType, ctx: Context): Iterable<TypeSpec> =
         when (type) {
             is Message ->
                 nonGrpc(ctx) {
                     nonDescriptors(ctx) {
-                        annotateMessage(
-                            type,
-                            ctx.copy(enclosing = ctx.enclosing + type)
+                        listOf(
+                            annotateMessage(
+                                type,
+                                ctx.copy(enclosing = ctx.enclosing + type)
+                            )
                         )
                     }
                 }
             is Enum ->
                 nonGrpc(ctx) {
                     nonDescriptors(ctx) {
-                        annotateEnum(type, ctx)
+                        listOf(EnumBuilder(type, ctx).build())
                     }
                 }
             is Service ->
@@ -93,14 +90,14 @@ object Annotator {
                 )
         }
 
-    private fun nonDescriptors(ctx: Context, gen: () -> String) =
-        nonDescriptors(ctx.desc.context, "", gen)
+    private fun <T> nonDescriptors(ctx: Context, gen: () -> Iterable<T>) =
+        nonDescriptors(ctx.desc.context, emptyList(), gen)
 
     fun <T> nonDescriptors(ctx: ProtocolContext, default: T, gen: () -> T) =
         boolGen(!ctx.onlyGenerateDescriptors, default, gen)
 
-    private fun nonGrpc(ctx: Context, gen: () -> String) =
-        nonGrpc(ctx.desc.context, "", gen)
+    private fun <T> nonGrpc(ctx: Context, gen: () -> Iterable<T>) =
+        nonGrpc(ctx.desc.context, emptyList(), gen)
 
     fun <T> nonGrpc(ctx: ProtocolContext, default: T, gen: () -> T) =
         boolGen(!ctx.onlyGenerateGrpc, default, gen)

diff --git a/...ab/protokt/codegen/impl/AnnotatorUtils.kt → ...tokt/codegen/annotators/AnnotatorUtils.kt b/...ab/protokt/codegen/impl/AnnotatorUtils.kt → ...tokt/codegen/annotators/AnnotatorUtils.kt
@@ -13,19 +13,16 @@
  * limitations under the License.
  */
 
-package com.toasttab.protokt.codegen.impl
+package com.toasttab.protokt.codegen.annotators
 
-import arrow.core.None
-import arrow.core.Some
-import com.toasttab.protokt.codegen.impl.Annotator.Context
+import com.squareup.kotlinpoet.TypeName
+import com.toasttab.protokt.codegen.annotators.Annotator.Context
 import com.toasttab.protokt.codegen.impl.Wrapper.interceptMapKeyTypeName
 import com.toasttab.protokt.codegen.impl.Wrapper.interceptMapValueTypeName
 import com.toasttab.protokt.codegen.protoc.MapEntry
 import com.toasttab.protokt.codegen.protoc.Message
 import com.toasttab.protokt.codegen.protoc.Oneof
 import com.toasttab.protokt.codegen.protoc.StandardField
-import com.toasttab.protokt.codegen.protoc.TypeDesc
-import com.toasttab.protokt.codegen.template.Renderers.ConcatWithScope
 
 fun resolveMapEntry(m: Message) =
     MapEntry(
@@ -36,36 +33,15 @@ fun resolveMapEntry(m: Message) =
 fun resolveMapEntryTypes(f: StandardField, ctx: Context) =
     f.mapEntry!!.let {
         MapTypeParams(
-            interceptMapKeyTypeName(f, it.key.unqualifiedTypeName, ctx)!!,
-            interceptMapValueTypeName(f, it.value.typePClass.renderName(ctx.pkg), ctx)!!
+            interceptMapKeyTypeName(f, it.key.typePClass.toTypeName(), ctx)!!,
+            interceptMapValueTypeName(f, it.value.typePClass.toTypeName(), ctx)!!
         )
     }
 
 class MapTypeParams(
-    val kType: String,
-    val vType: String
+    val kType: TypeName,
+    val vType: TypeName
 )
 
-fun oneOfScope(f: Oneof, type: String, ctx: Context) =
-    ctx.stripEnclosingMessageNamePrefix(
-        ctx.stripRootMessageNamePrefix(
-            ConcatWithScope.render(
-                scope = type,
-                value = f.name
-            )
-        )
-    )
-
-fun String.emptyToNone() =
-    if (isEmpty()) {
-        None
-    } else {
-        Some(this)
-    }
-
-fun kotlinPackage(data: TypeDesc) =
-    resolvePackage(
-        data.desc.options,
-        data.desc.packageName,
-        data.desc.context.respectJavaPackage
-    )
+fun oneOfScope(f: Oneof, type: String) =
+    "$type.${f.name}"