Repository: JetBrains/kotlin
File: compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/transform.kt (Kotlin)

Code before:
import org.jetbrains.kotlin.ir.declarations.IrDeclarationContainer
import org.jetbrains.kotlin.ir.visitors.IrElementTransformer

inline fun <reified T : IrElement> MutableList<T>.transformInPlace(transformation: (T) -> IrElement) {
    for (i in 0 until size) {
        set(i, transformation(get(i)) as T)
    }
}

fun <T : IrElement, D> MutableList<T>.transformInPlace(transformer: IrElementTransformer<D>, data: D) {
    for (i in 0 until size) {
        // Cast to IrElementBase to avoid casting to interface and invokeinterface, both of which are slow.
        @Suppress("UNCHECKED_CAST")
        set(i, (get(i) as IrElementBase).transform(transformer, data) as T)
    }
}

fun <T : IrElement, D> Array<T?>.transformInPlace(transformer: IrElementTransformer<D>, data: D) {
    for (i in 0 until size) {
        // Cast to IrElementBase to avoid casting to interface and invokeinterface, both of which are slow.
        (get(i) as IrElementBase?)?.let {
            @Suppress("UNCHECKED_CAST")
            set(i, it.transform(transformer, data) as T)
        }
    }
}

/**
 * Transforms a mutable list in place.
 * Each element `it` is replaced with a result of `transformation(it)`,
 * `null` means "keep existing element" (to avoid creating excessive singleton lists).
 */
inline fun <T> MutableList<T>.transformFlat(transformation: (T) -> List<T>?) {
    var i = 0
    while (i < size) {
        val item = get(i)

        i = replaceInPlace(transformation(item), i)
    }
}

/**
 * Transforms a subset of a mutable list in place.
 * Each element `it` that has a type S is replaced with a result of `transformation(it)`,
 * `null` means "keep existing element" (to avoid creating excessive singleton lists).
 */
inline fun <T, reified S : T> MutableList<T>.transformSubsetFlat(transformation: (S) -> List<S>?) {
    var i = 0
    while (i < size) {
        val item = get(i)


Review comment (performance):
Let's rewrite it a bit to reduce the bytecode size, which might be important because it's a hot spot.
```suggestion
    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)
        }
    }
```

Code after:
import org.jetbrains.kotlin.ir.declarations.IrDeclarationContainer
import org.jetbrains.kotlin.ir.visitors.IrElementTransformer

inline fun <reified T : IrElement> MutableList<T>.transformInPlace(transformation: (T) -> IrElement) {
    for (i in 0 until size) {
        set(i, transformation(get(i)) as T)
    }
}

fun <T : IrElement, D> MutableList<T>.transformInPlace(transformer: IrElementTransformer<D>, data: D) {
    for (i in 0 until size) {
        // Cast to IrElementBase to avoid casting to interface and invokeinterface, both of which are slow.
        @Suppress("UNCHECKED_CAST")
        set(i, (get(i) as IrElementBase).transform(transformer, data) as T)
    }
}

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)`,
 * `null` means "keep existing element" (to avoid creating excessive singleton lists).
 */
inline fun <T> MutableList<T>.transformFlat(transformation: (T) -> List<T>?) {
    var i = 0
    while (i < size) {
        val item = get(i)

        i = replaceInPlace(transformation(item), i)
    }
}

/**
 * Transforms a subset of a mutable list in place.
 * Each element `it` that has a type S is replaced with a result of `transformation(it)`,
 * `null` means "keep existing element" (to avoid creating excessive singleton lists).
 */
inline fun <T, reified S : T> MutableList<T>.transformSubsetFlat(transformation: (S) -> List<S>?) {
    var i = 0
    while (i < size) {
        val item = get(i)