From 9145cf703586f75edf693468a6ffb10550afc33a Mon Sep 17 00:00:00 2001
From: mcpiroman <mcpiroman@gmail.com>
Date: Sat, 4 Mar 2023 18:57:59 +0100
Subject: [PATCH] IR tree gen: allow lists of nullable elements and arrays

---
 .../src/org/jetbrains/kotlin/ir/util/transform.kt | 11 +++++++++++
 .../kotlin/ir/generator/config/ConfigModel.kt     |  3 ++-
 .../kotlin/ir/generator/model/Transformations.kt  | 15 +++++++++++----
 .../kotlin/ir/generator/print/Elements.kt         |  7 ++++++-
 4 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/transform.kt b/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/transform.kt
index e1e8d4a3409..74b44191a11 100644
--- a/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/transform.kt
+++ b/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/transform.kt
@@ -36,6 +36,17 @@ fun <T : IrElement, D> MutableList<T>.transformInPlace(transformer: IrElementTra
     }
 }
 
+fun <T : IrElement, D> Array<T?>.transformInPlace(transformer: IrElementTransformer<D>, data: D) {
+    for (i in indices) {
+        // Cast to IrElementBase to avoid casting to interface and invokeinterface, both of which are slow.
+        val element = get(i) as IrElementBase?
+        if (element != null) {
+            @Suppress("UNCHECKED_CAST")
+            set(i, element.transform(transformer, data) as T)
+        }
+    }
+}
+
 /**
  * Transforms a mutable list in place.
  * Each element `it` is replaced with a result of `transformation(it)`,
diff --git a/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/config/ConfigModel.kt b/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/config/ConfigModel.kt
index 8a8c23b6d91..8ffb02de058 100644
--- a/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/config/ConfigModel.kt
+++ b/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/config/ConfigModel.kt
@@ -116,6 +116,7 @@ class ListFieldConfig(
     enum class Mutability {
         Immutable,
         Var,
-        List
+        List,
+        Array
     }
 }
diff --git a/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/model/Transformations.kt b/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/model/Transformations.kt
index e29a0c4bf16..92254722fe6 100644
--- a/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/model/Transformations.kt
+++ b/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/model/Transformations.kt
@@ -32,10 +32,17 @@ fun config2model(config: Config): Model {
                     fc.baseGetter
                 )
                 is ListFieldConfig -> {
-                    val listType = if (fc.mutability == ListFieldConfig.Mutability.List) type(
-                        "kotlin.collections",
-                        "MutableList"
-                    ) else type("kotlin.collections", "List")
+                    val listType = when (fc.mutability) {
+                        ListFieldConfig.Mutability.List -> type(
+                            "kotlin.collections",
+                            "MutableList"
+                        )
+                        ListFieldConfig.Mutability.Array -> type(
+                            "kotlin.",
+                            "Array"
+                        )
+                        else -> type("kotlin.collections", "List")
+                    }
                     ListField(
                         fc,
                         fc.name,
diff --git a/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/print/Elements.kt b/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/print/Elements.kt
index 7ed60053bff..8b4e6bd54d9 100644
--- a/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/print/Elements.kt
+++ b/compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/print/Elements.kt
@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.ir.generator.elementTransformerType
 import org.jetbrains.kotlin.ir.generator.elementVisitorType
 import org.jetbrains.kotlin.ir.generator.model.*
 import org.jetbrains.kotlin.ir.generator.util.TypeKind
+import org.jetbrains.kotlin.ir.generator.util.TypeRefWithNullability
 import org.jetbrains.kotlin.ir.generator.util.tryParameterizedBy
 import java.io.File
 
@@ -161,7 +162,11 @@ fun printElements(generationPath: File, model: Model) = sequence {
                             if (child.nullable) append("?")
                             when (child) {
                                 is SingleField -> append(".%N(%N, %N)")
-                                is ListField -> append(".forEach { it.%N(%N, %N) }")
+                                is ListField -> {
+                                    append(".forEach { it")
+                                    if ((child.elementType as? TypeRefWithNullability)?.nullable == true) append("?")
+                                    append(".%N(%N, %N) }")
+                                }
                             }
                         }, child.name, acceptMethodName, visitorParam, dataParam)
                     }