[K/N] Add objcClassesIncludingCategories cinterop property

It allows to list Objective-C classes that should include
corresponding categories from the same file.
The current implementation is super-simple and slow, but
it is OK since it is not intended to be a general-purpose
solution for now.

kotlin-native/Interop/Indexer/src/main/kotlin/org/jetbrains/kotlin/native/interop/indexer/Indexer.kt

@@ -64,6 +64,7 @@ private class ObjCClassImpl(
     override val methods = mutableListOf<ObjCMethod>()
     override val properties = mutableListOf<ObjCProperty>()
     override var baseClass: ObjCClass? = null
+    override val includedCategories = mutableListOf<ObjCCategory>()
 }
 
 private class ObjCCategoryImpl(
@@ -392,10 +393,48 @@ public open class NativeIndexImpl(val library: NativeLibrary, val verbose: Boole
         return objCClassRegistry.getOrPut(cursor, {
             ObjCClassImpl(name, getLocation(cursor), isForwardDeclaration = false,
                     binaryName = getObjCBinaryName(cursor).takeIf { it != name })
-        }) {
+        }) { objcClass ->
-            addChildrenToObjCContainer(cursor, it)
+            addChildrenToObjCContainer(cursor, objcClass)
+            if (name in this.library.objCClassesIncludingCategories) {
+                // We don't include methods from categories to class during indexing
+                // because indexing does not care about how class is represented in Kotlin.
+                // Instead, it should be done during StubIR construction.
+                objcClass.includedCategories += collectClassCategories(cursor, name).mapNotNull { getObjCCategoryAt(it) }
+            }
         }
+    }
 
+    /**
+     * Find all categories for a class that is pointed by [classCursor] in the same file.
+     * NB: Current implementation is rather slow as it walks the whole translation unit.
+     */
+    private fun collectClassCategories(classCursor: CValue<CXCursor>, className: String): List<CValue<CXCursor>> {
+        assert(classCursor.kind == CXCursorKind.CXCursor_ObjCInterfaceDecl) { classCursor.kind }
+        val classFile = getContainingFile(classCursor)
+        val result = mutableListOf<CValue<CXCursor>>()
+        // Accessing the whole translation unit (TU) is overkill, but it is the simplest solution which is doable
+        // since we use this function for a narrow set of cases.
+        // Possible improvements:
+        // 1. Find/create a function that returns a file scope. `clang_findReferencesInFile` does not seem to work because for categories
+        // it returns `CXCursor_ObjCClassRef` (@interface >CLASS_REFERENCE<(CategoryName)) and there is no easy way to access category from
+        // there.
+        // 2. Extract categories collection into a separate TU pass and create Class -> [Category] mapping. This way we can avoid visiting
+        // TU for every class.
+        val translationUnit = clang_getCursorLexicalParent(classCursor)
+        visitChildren(translationUnit) { childCursor, _ ->
+            if (childCursor.kind == CXCursorKind.CXCursor_ObjCCategoryDecl) {
+                val categoryClassCursor = getObjCCategoryClassCursor(childCursor)
+                val categoryClassName = clang_getCursorDisplayName(categoryClassCursor).convertAndDispose()
+                if (className == categoryClassName) {
+                    val categoryFile = getContainingFile(childCursor)
+                    if (clang_File_isEqual(categoryFile, classFile) != 0) {
+                        result += childCursor
+                    }
+                }
+            }
+            CXChildVisitResult.CXChildVisit_Continue
+        }
+        return result
     }
 
     private fun getObjCProtocolAt(cursor: CValue<CXCursor>): ObjCProtocolImpl {

kotlin-native/Interop/Indexer/src/main/kotlin/org/jetbrains/kotlin/native/interop/indexer/NativeIndex.kt

@@ -100,16 +100,23 @@ data class CompilationWithPCH(
         get() = emptyList()
 }
 
-// TODO: Compilation hierarchy seems to require some refactoring.
+/**
-
+ *
-data class NativeLibrary(override val includes: List<IncludeInfo>,
+ *  @param objCClassesIncludingCategories Objective-C classes that should be merged with categories from the same file.
-                         override val additionalPreambleLines: List<String>,
+ *
-                         override val compilerArgs: List<String>,
+ * TODO: Compilation hierarchy seems to require some refactoring.
-                         val headerToIdMapper: HeaderToIdMapper,
+ */
-                         override val language: Language,
+data class NativeLibrary(
-                         val excludeSystemLibs: Boolean, // TODO: drop?
+        override val includes: List<IncludeInfo>,
-                         val headerExclusionPolicy: HeaderExclusionPolicy,
+        override val additionalPreambleLines: List<String>,
-                         val headerFilter: NativeLibraryHeaderFilter) : Compilation
+        override val compilerArgs: List<String>,
+        val headerToIdMapper: HeaderToIdMapper,
+        override val language: Language,
+        val excludeSystemLibs: Boolean, // TODO: drop?
+        val headerExclusionPolicy: HeaderExclusionPolicy,
+        val headerFilter: NativeLibraryHeaderFilter,
+        val objCClassesIncludingCategories: Set<String>,
+) : Compilation
 
 data class IndexerResult(val index: NativeIndex, val compilation: CompilationWithPCH)
 
@@ -270,6 +277,10 @@ data class ObjCProperty(val name: String, val getter: ObjCMethod, val setter: Ob
 abstract class ObjCClass(name: String) : ObjCClassOrProtocol(name) {
     abstract val binaryName: String?
     abstract val baseClass: ObjCClass?
+    /**
+     * Categories whose methods and properties should be generated as members of Kotlin class.
+     */
+    abstract val includedCategories: List<ObjCCategory>
 }
 abstract class ObjCProtocol(name: String) : ObjCClassOrProtocol(name)
 

kotlin-native/Interop/StubGenerator/src/main/kotlin/org/jetbrains/kotlin/native/interop/gen/jvm/main.kt

@@ -568,6 +568,7 @@ internal fun buildNativeLibrary(
 
     val headerExclusionPolicy = HeaderExclusionPolicyImpl(imports)
 
+    val objCClassesIncludingCategories = def.config.objcClassesIncludingCategories.toSet()
     return NativeLibrary(
             includes = includes,
             additionalPreambleLines = compilation.additionalPreambleLines,
@@ -576,7 +577,8 @@ internal fun buildNativeLibrary(
             language = compilation.language,
             excludeSystemLibs = excludeSystemLibs,
             headerExclusionPolicy = headerExclusionPolicy,
-            headerFilter = headerFilter
+            headerFilter = headerFilter,
+            objCClassesIncludingCategories = objCClassesIncludingCategories
     )
 }
 

kotlin-native/shared/src/main/kotlin/org/jetbrains/kotlin/konan/util/DefFile.kt

@@ -132,6 +132,9 @@ class DefFile(val file:File?, val config:DefFileConfig, val manifestAddendProper
             properties.getProperty("plugin")
         }
 
+        val objcClassesIncludingCategories by lazy {
+            properties.getSpaceSeparated("objcClassesIncludingCategories")
+        }
     }
 }