Parcelize: Allow Parcelize annotation on sealed classes

See https://issuetracker.google.com/177856520

plugins/parcelize/parcelize-compiler/src/org/jetbrains/kotlin/parcelize/ParcelizeDeclarationChecker.kt

@@ -9,10 +9,7 @@ import org.jetbrains.kotlin.codegen.ClassBuilderMode
 import org.jetbrains.kotlin.codegen.FrameMap
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.config.LanguageVersionSettings
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.descriptors.PropertyDescriptor
-import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.diagnostics.DiagnosticSink
 import org.jetbrains.kotlin.lexer.KtTokens
@@ -153,7 +150,7 @@ open class ParcelizeDeclarationChecker : DeclarationChecker {
             return
         }
 
-        if (declaration is KtClass && (declaration.isAnnotation() || declaration.isInterface())) {
+        if (declaration is KtClass && (declaration.isAnnotation() || declaration.isInterface() && !declaration.isSealed())) {
             val reportElement = declaration.nameIdentifier ?: declaration
             diagnosticHolder.report(ErrorsParcelize.PARCELABLE_SHOULD_BE_CLASS.on(reportElement))
             return
@@ -166,10 +163,9 @@ open class ParcelizeDeclarationChecker : DeclarationChecker {
             }
         }
 
-        val sealedOrAbstract =
-            declaration.modifierList?.let { it.getModifier(KtTokens.ABSTRACT_KEYWORD) ?: it.getModifier(KtTokens.SEALED_KEYWORD) }
-        if (sealedOrAbstract != null) {
-            diagnosticHolder.report(ErrorsParcelize.PARCELABLE_SHOULD_BE_INSTANTIABLE.on(sealedOrAbstract))
+        val abstractModifier = declaration.modifierList?.let { it.getModifier(KtTokens.ABSTRACT_KEYWORD) }
+        if (abstractModifier != null) {
+            diagnosticHolder.report(ErrorsParcelize.PARCELABLE_SHOULD_BE_INSTANTIABLE.on(abstractModifier))
         }
 
         if (declaration is KtClass && declaration.isInner()) {

plugins/parcelize/parcelize-compiler/src/org/jetbrains/kotlin/parcelize/ParcelizeResolveExtension.kt

@@ -30,8 +30,11 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.parcelize.ParcelizeSyntheticComponent.ComponentKind.DESCRIBE_CONTENTS
 import org.jetbrains.kotlin.parcelize.ParcelizeSyntheticComponent.ComponentKind.WRITE_TO_PARCEL
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperClassNotAny
+import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperInterfaces
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.extensions.SyntheticResolveExtension
 import org.jetbrains.kotlin.resolve.source.PsiSourceElement
@@ -122,6 +125,7 @@ open class ParcelizeResolveExtension : SyntheticResolveExtension {
 
         if (name.asString() == DESCRIBE_CONTENTS.methodName
             && thisDescriptor.isParcelize
+            && !thisDescriptor.isSealed()
             && isParcelizePluginEnabled()
             && result.none { it.isDescribeContents() }
             && fromSupertypes.none { it.isDescribeContents() }
@@ -129,6 +133,7 @@ open class ParcelizeResolveExtension : SyntheticResolveExtension {
             result += createMethod(thisDescriptor, DESCRIBE_CONTENTS, Modality.OPEN, thisDescriptor.builtIns.intType)
         } else if (name.asString() == WRITE_TO_PARCEL.methodName
             && thisDescriptor.isParcelize
+            && !thisDescriptor.isSealed()
             && isParcelizePluginEnabled()
             && result.none { it.isWriteToParcel() }
         ) {
@@ -197,16 +202,13 @@ val RAW_VALUE_ANNOTATION_FQ_NAMES = listOf(
 internal val PARCELER_FQNAME = FqName("kotlinx.parcelize.Parceler")
 internal val OLD_PARCELER_FQNAME = FqName("kotlinx.android.parcel.Parceler")
 
-val ClassDescriptor.isParcelize: Boolean
-    get() {
-        for (fqName in PARCELIZE_CLASS_FQ_NAMES) {
-            if (this.annotations.hasAnnotation(fqName)) {
-                return true
-            }
-        }
+val ClassDescriptor.hasParcelizeAnnotation: Boolean
+    get() = PARCELIZE_CLASS_FQ_NAMES.any(annotations::hasAnnotation)
 
-        return false
-    }
+val ClassDescriptor.isParcelize: Boolean
+    get() = hasParcelizeAnnotation
+            || getSuperClassNotAny()?.takeIf(DescriptorUtils::isSealedClass)?.hasParcelizeAnnotation == true
+            || getSuperInterfaces().any { DescriptorUtils.isSealedClass(it) && it.hasParcelizeAnnotation }
 
 val KotlinType.isParceler: Boolean
     get() = constructor.declarationDescriptor?.fqNameSafe == PARCELER_FQNAME

plugins/parcelize/parcelize-compiler/src/org/jetbrains/kotlin/parcelize/diagnostic/DefaultErrorMessagesParcelize.kt

@@ -43,7 +43,7 @@ object DefaultErrorMessagesParcelize : DefaultErrorMessages.Extension {
 
         MAP.put(
             ErrorsParcelize.PARCELABLE_SHOULD_BE_INSTANTIABLE,
-            "'Parcelable' should not be a 'sealed' or 'abstract' class"
+            "'Parcelable' should not be an 'abstract' class"
         )
 
         MAP.put(

plugins/parcelize/parcelize-compiler/src/org/jetbrains/kotlin/parcelize/ir/ParcelizeIrTransformer.kt

@@ -122,7 +122,10 @@ class ParcelizeIrTransformer(private val context: IrPluginContext, private val a
 
     override fun visitClass(declaration: IrClass) {
         declaration.acceptChildren(this, null)
-        if (!declaration.isParcelize)
+
+        // Sealed classes can be annotated with `@Parcelize`, but that only implies that we
+        // should process their immediate subclasses.
+        if (!declaration.isParcelize || declaration.modality == Modality.SEALED)
             return
 
         val parcelableProperties = declaration.parcelableProperties

plugins/parcelize/parcelize-compiler/src/org/jetbrains/kotlin/parcelize/ir/irUtils.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.parcelize.ir
 import org.jetbrains.kotlin.backend.common.ir.allOverridden
 import org.jetbrains.kotlin.backend.jvm.ir.erasedUpperBound
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.*
@@ -34,7 +35,12 @@ import org.jetbrains.kotlin.types.Variance
 
 // true if the class should be processed by the parcelize plugin
 val IrClass.isParcelize: Boolean
-    get() = kind in ParcelizeExtensionBase.ALLOWED_CLASS_KINDS && hasAnyAnnotation(PARCELIZE_CLASS_FQ_NAMES)
+    get() = kind in ParcelizeExtensionBase.ALLOWED_CLASS_KINDS &&
+            (hasAnyAnnotation(PARCELIZE_CLASS_FQ_NAMES) || superTypes.any { superType ->
+                superType.classOrNull?.owner?.let {
+                    it.modality == Modality.SEALED && it.hasAnyAnnotation(PARCELIZE_CLASS_FQ_NAMES)
+                } == true
+            })
 
 // Finds the getter for a pre-existing CREATOR field on the class companion, which is used for manual Parcelable implementations in Kotlin.
 val IrClass.creatorGetter: IrSimpleFunctionSymbol?

plugins/parcelize/parcelize-compiler/testData/box/sealedClass2.kt

@@ -0,0 +1,33 @@
+// WITH_RUNTIME
+
+@file:JvmName("TestKt")
+package test
+
+import kotlinx.parcelize.*
+import android.os.Parcel
+import android.os.Parcelable
+
+@Parcelize
+sealed class Foo : Parcelable {
+    data class A(val x: Int) : Foo()
+    object B : Foo()
+}
+
+data class C(val x: String) : Foo()
+
+@Parcelize
+data class Bar(val a: Foo.A, val b: Foo.B, val c: C, val foo: Foo) : Parcelable
+
+fun box() = parcelTest { parcel ->
+    val first = Bar(Foo.A(1024), Foo.B, C("OK"), Foo.A(1))
+
+    first.writeToParcel(parcel, 0)
+
+    val bytes = parcel.marshall()
+    parcel.unmarshall(bytes, 0, bytes.size)
+    parcel.setDataPosition(0)
+
+    val second = parcelableCreator<Bar>().createFromParcel(parcel)
+
+    assert(first == second)
+}

plugins/parcelize/parcelize-compiler/testData/box/sealedInterface.kt

@@ -0,0 +1,74 @@
+// WITH_RUNTIME
+
+@file:JvmName("TestKt")
+package test
+
+import kotlinx.parcelize.*
+import android.os.Parcel
+import android.os.Parcelable
+
+@Parcelize
+sealed interface I : Parcelable
+
+@Parcelize
+sealed class A : Parcelable
+
+abstract class B
+interface J
+
+data class AI(val x: String) : A(), I
+class I1 : J, I {
+    override fun equals(other: Any?): Boolean {
+        return other is I1
+    }
+}
+data class I2(val x: Float) : B(), I
+
+object A1 : A()
+open class A2(val x: Int) : A() {
+    override fun equals(other: Any?): Boolean {
+        return other is A2 && other::class == A2::class && x == other.x
+    }
+}
+
+@Parcelize
+class A3 : A2(3) {
+    override fun equals(other: Any?): Boolean {
+        return other is A3 && x == other.x
+    }
+}
+
+@Parcelize
+data class C(
+    val a: A,
+    val i: I,
+    val a1: A1,
+    val a2: A2,
+    val a3: A3,
+    val ai: AI,
+    val i1: I1,
+    val i2: I2,
+) : Parcelable
+
+fun box() = parcelTest { parcel ->
+    val first = C(
+        AI("0"),
+        AI("1"),
+        A1,
+        A2(2),
+        A3(),
+        AI("4"),
+        I1(),
+        I2(5.0f),
+    )
+
+    first.writeToParcel(parcel, 0)
+
+    val bytes = parcel.marshall()
+    parcel.unmarshall(bytes, 0, bytes.size)
+    parcel.setDataPosition(0)
+
+    val second = parcelableCreator<C>().createFromParcel(parcel)
+
+    assert(first == second)
+}

plugins/parcelize/parcelize-compiler/tests/org/jetbrains/kotlin/parcelize/test/ParcelizeBoxTestGenerated.java

@@ -350,6 +350,16 @@ public class ParcelizeBoxTestGenerated extends AbstractParcelizeBoxTest {
         runTest("plugins/parcelize/parcelize-compiler/testData/box/sealedClass.kt");
     }
 
+    @TestMetadata("sealedClass2.kt")
+    public void testSealedClass2() throws Exception {
+        runTest("plugins/parcelize/parcelize-compiler/testData/box/sealedClass2.kt");
+    }
+
+    @TestMetadata("sealedInterface.kt")
+    public void testSealedInterface() throws Exception {
+        runTest("plugins/parcelize/parcelize-compiler/testData/box/sealedInterface.kt");
+    }
+
     @TestMetadata("shortArray.kt")
     public void testShortArray() throws Exception {
         runTest("plugins/parcelize/parcelize-compiler/testData/box/shortArray.kt");

plugins/parcelize/parcelize-compiler/tests/org/jetbrains/kotlin/parcelize/test/ParcelizeIrBoxTestGenerated.java

@@ -350,6 +350,16 @@ public class ParcelizeIrBoxTestGenerated extends AbstractParcelizeIrBoxTest {
         runTest("plugins/parcelize/parcelize-compiler/testData/box/sealedClass.kt");
     }
 
+    @TestMetadata("sealedClass2.kt")
+    public void testSealedClass2() throws Exception {
+        runTest("plugins/parcelize/parcelize-compiler/testData/box/sealedClass2.kt");
+    }
+
+    @TestMetadata("sealedInterface.kt")
+    public void testSealedInterface() throws Exception {
+        runTest("plugins/parcelize/parcelize-compiler/testData/box/sealedInterface.kt");
+    }
+
     @TestMetadata("shortArray.kt")
     public void testShortArray() throws Exception {
         runTest("plugins/parcelize/parcelize-compiler/testData/box/shortArray.kt");