From acb52423680b5af8ffb0fdb7a5dfd18e406b8ec8 Mon Sep 17 00:00:00 2001
From: Mattia Iavarone <mattia@deepmedia.io>
Date: Sun, 8 Sep 2024 18:28:48 +0200
Subject: [PATCH] Support companion objects

---
 .../src/main/kotlin/utils/PoetUtils.kt               |  5 ++++-
 .../io/deepmedia/tools/knee/tests/ObjectTests.kt     | 12 ++++++++++++
 .../src/backendMain/kotlin/ObjectDefinitions.kt      | 10 ++++++++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/knee-compiler-plugin/src/main/kotlin/utils/PoetUtils.kt b/knee-compiler-plugin/src/main/kotlin/utils/PoetUtils.kt
index 7bec3f8..0b96e73 100644
--- a/knee-compiler-plugin/src/main/kotlin/utils/PoetUtils.kt
+++ b/knee-compiler-plugin/src/main/kotlin/utils/PoetUtils.kt
@@ -22,7 +22,10 @@ fun IrClass.asTypeSpec(rename: ((String) -> String)? = null): TypeSpec.Builder {
     val name = codegenName.map { rename?.invoke(it) ?: it }.asString()
     return when (kind) {
         ClassKind.ENUM_CLASS -> TypeSpec.enumBuilder(name)
-        ClassKind.OBJECT -> TypeSpec.objectBuilder(name)
+        ClassKind.OBJECT -> when {
+            isCompanion -> TypeSpec.companionObjectBuilder(if (name == "Companion") null else name)
+            else -> TypeSpec.objectBuilder(name)
+        }
         ClassKind.INTERFACE -> TypeSpec.interfaceBuilder(name)
         ClassKind.ANNOTATION_CLASS -> TypeSpec.annotationBuilder(name)
         ClassKind.ENUM_ENTRY -> error("Enum entries ($this) can't become a TypeSpec.")
diff --git a/tests/test-classes/src/androidInstrumentedTest/kotlin/io/deepmedia/tools/knee/tests/ObjectTests.kt b/tests/test-classes/src/androidInstrumentedTest/kotlin/io/deepmedia/tools/knee/tests/ObjectTests.kt
index 97e0f79..b4f81fa 100644
--- a/tests/test-classes/src/androidInstrumentedTest/kotlin/io/deepmedia/tools/knee/tests/ObjectTests.kt
+++ b/tests/test-classes/src/androidInstrumentedTest/kotlin/io/deepmedia/tools/knee/tests/ObjectTests.kt
@@ -41,4 +41,16 @@ class ObjectTests {
         check(TopLevelObject.value == 15) { TopLevelObject.value }
     }
 
+    @Test
+    fun testInner_property() {
+        ObjectParent.InnerObject.value = 15
+        check(ObjectParent.InnerObject.value == 15) { ObjectParent.InnerObject.value }
+    }
+
+    @Test
+    fun testCompanion_property() {
+        ObjectParent.value = 15
+        check(ObjectParent.value == 15) { ObjectParent.value }
+    }
+
 }
diff --git a/tests/test-classes/src/backendMain/kotlin/ObjectDefinitions.kt b/tests/test-classes/src/backendMain/kotlin/ObjectDefinitions.kt
index 9088278..bfe750c 100644
--- a/tests/test-classes/src/backendMain/kotlin/ObjectDefinitions.kt
+++ b/tests/test-classes/src/backendMain/kotlin/ObjectDefinitions.kt
@@ -13,3 +13,13 @@ object TopLevelObject {
     @Knee fun decrement() { value -= 1 }
     @Knee override fun toString(): String = "TopLevelObject($value)"
 }
+
+class ObjectParent {
+    @KneeObject object InnerObject {
+        @Knee var value: Int = 0
+    }
+
+    @KneeObject companion object {
+        @Knee var value: Int = 0
+    }
+}