Implement IR implementation for moshi-sealed (#192)

* Abstract out the generator interface * Fixup logic * Add controls to enable * Port sealed tests * Implement IR for sealed Resolves #185 * Update README
ZacSweers · Dec 24, 2021 · 7650423 · 7650423
1 parent 705eccf
commit 7650423
Show file tree

Hide file tree

Showing 21 changed files with 1,944 additions and 941 deletions.
diff --git a/moshi-ir/README.md b/moshi-ir/README.md
@@ -1,10 +1,10 @@
 Moshi-IR
 ========
 
-A Kotlin IR implementation of Moshi code gen.
+A Kotlin IR implementation of Moshi code gen and moshi-sealed code gen.
 
-The goal of this is to have functional parity with Moshi's native Kapt/KSP code gen but run as a fully embedded IR 
-plugin.
+The goal of this is to have functional parity with their native Kapt/KSP code gen analogues but run as a fully 
+embedded IR plugin.
 
 **Benefits**
 - Significantly faster build times.
@@ -35,7 +35,8 @@ these, you should be aware. If you have any issues, you can always fall back to
 
 ### Installation
 
-Simply apply the Gradle plugin in your project to use it.
+Simply apply the Gradle plugin in your project to use it. You can enable moshi-sealed code gen via the `moshi` 
+extension.
 
 The Gradle plugin is published to Maven Central, so ensure you have `mavenCentral()` visible to your buildscript 
 classpath.
@@ -46,6 +47,11 @@ plugins {
   kotlin("jvm")
   id("dev.zacsweers.moshix") version "x.y.z"
 }
+
+moshi {
+  // Opt-in to enable moshi-sealed, disabled by default.
+  enableSealed.set(true)
+}
 ```
 
 Snapshots of the development version are available in [Sonatype's snapshots repository][snapshots].

diff --git a/...iler-plugin/src/main/kotlin/dev/zacsweers/moshix/ir/compiler/MoshiCommandLineProcessor.kt b/...iler-plugin/src/main/kotlin/dev/zacsweers/moshix/ir/compiler/MoshiCommandLineProcessor.kt
@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.config.CompilerConfigurationKey
 
 internal val KEY_ENABLED = CompilerConfigurationKey<Boolean>("enabled")
 internal val KEY_GENERATED_ANNOTATION = CompilerConfigurationKey<String>("generatedAnnotation")
+internal val KEY_ENABLE_SEALED = CompilerConfigurationKey<Boolean>("enableSealed")
 
 @AutoService(CommandLineProcessor::class)
 public class MoshiCommandLineProcessor : CommandLineProcessor {
@@ -33,6 +34,7 @@ public class MoshiCommandLineProcessor : CommandLineProcessor {
   override val pluginOptions: Collection<AbstractCliOption> =
       listOf(
           CliOption("enabled", "<true | false>", "", required = true),
+          CliOption("enableSealed", "<true | false>", "", required = false),
           CliOption("generatedAnnotation", "String", "", required = false),
       )
 
@@ -43,6 +45,7 @@ public class MoshiCommandLineProcessor : CommandLineProcessor {
   ): Unit =
       when (option.optionName) {
         "enabled" -> configuration.put(KEY_ENABLED, value.toBoolean())
+        "enableSealed" -> configuration.put(KEY_ENABLE_SEALED, value.toBoolean())
         "generatedAnnotation" -> configuration.put(KEY_GENERATED_ANNOTATION, value)
         else -> error("Unknown plugin option: ${option.optionName}")
       }

diff --git a/...mpiler-plugin/src/main/kotlin/dev/zacsweers/moshix/ir/compiler/MoshiComponentRegistrar.kt b/...mpiler-plugin/src/main/kotlin/dev/zacsweers/moshix/ir/compiler/MoshiComponentRegistrar.kt
@@ -33,12 +33,13 @@ public class MoshiComponentRegistrar : ComponentRegistrar {
   ) {
 
     if (configuration[KEY_ENABLED] == false) return
+    val enableSealed = configuration[KEY_ENABLE_SEALED] ?: false
 
     val messageCollector =
         configuration.get(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, MessageCollector.NONE)
     val fqGeneratedAnnotation = configuration[KEY_GENERATED_ANNOTATION]?.let(::FqName)
 
     IrGenerationExtension.registerExtension(
-        project, MoshiIrGenerationExtension(messageCollector, fqGeneratedAnnotation))
+        project, MoshiIrGenerationExtension(messageCollector, fqGeneratedAnnotation, enableSealed))
   }
 }
diff --git a/...ler-plugin/src/main/kotlin/dev/zacsweers/moshix/ir/compiler/MoshiIrGenerationExtension.kt b/...ler-plugin/src/main/kotlin/dev/zacsweers/moshix/ir/compiler/MoshiIrGenerationExtension.kt
@@ -24,7 +24,8 @@ import org.jetbrains.kotlin.name.FqName
 
 internal class MoshiIrGenerationExtension(
     private val messageCollector: MessageCollector,
-    private val generatedAnnotationName: FqName?
+    private val generatedAnnotationName: FqName?,
+    private val enableSealed: Boolean
 ) : IrGenerationExtension {
 
   override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
@@ -40,7 +41,12 @@ internal class MoshiIrGenerationExtension(
     val deferred = mutableListOf<GeneratedAdapter>()
     val moshiTransformer =
         MoshiIrVisitor(
-            moduleFragment, pluginContext, messageCollector, generatedAnnotation, deferred)
+            moduleFragment,
+            pluginContext,
+            messageCollector,
+            generatedAnnotation,
+            enableSealed,
+            deferred)
     moduleFragment.transform(moshiTransformer, null)
     for ((file, adapters) in deferred.groupBy { it.irFile }) {
       for (adapter in adapters) {

diff --git a/.../moshi-compiler-plugin/src/main/kotlin/dev/zacsweers/moshix/ir/compiler/MoshiIrVisitor.kt b/.../moshi-compiler-plugin/src/main/kotlin/dev/zacsweers/moshix/ir/compiler/MoshiIrVisitor.kt
@@ -15,8 +15,10 @@
  */
 package dev.zacsweers.moshix.ir.compiler
 
-import dev.zacsweers.moshix.ir.compiler.api.AdapterGenerator
+import dev.zacsweers.moshix.ir.compiler.api.MoshiAdapterGenerator
 import dev.zacsweers.moshix.ir.compiler.api.PropertyGenerator
+import dev.zacsweers.moshix.ir.compiler.sealed.MoshiSealedSymbols
+import dev.zacsweers.moshix.ir.compiler.sealed.SealedAdapterGenerator
 import dev.zacsweers.moshix.ir.compiler.util.error
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
@@ -44,14 +46,19 @@ internal class MoshiIrVisitor(
     private val pluginContext: IrPluginContext,
     private val messageCollector: MessageCollector,
     private val generatedAnnotation: IrClassSymbol?,
+    private val enableSealed: Boolean,
     private val deferredAddedClasses: MutableList<GeneratedAdapter>
 ) : IrElementTransformerVoidWithContext() {
 
-  private val moshiSymbols = MoshiSymbols(pluginContext.irBuiltIns, moduleFragment, pluginContext)
+  private val moshiSymbols by lazy {
+    MoshiSymbols(pluginContext.irBuiltIns, moduleFragment, pluginContext)
+  }
+
+  private val moshiSealedSymbols by lazy { MoshiSealedSymbols(pluginContext) }
 
   private fun adapterGenerator(
       originalType: IrClass,
-  ): AdapterGenerator? {
+  ): MoshiAdapterGenerator? {
     val type = targetType(originalType, pluginContext, messageCollector) ?: return null
 
     val properties = mutableMapOf<String, PropertyGenerator>()
@@ -89,7 +96,7 @@ internal class MoshiIrVisitor(
           }
         }
 
-    return AdapterGenerator(pluginContext, moshiSymbols, type, sortedProperties)
+    return MoshiAdapterGenerator(pluginContext, moshiSymbols, type, sortedProperties)
   }
 
   override fun visitClassNew(declaration: IrClass): IrStatement {
@@ -101,34 +108,45 @@ internal class MoshiIrVisitor(
           return super.visitClassNew(declaration)
         }
 
-        if (call.valueArgumentsCount >= 2) {
-          // This is generator
-          call.getValueArgument(1)?.let { generator ->
-            @Suppress("UNCHECKED_CAST")
-            if ((generator as IrConst<String>).value.isNotBlank()) {
-              return super.visitClassNew(declaration)
+        // This is generator
+        @Suppress("UNCHECKED_CAST")
+        val generatorValue = (call.getValueArgument(1) as? IrConst<String>?)?.value.orEmpty()
+        val generator =
+            if (generatorValue.isNotBlank()) {
+              if (enableSealed && generatorValue.startsWith("sealed:")) {
+                val typeLabel = generatorValue.removePrefix("sealed:")
+                SealedAdapterGenerator(
+                    pluginContext,
+                    messageCollector,
+                    moshiSymbols,
+                    moshiSealedSymbols,
+                    declaration,
+                    typeLabel)
+              } else {
+                return super.visitClassNew(declaration)
+              }
+            } else {
+              // Unspecified/null - means it's empty/default.
+              adapterGenerator(declaration)
             }
-          }
-        }
 
-        val adapterGenerator =
-            adapterGenerator(declaration) ?: return super.visitClassNew(declaration)
-        pluginContext.irFactory.run {
-          try {
-            val adapterClass = adapterGenerator.prepare()
-            if (generatedAnnotation != null) {
-              // TODO add generated annotation
-            }
-            // Uncomment for debugging generated code
-            //            println("Dumping current IR src")
-            //            println(adapterClass.adapterClass.dumpSrc())
-            deferredAddedClasses += GeneratedAdapter(adapterClass.adapterClass, declaration.file)
-          } catch (e: Exception) {
-            messageCollector.error(declaration) {
-              "Error preparing adapter for ${declaration.fqNameWhenAvailable}"
-            }
-            throw e
+        val adapterGenerator = generator ?: return super.visitClassNew(declaration)
+        try {
+          val adapterClass = adapterGenerator.prepare() ?: return super.visitClassNew(declaration)
+          if (generatedAnnotation != null) {
+            // TODO add generated annotation
+          }
+          // Uncomment for debugging generated code
+          //            println("Dumping current IR src")
+          //          messageCollector.report(
+          //              CompilerMessageSeverity.STRONG_WARNING,
+          // adapterClass.adapterClass.dumpSrc())
+          deferredAddedClasses += GeneratedAdapter(adapterClass.adapterClass, declaration.file)
+        } catch (e: Exception) {
+          messageCollector.error(declaration) {
+            "Error preparing adapter for ${declaration.fqNameWhenAvailable}"
           }
+          throw e
         }
       }
     }

diff --git a/...ir/moshi-compiler-plugin/src/main/kotlin/dev/zacsweers/moshix/ir/compiler/MoshiSymbols.kt b/...ir/moshi-compiler-plugin/src/main/kotlin/dev/zacsweers/moshix/ir/compiler/MoshiSymbols.kt
@@ -21,10 +21,13 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.builders.declarations.addConstructor
+import org.jetbrains.kotlin.ir.builders.declarations.addExtensionReceiver
 import org.jetbrains.kotlin.ir.builders.declarations.addFunction
+import org.jetbrains.kotlin.ir.builders.declarations.addGetter
 import org.jetbrains.kotlin.ir.builders.declarations.addTypeParameter
 import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
 import org.jetbrains.kotlin.ir.builders.declarations.buildClass
+import org.jetbrains.kotlin.ir.builders.declarations.buildProperty
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationParent
 import org.jetbrains.kotlin.ir.declarations.IrFactory
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
@@ -34,6 +37,7 @@ import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.declarations.impl.IrExternalPackageFragmentImpl
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
+import org.jetbrains.kotlin.ir.symbols.IrPropertySymbol
 import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.IrTypeArgument
@@ -43,10 +47,12 @@ import org.jetbrains.kotlin.ir.types.createType
 import org.jetbrains.kotlin.ir.types.defaultType
 import org.jetbrains.kotlin.ir.types.makeNotNull
 import org.jetbrains.kotlin.ir.types.makeNullable
+import org.jetbrains.kotlin.ir.types.starProjectedType
 import org.jetbrains.kotlin.ir.types.typeWith
 import org.jetbrains.kotlin.ir.util.constructors
 import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.util.functions
+import org.jetbrains.kotlin.ir.util.getSimpleFunction
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
@@ -57,6 +63,8 @@ internal class MoshiSymbols(
 ) {
   private val irFactory: IrFactory = pluginContext.irFactory
 
+  private val javaLang: IrPackageFragment by lazy { createPackage("java.lang") }
+  private val kotlinJvm: IrPackageFragment by lazy { createPackage("kotlin.jvm") }
   private val moshiPackage: IrPackageFragment by lazy { createPackage("com.squareup.moshi") }
   private val javaReflectPackage: IrPackageFragment by lazy { createPackage("java.lang.reflect") }
 
@@ -122,6 +130,12 @@ internal class MoshiSymbols(
         .symbol
   }
 
+  val moshiBuilder by lazy {
+    pluginContext.referenceClass(FqName("com.squareup.moshi.Moshi.Builder"))!!
+  }
+
+  val moshiBuilderBuild by lazy { moshiBuilder.getSimpleFunction("build")!! }
+
   val moshi: IrClassSymbol by lazy {
     irFactory
         .buildClass {
@@ -144,10 +158,22 @@ internal class MoshiSymbols(
                     "annotations", irBuiltIns.setClass.typeWith(irBuiltIns.annotationType))
                 addValueParameter("fieldName", irBuiltIns.stringType)
               }
+
+          addFunction(
+              Name.identifier("newBuilder").identifier, moshiBuilder.defaultType, Modality.FINAL)
         }
         .symbol
   }
 
+  val moshiThreeArgAdapter by lazy { moshi.getSimpleFunction("adapter")!! }
+
+  val moshiNewBuilder by lazy { moshi.getSimpleFunction("newBuilder")!! }
+
+  val jsonAdapterFactoryCreate by lazy {
+    pluginContext.referenceClass(FqName("com.squareup.moshi.JsonAdapter.Factory"))!!
+        .getSimpleFunction("create")!!
+  }
+
   internal val jsonAdapter: IrClassSymbol by lazy {
     irFactory
         .buildClass {
@@ -177,6 +203,10 @@ internal class MoshiSymbols(
         .symbol
   }
 
+  val addAdapter by lazy {
+    pluginContext.referenceFunctions(FqName("com.squareup.moshi.addAdapter")).first()
+  }
+
   val jsonDataException: IrClassSymbol by lazy {
     pluginContext.referenceClass(FqName("com.squareup.moshi.JsonDataException"))!!
   }
@@ -247,6 +277,22 @@ internal class MoshiSymbols(
     }
   }
 
+  val javaLangClass: IrClassSymbol by lazy {
+    createClass(javaLang, "Class", ClassKind.CLASS, Modality.FINAL)
+  }
+
+  val kotlinKClassJava: IrPropertySymbol =
+      irFactory
+          .buildProperty() { name = Name.identifier("java") }
+          .apply {
+            parent = kotlinJvm
+            addGetter().apply {
+              addExtensionReceiver(irBuiltIns.kClassClass.starProjectedType)
+              returnType = javaLangClass.defaultType
+            }
+          }
+          .symbol
+
   private fun createPackage(packageName: String): IrPackageFragment =
       IrExternalPackageFragmentImpl.createEmptyExternalPackageFragment(
           moduleFragment.descriptor, FqName(packageName))