diff --git a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/ElementUtils.kt b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/ElementUtils.kt
index 3f10078b590..c5f3dce14d1 100644
--- a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/ElementUtils.kt
+++ b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/ElementUtils.kt
@@ -1,20 +1,37 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
 package org.jetbrains.kotlin.fir.tree.generator.model
 
 import org.jetbrains.kotlin.generators.tree.*
+import org.jetbrains.kotlin.generators.tree.ElementOrRef as GenericElementOrRef
 
 // ----------- Simple field -----------
 
-fun field(name: String, type: TypeRefWithNullability, nullable: Boolean = false, withReplace: Boolean = false): Field {
-    return SimpleField(name, type.copy(nullable), withReplace = withReplace)
+fun field(
+    name: String,
+    type: TypeRefWithNullability,
+    nullable: Boolean = false,
+    withReplace: Boolean = false,
+    isChild: Boolean = true,
+): Field {
+    val isMutable = type is GenericElementOrRef<*> || withReplace
+    return SimpleField(name, type.copy(nullable), isChild = isChild, isMutable = isMutable, withReplace = withReplace)
 }
 
-fun field(type: ClassRef<*>, nullable: Boolean = false, withReplace: Boolean = false): Field {
-    return SimpleField(type.simpleName.replaceFirstChar(Char::lowercaseChar), type.copy(nullable), withReplace = withReplace)
+fun field(
+    type: ClassOrElementRef,
+    nullable: Boolean = false,
+    withReplace: Boolean = false,
+    isChild: Boolean = true,
+): Field {
+    val name = when (type) {
+        is ClassRef<*> -> type.simpleName
+        is GenericElementOrRef<*> -> type.element.name
+    }.replaceFirstChar(Char::lowercaseChar)
+    return field(name, type, nullable, withReplace, isChild)
 }
 
 fun booleanField(name: String, withReplace: Boolean = false): Field {
@@ -29,23 +46,24 @@ fun intField(name: String, withReplace: Boolean = false): Field {
     return field(name, StandardTypes.int, withReplace = withReplace)
 }
 
-// ----------- Fir field -----------
-
-fun field(name: String, element: ElementOrRef, nullable: Boolean = false, withReplace: Boolean = false, isChild: Boolean = true): Field {
-    return FirField(name, element.copy(nullable), withReplace = withReplace, isChild = isChild)
-}
-
-fun field(element: Element, nullable: Boolean = false, withReplace: Boolean = false, isChild: Boolean = true): Field {
-    return FirField(element.name.replaceFirstChar(Char::lowercaseChar), element.copy(nullable), withReplace = withReplace, isChild = isChild)
-}
-
 // ----------- Field list -----------
 
-fun fieldList(name: String, type: TypeRef, withReplace: Boolean = false, useMutableOrEmpty: Boolean = false, isChild: Boolean = true): Field {
+fun fieldList(
+    name: String,
+    type: TypeRef,
+    withReplace: Boolean = false,
+    useMutableOrEmpty: Boolean = false,
+    isChild: Boolean = true,
+): Field {
     return FieldList(name, type, withReplace = withReplace, isChild = isChild, useMutableOrEmpty = useMutableOrEmpty)
 }
 
-fun fieldList(elementOrRef: ElementOrRef, withReplace: Boolean = false, useMutableOrEmpty: Boolean = false, isChild: Boolean = true): Field {
+fun fieldList(
+    elementOrRef: ElementOrRef,
+    withReplace: Boolean = false,
+    useMutableOrEmpty: Boolean = false,
+    isChild: Boolean = true,
+): Field {
     val name = elementOrRef.element.name.replaceFirstChar(Char::lowercaseChar) + "s"
     return FieldList(name, elementOrRef, withReplace = withReplace, isChild = isChild, useMutableOrEmpty = useMutableOrEmpty)
 }
diff --git a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Field.kt b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Field.kt
index 312392e5ea1..9922b01b21f 100644
--- a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Field.kt
+++ b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Field.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -136,21 +136,21 @@ class FieldWithDefault(override val origin: Field) : Field(), AbstractFieldWithD
 class SimpleField(
     override val name: String,
     override val typeRef: TypeRefWithNullability,
+    override val isChild: Boolean,
+    override var isMutable: Boolean,
     override var withReplace: Boolean,
     override var isVolatile: Boolean = false,
     override var isFinal: Boolean = false,
     override var isLateinit: Boolean = false,
     override var isParameter: Boolean = false,
 ) : Field() {
-    override var isMutable: Boolean = withReplace
-
-    override val isChild: Boolean
-        get() = false
 
     override fun internalCopy(): Field {
         return SimpleField(
             name = name,
             typeRef = typeRef,
+            isChild = isChild,
+            isMutable = isMutable,
             withReplace = withReplace,
             isVolatile = isVolatile,
             isFinal = isFinal,
@@ -164,6 +164,8 @@ class SimpleField(
     override fun replaceType(newType: TypeRefWithNullability) = SimpleField(
         name = name,
         typeRef = newType,
+        isChild = isChild,
+        isMutable = isMutable,
         withReplace = withReplace,
         isVolatile = isVolatile,
         isFinal = isFinal,
@@ -174,35 +176,6 @@ class SimpleField(
         updateFieldsInCopy(it)
     }
 }
-
-class FirField(
-    override val name: String,
-    val element: ElementRef,
-    override var withReplace: Boolean,
-    override val isChild: Boolean,
-) : Field() {
-
-    override val typeRef: ElementRef
-        get() = element
-    override var isVolatile: Boolean = false
-    override var isFinal: Boolean = false
-
-    override var isMutable: Boolean = true
-    override var isLateinit: Boolean = false
-    override var isParameter: Boolean = false
-
-    override fun internalCopy(): Field {
-        return FirField(
-            name,
-            element,
-            withReplace,
-            isChild,
-        ).apply {
-            withBindThis = this@FirField.withBindThis
-        }
-    }
-}
-
 // ----------- Field list -----------
 
 class FieldList(
diff --git a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/ImplementationPrinter.kt b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/ImplementationPrinter.kt
index ec18dc42b4f..8248265c257 100644
--- a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/ImplementationPrinter.kt
+++ b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/ImplementationPrinter.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -51,15 +51,13 @@ internal class ImplementationPrinter(
             when (this) {
                 is FieldWithDefault -> origin.transform()
 
-                is FirField ->
+                is SimpleField ->
                     println("$name = ${name}${call()}transform(transformer, data)")
 
                 is FieldList -> {
                     addImport(transformInPlaceImport)
                     println("${name}.transformInplace(transformer, data)")
                 }
-
-                else -> throw IllegalStateException()
             }
         }
         with(implementation) {
@@ -138,7 +136,7 @@ internal class ImplementationPrinter(
                                         )
                                     } else {
                                         when (field.origin) {
-                                            is FirField -> {
+                                            is SimpleField -> {
                                                 println(field.acceptString())
                                             }
 
diff --git a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/TypeRef.kt b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/TypeRef.kt
index 8a730ea3fd5..5537d01aede 100644
--- a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/TypeRef.kt
+++ b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/TypeRef.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -128,7 +128,7 @@ data class TypeRefWithVariance<out T : TypeRef>(val variance: Variance, val type
         TypeRefWithVariance(variance, typeRef.substitute(map))
 }
 
-interface ElementOrRef<Element> : ParametrizedTypeRef<ElementOrRef<Element>, NamedTypeParameterRef>, ClassOrElementRef
+sealed interface ElementOrRef<Element> : ParametrizedTypeRef<ElementOrRef<Element>, NamedTypeParameterRef>, ClassOrElementRef
         where Element : AbstractElement<Element, *, *> {
     val element: Element