Refactor ObjCExportHeaderGenerator to allow per declaration translation

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/objcexport/ObjCExportCodeGenerator.kt

@@ -33,7 +33,7 @@ internal fun TypeBridge.makeNothing() = when (this) {
 
 internal class ObjCExportCodeGenerator(
         codegen: CodeGenerator,
-        val namer: ObjCExportNamerImpl,
+        val namer: ObjCExportNamer,
         val mapper: ObjCExportMapper
 ) : ObjCCodeGenerator(codegen) {
 
@@ -155,14 +155,15 @@ internal class ObjCExportCodeGenerator(
         return callFromBridge(conversion.owner.llvmFunction, listOf(value), resultLifetime)
     }
 
-    internal fun emitRtti(
+    private val objCTypeAdapters = mutableListOf<ObjCTypeAdapter>()
+
+    internal fun generate(
             generatedClasses: Collection<ClassDescriptor>,
+            categoryMembers: Map<ClassDescriptor, List<CallableMemberDescriptor>>,
             topLevel: Map<SourceFile, List<CallableMemberDescriptor>>
     ) {
-        val objCTypeAdapters = mutableListOf<ObjCTypeAdapter>()
-
         generatedClasses.forEach {
-            objCTypeAdapters += createTypeAdapter(it)
+            objCTypeAdapters += createTypeAdapter(it, categoryMembers[it].orEmpty())
 
             if (!it.isInterface) {
                 val className = namer.getClassOrProtocolName(it).binaryName
@@ -180,7 +181,9 @@ internal class ObjCExportCodeGenerator(
             val name = namer.getFileClassName(sourceFile).binaryName
             dataGenerator.emitEmptyClass(name, namer.kotlinAnyName.binaryName)
         }
+    }
 
+    internal fun emitRtti() {
         NSNumberKind.values().mapNotNull { it.mappedKotlinClassId }.forEach {
             dataGenerator.exportClass("Kotlin${it.shortClassName}")
         }
@@ -189,7 +192,7 @@ internal class ObjCExportCodeGenerator(
 
         emitSpecialClassesConvertions()
 
-        objCTypeAdapters += createTypeAdapter(context.builtIns.any)
+        objCTypeAdapters += createTypeAdapter(context.builtIns.any, categoryMembers = emptyList())
 
         val placedClassAdapters = mutableMapOf<String, ConstPointer>()
         val placedInterfaceAdapters = mutableMapOf<String, ConstPointer>()
@@ -397,7 +400,7 @@ private fun ObjCExportCodeGenerator.emitBoxConverter(
 private fun ObjCExportCodeGenerator.emitFunctionConverters() {
     val generator = BlockAdapterToFunctionGenerator(this)
 
-    (0 .. mapper.maxFunctionTypeParameterCount).forEach { numberOfParameters ->
+    (0 .. ObjCExportMapper.maxFunctionTypeParameterCount).forEach { numberOfParameters ->
         val converter = generator.run { generateConvertFunctionToBlock(numberOfParameters) }
         setObjCExportTypeInfo(context.builtIns.getFunction(numberOfParameters), constPointer(converter))
     }
@@ -421,7 +424,7 @@ private fun ObjCExportCodeGenerator.emitKotlinFunctionAdaptersToBlock() {
     val ptr = staticData.placeGlobalArray(
             "",
             pointerType(runtime.typeInfoType),
-            (0 .. mapper.maxFunctionTypeParameterCount).map {
+            (0 .. ObjCExportMapper.maxFunctionTypeParameterCount).map {
                 generateKotlinFunctionAdapterToBlock(it)
             }
     ).pointer.getElementPtr(0)
@@ -866,7 +869,8 @@ private fun ObjCExportCodeGenerator.createTypeAdapterForFileClass(
 }
 
 private fun ObjCExportCodeGenerator.createTypeAdapter(
-        descriptor: ClassDescriptor
+        descriptor: ClassDescriptor,
+        categoryMembers: List<CallableMemberDescriptor>
 ): ObjCExportCodeGenerator.ObjCTypeAdapter {
     val adapters = mutableListOf<ObjCExportCodeGenerator.ObjCToKotlinMethodAdapter>()
     val classAdapters = mutableListOf<ObjCExportCodeGenerator.ObjCToKotlinMethodAdapter>()
@@ -883,7 +887,7 @@ private fun ObjCExportCodeGenerator.createTypeAdapter(
         }
     }
 
-    val categoryMethods = mapper.getCategoryMembersFor(descriptor).toMethods()
+    val categoryMethods = categoryMembers.toMethods()
 
     val exposedMethods = descriptor.contributedMethods.filter { mapper.shouldBeExposed(it) } + categoryMethods
 

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/objcexport/CustomTypeMapper.kt

@@ -5,58 +5,117 @@
 
 package org.jetbrains.kotlin.backend.konan.objcexport
 
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.builtins.getReceiverTypeFromFunctionType
 import org.jetbrains.kotlin.builtins.getReturnTypeFromFunctionType
 import org.jetbrains.kotlin.builtins.getValueParameterTypesFromFunctionType
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.resolve.descriptorUtil.classId
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeUtils
 
 internal interface CustomTypeMapper {
     val mappedClassId: ClassId
-    fun mapType(mappedSuperType: KotlinType): ObjCNonNullReferenceType
+    fun mapType(mappedSuperType: KotlinType, translator: ObjCExportTranslatorImpl): ObjCNonNullReferenceType
+}
 
-    class Simple(
+internal object CustomTypeMappers {
-            override val mappedClassId: ClassId,
+    /**
-            private val objCClassName: String
+     * Custom type mappers.
-    ) : CustomTypeMapper {
+     *
+     * Don't forget to update [hiddenTypes] after adding new one.
+     */
+    val byClassId: Map<ClassId, CustomTypeMapper> = with(KotlinBuiltIns.FQ_NAMES) {
+        val result = mutableListOf<CustomTypeMapper>()
 
-        override fun mapType(mappedSuperType: KotlinType): ObjCNonNullReferenceType =
+        result += Collection(list, "NSArray")
-                ObjCClassType(objCClassName)
+        result += Collection(mutableList, "NSMutableArray")
+        result += Collection(set, "NSSet")
+        result += Collection(mutableSet, { namer.mutableSetName.objCName })
+        result += Collection(map, "NSDictionary")
+        result += Collection(mutableMap, { namer.mutableMapName.objCName })
+
+        NSNumberKind.values().forEach {
+            // TODO: NSNumber seem to have different equality semantics.
+            val classId = it.mappedKotlinClassId
+            if (classId != null) {
+                result += Simple(classId, { namer.numberBoxName(classId).objCName })
+            }
+
+        }
+
+        result += Simple(ClassId.topLevel(string.toSafe()), "NSString")
+
+        (0..ObjCExportMapper.maxFunctionTypeParameterCount).forEach {
+            result += Function(it)
+        }
+
+        result.associateBy { it.mappedClassId }
     }
 
-    class Collection(
+    /**
-            private val generator: ObjCExportHeaderGenerator,
+     * Types to be "hidden" during mapping, i.e. represented as `id`.
-            mappedClassDescriptor: ClassDescriptor,
+     *
-            private val objCClassName: String
+     * Currently contains super types of classes handled by custom type mappers.
+     * Note: can be generated programmatically, but requires stdlib in this case.
+     */
+    val hiddenTypes: Set<ClassId> = listOf(
+            "kotlin.Any",
+            "kotlin.CharSequence",
+            "kotlin.Comparable",
+            "kotlin.Function",
+            "kotlin.Number",
+            "kotlin.collections.Collection",
+            "kotlin.collections.Iterable",
+            "kotlin.collections.MutableCollection",
+            "kotlin.collections.MutableIterable"
+    ).map { ClassId.topLevel(FqName(it)) }.toSet()
+
+    private class Simple(
+            override val mappedClassId: ClassId,
+            private val getObjCClassName: ObjCExportTranslatorImpl.() -> String
     ) : CustomTypeMapper {
 
-        override val mappedClassId = mappedClassDescriptor.classId!!
+        constructor(
+                mappedClassId: ClassId,
+                objCClassName: String
+        ) : this(mappedClassId, { objCClassName })
 
-        override fun mapType(mappedSuperType: KotlinType): ObjCNonNullReferenceType {
+        override fun mapType(mappedSuperType: KotlinType, translator: ObjCExportTranslatorImpl): ObjCNonNullReferenceType =
+                ObjCClassType(translator.getObjCClassName())
+    }
+
+    private class Collection(
+            mappedClassFqName: FqName,
+            private val getObjCClassName: ObjCExportTranslatorImpl.() -> String
+    ) : CustomTypeMapper {
+
+        constructor(
+                mappedClassFqName: FqName,
+                objCClassName: String
+        ) : this(mappedClassFqName, { objCClassName })
+
+        override val mappedClassId = ClassId.topLevel(mappedClassFqName)
+
+        override fun mapType(mappedSuperType: KotlinType, translator: ObjCExportTranslatorImpl): ObjCNonNullReferenceType {
             val typeArguments = mappedSuperType.arguments.map {
                 val argument = it.type
                 if (TypeUtils.isNullableType(argument)) {
                     // Kotlin `null` keys and values are represented as `NSNull` singleton.
                     ObjCIdType
                 } else {
-                    generator.mapReferenceTypeIgnoringNullability(argument)
+                    translator.mapReferenceTypeIgnoringNullability(argument)
                 }
             }
 
-            return ObjCClassType(objCClassName, typeArguments)
+            return ObjCClassType(translator.getObjCClassName(), typeArguments)
         }
     }
 
-    class Function(
+    private class Function(parameterCount: Int) : CustomTypeMapper {
-            private val generator: ObjCExportHeaderGenerator,
+        override val mappedClassId: ClassId = KotlinBuiltIns.getFunctionClassId(parameterCount)
-            parameterCount: Int
-    ) : CustomTypeMapper {
-        override val mappedClassId: ClassId = generator.builtIns.getFunction(parameterCount).classId!!
 
-        override fun mapType(mappedSuperType: KotlinType): ObjCNonNullReferenceType {
+        override fun mapType(mappedSuperType: KotlinType, translator: ObjCExportTranslatorImpl): ObjCNonNullReferenceType {
             val functionType = mappedSuperType
 
             val returnType = functionType.getReturnTypeFromFunctionType()
@@ -64,8 +123,8 @@ internal interface CustomTypeMapper {
                     functionType.getValueParameterTypesFromFunctionType().map { it.type }
 
             return ObjCBlockPointerType(
-                    generator.mapReferenceType(returnType),
+                    translator.mapReferenceType(returnType),
-                    parameterTypes.map { generator.mapReferenceType(it) }
+                    parameterTypes.map { translator.mapReferenceType(it) }
             )
         }
     }

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/objcexport/ObjCExport.kt

@@ -6,13 +6,11 @@
 package org.jetbrains.kotlin.backend.konan.objcexport
 
 import org.jetbrains.kotlin.backend.konan.descriptors.getPackageFragments
+import org.jetbrains.kotlin.backend.konan.getExportedDependencies
 import org.jetbrains.kotlin.backend.konan.isNativeBinary
 import org.jetbrains.kotlin.backend.konan.llvm.CodeGenerator
 import org.jetbrains.kotlin.backend.konan.llvm.objcexport.ObjCExportCodeGenerator
-import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
-import org.jetbrains.kotlin.descriptors.SourceFile
 import org.jetbrains.kotlin.konan.file.File
 import org.jetbrains.kotlin.konan.target.CompilerOutputKind
 import org.jetbrains.kotlin.konan.target.KonanTarget
@@ -32,36 +30,33 @@ internal class ObjCExport(val codegen: CodeGenerator) {
 
         if (!context.config.produce.isNativeBinary) return // TODO: emit RTTI to the same modules as classes belong to.
 
-        val objCCodeGenerator: ObjCExportCodeGenerator
+        val produceFramework = context.config.produce == CompilerOutputKind.FRAMEWORK
-        val generatedClasses: Set<ClassDescriptor>
-        val topLevelDeclarations: Map<SourceFile, List<CallableMemberDescriptor>>
 
-        if (context.config.produce == CompilerOutputKind.FRAMEWORK) {
+        val mapper = ObjCExportMapper()
-            val headerGenerator = ObjCExportHeaderGeneratorImpl(context)
+        val exportedDependencies = if (produceFramework) context.getExportedDependencies() else emptyList()
+        val moduleDescriptors = listOf(context.moduleDescriptor) + exportedDependencies
+        val namer = ObjCExportNamerImpl(
+                moduleDescriptors.toSet(),
+                context.moduleDescriptor.builtIns,
+                mapper,
+                context.moduleDescriptor.namePrefix,
+                local = false
+        )
+
+        val objCCodeGenerator = ObjCExportCodeGenerator(codegen, namer, mapper)
+
+        if (produceFramework) {
+            val headerGenerator = ObjCExportHeaderGeneratorImpl(context, moduleDescriptors, mapper, namer)
             produceFrameworkSpecific(headerGenerator)
 
-            generatedClasses = headerGenerator.generatedClasses
+            objCCodeGenerator.generate(
-            topLevelDeclarations = headerGenerator.topLevel
+                    generatedClasses = headerGenerator.generatedClasses,
-            objCCodeGenerator = ObjCExportCodeGenerator(codegen, headerGenerator.namer, headerGenerator.mapper)
+                    categoryMembers = headerGenerator.extensions,
-        } else {
+                    topLevel = headerGenerator.topLevel
-            // TODO: refactor ObjCExport* to handle this case on a general basis.
+            )
-            val mapper = object : ObjCExportMapper() {
-                override fun getCategoryMembersFor(descriptor: ClassDescriptor): List<CallableMemberDescriptor> =
-                        emptyList()
-
-                override fun isSpecialMapped(descriptor: ClassDescriptor): Boolean =
-                        error("shouldn't reach here")
-
-            }
-
-            val namer = ObjCExportNamerImpl(emptySet(), context.builtIns, mapper, context.moduleDescriptor.namePrefix)
-            objCCodeGenerator = ObjCExportCodeGenerator(codegen, namer, mapper)
-
-            generatedClasses = emptySet()
-            topLevelDeclarations = emptyMap()
         }
 
-        objCCodeGenerator.emitRtti(generatedClasses = generatedClasses, topLevel = topLevelDeclarations)
+        objCCodeGenerator.emitRtti()
     }
 
     private fun produceFrameworkSpecific(headerGenerator: ObjCExportHeaderGenerator) {

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/objcexport/ObjCExportHeaderGeneratorImpl.kt

@@ -6,16 +6,17 @@
 package org.jetbrains.kotlin.backend.konan.objcexport
 
 import org.jetbrains.kotlin.backend.konan.Context
-import org.jetbrains.kotlin.backend.konan.getExportedDependencies
 import org.jetbrains.kotlin.backend.konan.reportCompilationWarning
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.ir.util.report
 
-internal class ObjCExportHeaderGeneratorImpl(val context: Context) : ObjCExportHeaderGenerator(
+internal class ObjCExportHeaderGeneratorImpl(
-        context.moduleDescriptor,
+        val context: Context,
-        context.getExportedDependencies(),
+        moduleDescriptors: List<ModuleDescriptor>,
-        context.builtIns
+        mapper: ObjCExportMapper,
-) {
+        namer: ObjCExportNamer
+) : ObjCExportHeaderGenerator(moduleDescriptors, context.builtIns, mapper, namer) {
 
     override fun reportWarning(text: String) {
         context.reportCompilationWarning(text)

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/objcexport/ObjCExportMapper.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.backend.konan.*
 import org.jetbrains.kotlin.backend.konan.descriptors.allOverriddenDescriptors
 import org.jetbrains.kotlin.backend.konan.descriptors.isArray
 import org.jetbrains.kotlin.backend.konan.descriptors.isInterface
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.builtins.PrimitiveType
 import org.jetbrains.kotlin.builtins.UnsignedType
 import org.jetbrains.kotlin.descriptors.*
@@ -21,10 +22,19 @@ import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.typeUtil.isNothing
 import org.jetbrains.kotlin.types.typeUtil.isUnit
 
-internal abstract class ObjCExportMapper {
+internal class ObjCExportMapper {
-    abstract fun getCategoryMembersFor(descriptor: ClassDescriptor): List<CallableMemberDescriptor>
+    companion object {
-    val maxFunctionTypeParameterCount get() = KONAN_FUNCTION_INTERFACES_MAX_PARAMETERS
+        val maxFunctionTypeParameterCount get() = KONAN_FUNCTION_INTERFACES_MAX_PARAMETERS
-    abstract fun isSpecialMapped(descriptor: ClassDescriptor): Boolean
+    }
+
+    val customTypeMappers: Map<ClassId, CustomTypeMapper> get() = CustomTypeMappers.byClassId
+    val hiddenTypes: Set<ClassId> get() = CustomTypeMappers.hiddenTypes
+
+    fun isSpecialMapped(descriptor: ClassDescriptor): Boolean {
+        // TODO: this method duplicates some of the [ObjCExportTranslatorImpl.mapReferenceType] logic.
+        return KotlinBuiltIns.isAny(descriptor) ||
+                descriptor.getAllSuperClassifiers().any { it.classId in customTypeMappers }
+    }
 
     private val methodBridgeCache = mutableMapOf<FunctionDescriptor, MethodBridge>()
 

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/objcexport/ObjCExportNamer.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.backend.konan.objcexport
 
+import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.backend.konan.descriptors.isArray
 import org.jetbrains.kotlin.backend.konan.descriptors.isInterface
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
@@ -14,11 +15,11 @@ import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.load.kotlin.PackagePartClassUtils
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.resolve.descriptorUtil.*
+import org.jetbrains.kotlin.resolve.descriptorUtil.isSubclassOf
+import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.source.PsiSourceFile
 
-
 interface ObjCExportNamer {
     data class ClassOrProtocolName(val swiftName: String, val objCName: String, val binaryName: String = objCName)
 
@@ -29,32 +30,37 @@ interface ObjCExportNamer {
     fun getPropertyName(property: PropertyDescriptor): String
     fun getObjectInstanceSelector(descriptor: ClassDescriptor): String
     fun getEnumEntrySelector(descriptor: ClassDescriptor): String
+
+    fun numberBoxName(classId: ClassId): ClassOrProtocolName
+
+    val kotlinAnyName: ClassOrProtocolName
+    val mutableSetName: ClassOrProtocolName
+    val mutableMapName: ClassOrProtocolName
+    val kotlinNumberName: ClassOrProtocolName
 }
 
 fun createNamer(moduleDescriptor: ModuleDescriptor,
                 topLevelNamePrefix: String = moduleDescriptor.namePrefix): ObjCExportNamer =
         createNamer(moduleDescriptor, emptyList(), topLevelNamePrefix)
 
-fun createNamer(moduleDescriptor: ModuleDescriptor,
+fun createNamer(
-                exportedDependencies: List<ModuleDescriptor>,
+        moduleDescriptor: ModuleDescriptor,
-                topLevelNamePrefix: String = moduleDescriptor.namePrefix): ObjCExportNamer {
+        exportedDependencies: List<ModuleDescriptor>,
-    val generator = object : ObjCExportHeaderGenerator(
+        topLevelNamePrefix: String = moduleDescriptor.namePrefix
-            moduleDescriptor,
+): ObjCExportNamer = ObjCExportNamerImpl(
-            exportedDependencies,
+        (exportedDependencies + moduleDescriptor).toSet(),
-            moduleDescriptor.builtIns,
+        moduleDescriptor.builtIns,
-            topLevelNamePrefix
+        ObjCExportMapper(),
-    ) {
+        topLevelNamePrefix,
-        override fun reportWarning(text: String) {}
+        local = true
-        override fun reportWarning(method: FunctionDescriptor, text: String) {}
+)
-    }
-    return generator.namer
-}
 
 internal class ObjCExportNamerImpl(
-        val moduleDescriptors: Set<ModuleDescriptor>,
+        private val moduleDescriptors: Set<ModuleDescriptor>,
         builtIns: KotlinBuiltIns,
-        val mapper: ObjCExportMapper,
+        private val mapper: ObjCExportMapper,
-        private val topLevelNamePrefix: String
+        private val topLevelNamePrefix: String,
+        private val local: Boolean
 ) : ObjCExportNamer {
 
     private fun String.toUnmangledClassOrProtocolName(): ObjCExportNamer.ClassOrProtocolName =
@@ -66,15 +72,15 @@ internal class ObjCExportNamerImpl(
             binaryName = "Kotlin$this"
     )
 
-    val kotlinAnyName = "KotlinBase".toUnmangledClassOrProtocolName()
+    override val kotlinAnyName = "KotlinBase".toUnmangledClassOrProtocolName()
 
-    val mutableSetName = "MutableSet".toSpecialStandardClassOrProtocolName()
+    override val mutableSetName = "MutableSet".toSpecialStandardClassOrProtocolName()
-    val mutableMapName = "MutableDictionary".toSpecialStandardClassOrProtocolName()
+    override val mutableMapName = "MutableDictionary".toSpecialStandardClassOrProtocolName()
 
-    fun numberBoxName(classId: ClassId): ObjCExportNamer.ClassOrProtocolName =
+    override fun numberBoxName(classId: ClassId): ObjCExportNamer.ClassOrProtocolName =
             classId.shortClassName.asString().toSpecialStandardClassOrProtocolName()
 
-    val kotlinNumberName = "Number".toSpecialStandardClassOrProtocolName()
+    override val kotlinNumberName = "Number".toSpecialStandardClassOrProtocolName()
 
     private val methodSelectors = object : Mapping<FunctionDescriptor, String>() {
 
@@ -204,8 +210,10 @@ internal class ObjCExportNamerImpl(
                 } else {
                     append(descriptor.name.asString().capitalize())
                 }
-            } else {
+            } else if (containingDeclaration is PackageFragmentDescriptor) {
                 appendTopLevelClassBaseName(descriptor)
+            } else {
+                error("unexpected class parent: $containingDeclaration")
             }
         }.mangledBySuffixUnderscores()
     }
@@ -219,8 +227,10 @@ internal class ObjCExportNamerImpl(
                     append(getClassOrProtocolObjCName(containingDeclaration))
                             .append(descriptor.name.asString().capitalize())
 
-                } else {
+                } else if (containingDeclaration is PackageFragmentDescriptor) {
                     append(topLevelNamePrefix).appendTopLevelClassBaseName(descriptor)
+                } else {
+                    error("unexpected class parent: $containingDeclaration")
                 }
             }.mangledBySuffixUnderscores()
         }
@@ -430,21 +440,22 @@ internal class ObjCExportNamerImpl(
             error("name candidates run out")
         }
 
-        fun getIfAssigned(element: T): N? = elementToName[element]
+        private fun getIfAssigned(element: T): N? = elementToName[element]
 
-        fun tryAssign(element: T, name: N): Boolean {
+        private fun tryAssign(element: T, name: N): Boolean {
             if (element in elementToName) error(element)
 
             if (reserved(name)) return false
 
-            val elements = nameToElements.getOrPut(name) { mutableListOf() }
+            if (nameToElements[name].orEmpty().any { conflict(element, it) }) {
-            if (elements.any { conflict(element, it) }) {
                 return false
             }
 
-            elements += element
+            if (!local) {
+                nameToElements.getOrPut(name) { mutableListOf() } += element
 
-            elementToName[element] = name
+                elementToName[element] = name
+            }
 
             return true
         }