FIR2IR: cache properties & fields more correctly

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrDeclarationStorage.kt

@@ -756,24 +756,16 @@ class Fir2IrDeclarationStorage(
         }
     }
 
-    fun getIrProperty(
+    private fun createIrProperty(
         property: FirProperty,
         irParent: IrDeclarationParent?,
         origin: IrDeclarationOrigin = IrDeclarationOrigin.DEFINED,
         shouldLeaveScope: Boolean = true
     ): IrProperty {
-        val cached = propertyCache[property]
-        if (cached != null) {
-            if (!shouldLeaveScope) {
-                enterScope(cached.descriptor)
-            }
-            return cached
-        }
-        return propertyCache.getOrPut(property) {
         val containerSource = property.containerSource
         val descriptor = containerSource?.let { WrappedPropertyDescriptorWithContainerSource(it) } ?: WrappedPropertyDescriptor()
         preCacheTypeParameters(property)
-            property.convertWithOffsets { startOffset, endOffset ->
+        return property.convertWithOffsets { startOffset, endOffset ->
             enterScope(descriptor)
             val result = irSymbolTable.declareProperty(
                 startOffset, endOffset,
@@ -821,17 +813,32 @@ class Fir2IrDeclarationStorage(
             if (shouldLeaveScope) {
                 leaveScope(descriptor)
             }
+            propertyCache[property] = result
             result
         }
     }
+
+    fun getIrProperty(
+        property: FirProperty,
+        irParent: IrDeclarationParent?,
+        origin: IrDeclarationOrigin = IrDeclarationOrigin.DEFINED,
+        shouldLeaveScope: Boolean = true
+    ): IrProperty {
+        val cached = propertyCache[property]
+        if (cached != null) {
+            if (!shouldLeaveScope) {
+                enterScope(cached.descriptor)
+            }
+            return cached
+        }
+        return createIrProperty(property, irParent, origin, shouldLeaveScope)
     }
 
-    private fun getIrField(field: FirField): IrField {
+    private fun createIrField(field: FirField): IrField {
-        return fieldCache.getOrPut(field) {
         val descriptor = WrappedFieldDescriptor()
         val origin = IrDeclarationOrigin.IR_EXTERNAL_JAVA_DECLARATION_STUB
         val type = field.returnTypeRef.toIrType()
-            field.convertWithOffsets { startOffset, endOffset ->
+        return field.convertWithOffsets { startOffset, endOffset ->
             irSymbolTable.declareField(
                 startOffset, endOffset,
                 origin, descriptor, type
@@ -845,7 +852,7 @@ class Fir2IrDeclarationStorage(
                     isFakeOverride = false
                 ).apply {
                     descriptor.bind(this)
-                    }
+                    fieldCache[field] = this
                 }
             }
         }
@@ -996,14 +1003,16 @@ class Fir2IrDeclarationStorage(
     fun getIrPropertyOrFieldSymbol(firVariableSymbol: FirVariableSymbol<*>): IrSymbol {
         return when (val fir = firVariableSymbol.fir) {
             is FirProperty -> {
+                propertyCache[fir]?.let { return irSymbolTable.referenceProperty(it.descriptor) }
                 val irParent = findIrParent(fir)
-                val irProperty = getIrProperty(fir, irParent).apply {
+                val irProperty = createIrProperty(fir, irParent).apply {
                     setAndModifyParent(irParent)
                 }
                 irSymbolTable.referenceProperty(irProperty.descriptor)
             }
             is FirField -> {
-                val irField = getIrField(fir).apply {
+                fieldCache[fir]?.let { return irSymbolTable.referenceField(it.descriptor) }
+                val irField = createIrField(fir).apply {
                     setAndModifyParent(findIrParent(fir))
                 }
                 irSymbolTable.referenceField(irField.descriptor)
@@ -1015,8 +1024,9 @@ class Fir2IrDeclarationStorage(
     fun getIrBackingFieldSymbol(firVariableSymbol: FirVariableSymbol<*>): IrSymbol {
         return when (val fir = firVariableSymbol.fir) {
             is FirProperty -> {
+                propertyCache[fir]?.let { return irSymbolTable.referenceField(it.backingField!!.descriptor) }
                 val irParent = findIrParent(fir)
-                val irProperty = getIrProperty(fir, irParent).apply {
+                val irProperty = createIrProperty(fir, irParent).apply {
                     setAndModifyParent(irParent)
                 }
                 irSymbolTable.referenceField(irProperty.backingField!!.descriptor)