From 4f4a27c18ae223d6bdd9674931623fce0320803e Mon Sep 17 00:00:00 2001
From: Alejandro Serrano <trupill@gmail.com>
Date: Wed, 28 Feb 2024 10:37:14 +0100
Subject: [PATCH] Ignore properties in sealed hierarchies where the type
 changes across children (#3382)

* Fixes #3380
* Add test for #3881
---
 .../arrow/optics/plugin/internals/info.kt     |  6 +++
 .../optics/plugin/internals/processor.kt      | 20 +++++++-
 .../kotlin/arrow/optics/plugin/LensTests.kt   | 48 +++++++++++++++++++
 3 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/arrow-libs/optics/arrow-optics-ksp-plugin/src/main/kotlin/arrow/optics/plugin/internals/info.kt b/arrow-libs/optics/arrow-optics-ksp-plugin/src/main/kotlin/arrow/optics/plugin/internals/info.kt
index 2f2d8058c70..e3f25400828 100644
--- a/arrow-libs/optics/arrow-optics-ksp-plugin/src/main/kotlin/arrow/optics/plugin/internals/info.kt
+++ b/arrow-libs/optics/arrow-optics-ksp-plugin/src/main/kotlin/arrow/optics/plugin/internals/info.kt
@@ -23,3 +23,9 @@ val String.sealedLensConstructorOverridesOnly
       |                                       ^
       |To generate Lens for a sealed class make sure all children override target properties in constructors.
     """.trimMargin()
+
+val String.sealedLensChangedProperties
+  get() =
+    """
+      |Ignoring $this when generating lenses; the type is not uniform across all children.
+    """.trimMargin()
diff --git a/arrow-libs/optics/arrow-optics-ksp-plugin/src/main/kotlin/arrow/optics/plugin/internals/processor.kt b/arrow-libs/optics/arrow-optics-ksp-plugin/src/main/kotlin/arrow/optics/plugin/internals/processor.kt
index 17da0477d4d..db697683b31 100644
--- a/arrow-libs/optics/arrow-optics-ksp-plugin/src/main/kotlin/arrow/optics/plugin/internals/processor.kt
+++ b/arrow-libs/optics/arrow-optics-ksp-plugin/src/main/kotlin/arrow/optics/plugin/internals/processor.kt
@@ -8,6 +8,7 @@ import com.google.devtools.ksp.isAbstract
 import com.google.devtools.ksp.processing.KSPLogger
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.KSDeclaration
+import com.google.devtools.ksp.symbol.KSPropertyDeclaration
 import com.google.devtools.ksp.symbol.KSType
 import com.google.devtools.ksp.symbol.KSTypeArgument
 import com.google.devtools.ksp.symbol.KSTypeParameter
@@ -119,7 +120,13 @@ internal fun evalAnnotatedClass(
         return null
       }
 
-      val propertyNames = properties
+      val (goodProperties, badProperties) = properties.partition { it.sameInChildren(subclasses) }
+
+      badProperties.forEach {
+        logger.info(it.qualifiedNameOrSimpleName.sealedLensChangedProperties, element)
+      }
+
+      val propertyNames = goodProperties
         .map { it.simpleName.asString() }
 
       val nonConstructorOverrides = subclasses
@@ -133,7 +140,7 @@ internal fun evalAnnotatedClass(
         return null
       }
 
-      properties
+      goodProperties
         .map { it.type.resolve().qualifiedString() }
         .zip(propertyNames) { type, name ->
           Focus(
@@ -151,6 +158,15 @@ internal fun evalAnnotatedClass(
   }
 }
 
+fun KSPropertyDeclaration.sameInChildren(subclasses: Sequence<KSClassDeclaration>): Boolean =
+  subclasses.all { subclass ->
+    val property = subclass.getAllProperties().singleOrNull { it.simpleName == this.simpleName }
+    when (property) {
+      null -> false
+      else -> property.type.resolve() == this.type.resolve()
+    }
+  }
+
 internal fun evalAnnotatedDslElement(element: KSClassDeclaration, logger: KSPLogger): Target =
   when {
     element.isDataClass ->
diff --git a/arrow-libs/optics/arrow-optics-ksp-plugin/src/test/kotlin/arrow/optics/plugin/LensTests.kt b/arrow-libs/optics/arrow-optics-ksp-plugin/src/test/kotlin/arrow/optics/plugin/LensTests.kt
index e352498f77d..bcbeb87d7b0 100755
--- a/arrow-libs/optics/arrow-optics-ksp-plugin/src/test/kotlin/arrow/optics/plugin/LensTests.kt
+++ b/arrow-libs/optics/arrow-optics-ksp-plugin/src/test/kotlin/arrow/optics/plugin/LensTests.kt
@@ -208,4 +208,52 @@ class LensTests {
       |val r = l != null
       """.compilationFails()
   }
+
+  @Test
+  fun `Lens for sealed class property, ignoring changed nullability`() {
+    """
+      |$`package`
+      |$imports
+      |@optics
+      |sealed class LensSealed {
+      |  abstract val property1: String?
+      |  
+      |  data class dataChild1(override val property1: String?) : LensSealed()
+      |  data class dataChild2(override val property1: String?, val number: Int) : LensSealed()
+      |  data class dataChild3(override val property1: String, val enabled: Boolean) : LensSealed()
+      |   
+      |  companion object 
+      |}
+      |
+      |val l: Lens<LensSealed, String>? = LensSealed.property1
+      |val r = l != null
+      """.compilationFails()
+  }
+
+  @Test
+  fun `Lens for sealed class property, ignoring changed types`() {
+    """
+      |$`package`
+      |$imports
+      |@optics
+      |sealed interface Base<out T> {
+      |    val prop: T
+      |
+      |    companion object
+      |}
+      |
+      |@optics
+      |data class Child1(override val prop: String) : Base<String> {
+      |    companion object
+      |}
+      |
+      |@optics
+      |data class Child2(override val prop: Int) : Base<Int> {
+      |    companion object
+      |}
+      |
+      |val l: Lens<Base<String>, String> = Base.prop()
+      |val r = l != null
+      """.compilationFails()
+  }
 }