diff --git a/compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeAttributes.kt b/compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeAttributes.kt
index 81b224c8d62..9f43777682a 100644
--- a/compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeAttributes.kt
+++ b/compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeAttributes.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.types
 
+import org.jetbrains.kotlin.fir.util.ConeTypeRegistry
 import org.jetbrains.kotlin.util.AttributeArrayOwner
 import org.jetbrains.kotlin.util.TypeRegistry
 import org.jetbrains.kotlin.types.model.AnnotationMarker
@@ -36,7 +37,7 @@ abstract class ConeAttribute<T : ConeAttribute<T>> : AnnotationMarker {
 class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : AttributeArrayOwner<ConeAttribute<*>, ConeAttribute<*>>(),
     Iterable<ConeAttribute<*>> {
 
-    companion object : TypeRegistry<ConeAttribute<*>, ConeAttribute<*>>() {
+    companion object : ConeTypeRegistry<ConeAttribute<*>, ConeAttribute<*>>() {
         inline fun <reified T : ConeAttribute<T>> attributeAccessor(): ReadOnlyProperty<ConeAttributes, T?> {
             @Suppress("UNCHECKED_CAST")
             return generateNullableAccessor<ConeAttribute<*>, T>(T::class) as ReadOnlyProperty<ConeAttributes, T?>
diff --git a/compiler/fir/cones/src/org/jetbrains/kotlin/fir/util/ConeTypeRegistry.kt b/compiler/fir/cones/src/org/jetbrains/kotlin/fir/util/ConeTypeRegistry.kt
new file mode 100644
index 00000000000..194235f427c
--- /dev/null
+++ b/compiler/fir/cones/src/org/jetbrains/kotlin/fir/util/ConeTypeRegistry.kt
@@ -0,0 +1,19 @@
+/*
+ * Copyright 2010-2021 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.util
+
+import org.jetbrains.kotlin.util.TypeRegistry
+import java.util.concurrent.ConcurrentHashMap
+import kotlin.reflect.KClass
+
+abstract class ConeTypeRegistry<K : Any, V : Any> : TypeRegistry<K, V>() {
+    override fun <T : K> ConcurrentHashMap<KClass<out K>, Int>.customComputeIfAbsent(
+        kClass: KClass<T>,
+        compute: (KClass<out K>) -> Int
+    ): Int {
+        return this.computeIfAbsent(kClass, compute)
+    }
+}
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt
index 6a2c30f7368..8283546d468 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir
 
 import org.jetbrains.kotlin.fir.types.impl.*
+import org.jetbrains.kotlin.fir.util.ConeTypeRegistry
 import org.jetbrains.kotlin.util.ArrayMapAccessor
 import org.jetbrains.kotlin.util.ComponentArrayOwner
 import org.jetbrains.kotlin.util.NullableArrayMapAccessor
@@ -18,7 +19,7 @@ abstract class FirSession @PrivateSessionConstructor constructor(
     val sessionProvider: FirSessionProvider?,
     val kind: Kind
 ) : ComponentArrayOwner<FirSessionComponent, FirSessionComponent>() {
-    companion object : TypeRegistry<FirSessionComponent, FirSessionComponent>() {
+    companion object : ConeTypeRegistry<FirSessionComponent, FirSessionComponent>() {
         inline fun <reified T : FirSessionComponent> sessionComponentAccessor(): ArrayMapAccessor<FirSessionComponent, FirSessionComponent, T> {
             return generateAccessor(T::class)
         }
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationAttributes.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationAttributes.kt
index 794af7e9f91..a6d6f1819bc 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationAttributes.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationAttributes.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.declarations
 
+import org.jetbrains.kotlin.fir.util.ConeTypeRegistry
 import org.jetbrains.kotlin.util.ArrayMap
 import org.jetbrains.kotlin.util.AttributeArrayOwner
 import org.jetbrains.kotlin.util.NullableArrayMapAccessor
@@ -47,7 +48,7 @@ class FirDeclarationAttributes : AttributeArrayOwner<FirDeclarationDataKey, Any>
  *    object SomeKey : FirDeclarationDataKey()
  *    var FirDeclaration.someString: String? by FirDeclarationDataRegistry.data(SomeKey)
  */
-object FirDeclarationDataRegistry : TypeRegistry<FirDeclarationDataKey, Any>() {
+object FirDeclarationDataRegistry : ConeTypeRegistry<FirDeclarationDataKey, Any>() {
     fun <K : FirDeclarationDataKey> data(key: K): DeclarationDataAccessor {
         val kClass = key::class
         return DeclarationDataAccessor(generateAnyNullableAccessor(kClass), kClass)
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/FirExtensionService.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/FirExtensionService.kt
index a95689ea53f..826029af466 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/FirExtensionService.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/FirExtensionService.kt
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.extensions
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.NoMutableState
+import org.jetbrains.kotlin.fir.util.ConeTypeRegistry
 import org.jetbrains.kotlin.util.ArrayMapAccessor
 import org.jetbrains.kotlin.util.ComponentArrayOwner
 import org.jetbrains.kotlin.util.TypeRegistry
@@ -18,7 +19,7 @@ annotation class PluginServicesInitialization
 
 @NoMutableState
 class FirExtensionService(val session: FirSession) : ComponentArrayOwner<FirExtension, List<FirExtension>>(), FirSessionComponent {
-    companion object : TypeRegistry<FirExtension, List<FirExtension>>() {
+    companion object : ConeTypeRegistry<FirExtension, List<FirExtension>>() {
         inline fun <reified P : FirExtension, V : List<P>> registeredExtensions(): ArrayMapAccessor<FirExtension, List<FirExtension>, V> {
             @Suppress("UNCHECKED_CAST")
             return generateAccessor(P::class, default = emptyList<P>() as V)
diff --git a/compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/TestServices.kt b/compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/TestServices.kt
index 9af8e0901cd..38a90eb9506 100644
--- a/compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/TestServices.kt
+++ b/compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/TestServices.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.test.services
 
+import org.jetbrains.kotlin.fir.util.ConeTypeRegistry
 import org.jetbrains.kotlin.util.ArrayMapAccessor
 import org.jetbrains.kotlin.util.ComponentArrayOwner
 import org.jetbrains.kotlin.util.NullableArrayMapAccessor
@@ -28,7 +29,7 @@ class TestServices : ComponentArrayOwner<TestService, TestService>(){
     override val typeRegistry: TypeRegistry<TestService, TestService>
         get() = Companion
 
-    companion object : TypeRegistry<TestService, TestService>() {
+    companion object : ConeTypeRegistry<TestService, TestService>() {
         inline fun <reified T : TestService> testServiceAccessor(): ArrayMapAccessor<TestService, TestService, T> {
             return generateAccessor(T::class)
         }
diff --git a/core/compiler.common/src/org/jetbrains/kotlin/util/ArrayMapOwner.kt b/core/compiler.common/src/org/jetbrains/kotlin/util/ArrayMapOwner.kt
index a0a96a5d85b..33674768728 100644
--- a/core/compiler.common/src/org/jetbrains/kotlin/util/ArrayMapOwner.kt
+++ b/core/compiler.common/src/org/jetbrains/kotlin/util/ArrayMapOwner.kt
@@ -78,9 +78,14 @@ abstract class TypeRegistry<K : Any, V : Any> {
     }
 
     fun <T : K> getId(kClass: KClass<T>): Int {
-        return idPerType.computeIfAbsent(kClass) { idCounter.getAndIncrement() }
+        return idPerType.customComputeIfAbsent(kClass) { idCounter.getAndIncrement() }
     }
 
+    abstract fun <T : K> ConcurrentHashMap<KClass<out K>, Int>.customComputeIfAbsent(
+        kClass: KClass<T>,
+        compute: (KClass<out K>) -> Int
+    ): Int
+
     fun allValuesThreadUnsafeForRendering(): Map<KClass<out K>, Int> {
         return idPerType
     }
diff --git a/core/descriptors/src/org/jetbrains/kotlin/types/TypeAttributes.kt b/core/descriptors/src/org/jetbrains/kotlin/types/TypeAttributes.kt
index ffc80b685c3..ad406aa5dfa 100644
--- a/core/descriptors/src/org/jetbrains/kotlin/types/TypeAttributes.kt
+++ b/core/descriptors/src/org/jetbrains/kotlin/types/TypeAttributes.kt
@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.types.model.AnnotationMarker
 import org.jetbrains.kotlin.util.AttributeArrayOwner
 import org.jetbrains.kotlin.util.TypeRegistry
 import org.jetbrains.kotlin.utils.addIfNotNull
+import java.util.concurrent.ConcurrentHashMap
 import kotlin.properties.ReadOnlyProperty
 import kotlin.reflect.KClass
 
@@ -41,6 +42,15 @@ class TypeAttributes private constructor(attributes: List<TypeAttribute<*>>) : A
             return generateNullableAccessor<TypeAttribute<*>, T>(T::class) as ReadOnlyProperty<TypeAttributes, T?>
         }
 
+        override fun <T : TypeAttribute<*>> ConcurrentHashMap<KClass<out TypeAttribute<*>>, Int>.customComputeIfAbsent(
+            kClass: KClass<T>,
+            compute: (KClass<out TypeAttribute<*>>) -> Int
+        ): Int {
+            return this[kClass] ?: synchronized(this) {
+                this[kClass] ?: compute(kClass).also { this.putIfAbsent(kClass, it) }
+            }
+        }
+
         val Empty: TypeAttributes = TypeAttributes(listOf(AnnotationsTypeAttribute(Annotations.EMPTY)))
 
         fun create(attributes: List<TypeAttribute<*>>): TypeAttributes {