[FIR/IR generator] Introduce ImportCollector

This class enables printing the import list in generated files in a smarter way. Also, refactor `Importable` interface hierarchy, namely, don't inherit `TypeRef` from `Importable`, since we have types like `TypeRef.Star` which are not really importable. Replace the `Importable#typeWithArguments` property with the `TypeRef#render` method to utilize `ImportCollector` while rendering types.
2023-10-16 14:57:22 +02:00
parent bec6d38490
commit 7b7bcb8ffa
451 changed files with 1592 additions and 3017 deletions
@@ -5,8 +5,6 @@

 package org.jetbrains.kotlin.generators.tree

-import org.jetbrains.kotlin.generators.tree.printer.generics
-
 /**
 * A class representing a FIR or IR tree element.
 */
@@ -43,8 +41,14 @@ abstract class AbstractElement<Element, Field> : ElementOrRef<Element, Field>, F
    override val allParents: List<Element>
        get() = elementParents.map { it.element }

-    final override val typeWithArguments: String
-        get() = type + generics
+    context(ImportCollector)
+    final override fun renderTo(appendable: Appendable) {
+        addImport(this)
+        appendable.append(typeName)
+        if (params.isNotEmpty()) {
+            params.joinTo(appendable, prefix = "<", postfix = ">") { it.name }
+        }
+    }

    abstract override val allFields: List<Field>

@@ -66,4 +70,6 @@ abstract class AbstractElement<Element, Field> : ElementOrRef<Element, Field>, F

    @Suppress("UNCHECKED_CAST")
    override fun substitute(map: TypeParameterSubstitutionMap): Element = this as Element
+
+    fun withStarArgs(): ElementRef<Element, Field> = copy(params.associateWith { TypeRef.Star })
 }
@@ -26,7 +26,7 @@ abstract class AbstractField {

    open val arbitraryImportables: MutableList<Importable> = mutableListOf()

-    open var optInAnnotation: ArbitraryImportable? = null
+    open var optInAnnotation: ClassRef<*>? = null

    abstract val isMutable: Boolean
    open val withGetter: Boolean get() = false
@@ -5,7 +5,4 @@

 package org.jetbrains.kotlin.generators.tree

-data class ArbitraryImportable(override val packageName: String, override val type: String) : Importable {
-    override val typeWithArguments: String
-        get() = type
-}
+data class ArbitraryImportable(override val packageName: String, override val typeName: String) : Importable
@@ -5,7 +5,7 @@

 package org.jetbrains.kotlin.generators.tree

-interface ImplementationKindOwner : Importable {
+interface ImplementationKindOwner : TypeRef, Importable {
    var kind: ImplementationKind?
    val allParents: List<ImplementationKindOwner>
 }
@@ -0,0 +1,82 @@
+/*
+ * Copyright 2010-2023 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.generators.tree
+
+import java.util.SortedMap
+import java.util.SortedSet
+
+/**
+ * Used to collect [TypeRef]s for printing a list of imports for those [TypeRef]s.
+ */
+class ImportCollector(currentPackage: String) {
+
+    companion object {
+        private val STAR = sortedSetOf("*")
+
+        /**
+         * The maximum number of imports from a single package before collapsing those imports to a star-import.
+         */
+        private const val STAR_COLLAPSE_THRESHOLD = 4
+    }
+
+    /**
+     * A map of package names to a list of entities to import from that package.
+     */
+    private val imports: SortedMap<String, SortedSet<String>> = sortedMapOf()
+
+    /**
+     * Entities from these packages will not be imported explicitly.
+     *
+     * See [the list of default imports](https://kotlinlang.org/docs/packages.html#default-imports).
+     */
+    private val ignoredPackages = hashSetOf(
+        currentPackage,
+        "kotlin",
+        "kotlin.annotation",
+        "kotlin.collections",
+        "kotlin.comparisons",
+        "kotlin.io",
+        "kotlin.ranges",
+        "kotlin.sequences",
+        "kotlin.text",
+        "java.lang",
+        "kotlin.jvm",
+    )
+
+    private fun addImport(packageName: String, entity: String) {
+        if (packageName in ignoredPackages) return
+        val entities = imports.computeIfAbsent(packageName) { sortedSetOf() }
+        if (entities === STAR) return
+        if (entity == "*") {
+            imports[packageName] = STAR
+            return
+        }
+        entities.add(entity)
+        if (entities.size > STAR_COLLAPSE_THRESHOLD) {
+            imports[packageName] = STAR
+        }
+    }
+
+    fun addImport(importable: Importable) {
+        addImport(importable.packageName, importable.typeName)
+    }
+
+    fun addStarImport(packageName: String) {
+        addImport(packageName, "*")
+    }
+
+    fun printAllImports(printer: Appendable) {
+        for ((packageName, entities) in imports) {
+            for (entity in entities) {
+                printer.append("import ", packageName, ".", entity, "\n")
+            }
+        }
+    }
+
+    fun addAllImports(importables: Collection<Importable>) {
+        importables.forEach(this::addImport)
+    }
+}
@@ -6,9 +6,11 @@
 package org.jetbrains.kotlin.generators.tree

 interface Importable {
-    val type: String
-    val packageName: String?
-    val fullQualifiedName: String? get() = packageName?.let { "$it.${type.replace(Regex("<.+>"), "")}" }

-    val typeWithArguments: String
+    /**
+     * The name of the entity to import.
+     */
+    val typeName: String
+
+    val packageName: String
 }
@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.types.Variance
 import java.util.*
 import kotlin.reflect.KClass

-interface TypeRef : Importable {
+interface TypeRef {

    /**
     * Constructs a new [TypeRef] by recursively replacing referenced [TypeParameterRef]s with other types according to the provided
@@ -18,23 +18,34 @@ interface TypeRef : Importable {
     */
    fun substitute(map: TypeParameterSubstitutionMap): TypeRef

+    /**
+     * Prints this type to [appendable] with all its arguments and question marks, while recursively collecting
+     * `this` and other referenced types into the import collector passed as context.
+     */
+    context(ImportCollector)
+    fun renderTo(appendable: Appendable)
+
    object Star : TypeRef {
-        override val type: String
-            get() = "*"
-
-        override val packageName: String?
-            get() = null
-
-        override val typeWithArguments: String
-            get() = type

        override fun substitute(map: TypeParameterSubstitutionMap) = this

-        override fun toString(): String = type
+        override fun toString(): String = "*"
+
+        context(ImportCollector)
+        override fun renderTo(appendable: Appendable) {
+            appendable.append(toString())
+        }
    }
 }

-sealed interface ClassOrElementRef : TypeRefWithNullability
+/**
+ * Prints this type as a string with all its arguments and question marks, while recursively collecting
+ * `this` and other referenced types into the import collector passed as context.
+ */
+context(ImportCollector)
+fun TypeRef.render(): String = buildString { renderTo(this) }
+
+sealed interface ClassOrElementRef : TypeRefWithNullability, Importable

 // Based on com.squareup.kotlinpoet.ClassName
 class ClassRef<P : TypeParameterRef> private constructor(
@@ -59,21 +70,25 @@ class ClassRef<P : TypeParameterRef> private constructor(
    private val names = Collections.unmodifiableList(names)

    /** Fully qualified name using `.` as a separator, like `kotlin.collections.Map.Entry`. */
-    val canonicalName: String = if (names[0].isEmpty())
-        names.subList(1, names.size).joinToString(".") else
-        names.joinToString(".")
+    val canonicalName: String = if (names[0].isEmpty()) typeName else names.joinToString(".")

    /** Package name, like `"kotlin.collections"` for `Map.Entry`. */
-    override val packageName: String get() = names[0]
+    override val packageName: String
+        get() = names[0]

    /** Simple name of this class, like `"Entry"` for `Map.Entry`. */
    val simpleName: String get() = names[names.size - 1]

-    override val type: String
-        get() = simpleName
+    override val typeName: String
+        get() = simpleNames.joinToString(separator = ".")

-    override val fullQualifiedName: String
-        get() = canonicalName
+    context(ImportCollector)
+    override fun renderTo(appendable: Appendable) {
+        addImport(this)
+        simpleNames.joinTo(appendable, separator = ".")
+        renderArgsTo(appendable)
+        renderNullabilityTo(appendable)
+    }

    /**
     * The enclosing classes, outermost first, followed by the simple name. This is `["Map", "Entry"]`
@@ -99,20 +114,14 @@ class ClassRef<P : TypeParameterRef> private constructor(
 */
 data class TypeRefWithVariance<out T : TypeRef>(val variance: Variance, val typeRef: T) : TypeRef {

-    override val type: String
-        get() = buildString {
-            if (variance != Variance.INVARIANT) {
-                append(variance)
-                append(" ")
-            }
-            append(typeRef.typeWithArguments)
+    context(ImportCollector)
+    override fun renderTo(appendable: Appendable) {
+        if (variance != Variance.INVARIANT) {
+            appendable.append(variance.label)
+            appendable.append(' ')
        }
-
-    override val packageName: String?
-        get() = null
-
-    override val typeWithArguments: String
-        get() = type
+        typeRef.renderTo(appendable)
+    }

    override fun substitute(map: TypeParameterSubstitutionMap): TypeRefWithVariance<*> =
        TypeRefWithVariance(variance, typeRef.substitute(map))
@@ -134,13 +143,22 @@ data class ElementRef<Element : AbstractElement<Element, Field>, Field : Abstrac
    override fun copy(args: Map<NamedTypeParameterRef, TypeRef>) = ElementRef(element, args, nullable)
    override fun copy(nullable: Boolean) = ElementRef(element, args, nullable)

-    override val type: String
-        get() = element.type
-    override val packageName: String?
+    override val typeName: String
+        get() = element.typeName
+
+    override val packageName: String
        get() = element.packageName

+    context(ImportCollector)
+    override fun renderTo(appendable: Appendable) {
+        addImport(element)
+        appendable.append(element.typeName)
+        renderArgsTo(appendable)
+        renderNullabilityTo(appendable)
+    }
+
    override fun toString() = buildString {
-        append(element.name)
+        append(element.typeName)
        append("<")
        append(args)
        append(">")
@@ -161,14 +179,10 @@ data class PositionTypeParameterRef(
 ) : TypeParameterRef {
    override fun toString() = index.toString()

-    override val type: String
-        get() = error("Getting type from ${this::class.simpleName} is not supported")
-
-    override val packageName: String?
-        get() = null
-
-    override val typeWithArguments: String
-        get() = type
+    context(ImportCollector)
+    override fun renderTo(appendable: Appendable) {
+        renderingIsNotSupported()
+    }

    override fun copy(nullable: Boolean) = PositionTypeParameterRef(index, nullable)
 }
@@ -187,13 +201,11 @@ open class NamedTypeParameterRef(

    override fun toString() = name

-    override val type: String
-        get() = name
-
-    override val packageName: String?
-        get() = null
-
-    override val typeWithArguments get() = name
+    context(ImportCollector)
+    override fun renderTo(appendable: Appendable) {
+        appendable.append(name)
+        renderNullabilityTo(appendable)
+    }

    final override fun copy(nullable: Boolean) = NamedTypeParameterRef(name, nullable)
 }
@@ -204,6 +216,12 @@ interface TypeRefWithNullability : TypeRef {
    fun copy(nullable: Boolean): TypeRefWithNullability
 }

+private fun TypeRefWithNullability.renderNullabilityTo(appendable: Appendable) {
+    if (nullable) {
+        appendable.append('?')
+    }
+}
+
 interface ParametrizedTypeRef<Self : ParametrizedTypeRef<Self, P>, P : TypeParameterRef> : TypeRef {
    val args: Map<P, TypeRef>

@@ -211,16 +229,20 @@ interface ParametrizedTypeRef<Self : ParametrizedTypeRef<Self, P>, P : TypeParam

    override fun substitute(map: TypeParameterSubstitutionMap): Self =
        copy(args.mapValues { it.value.substitute(map) })
+}

-    override val typeWithArguments: String
-        get() = type + generics + (if (nullable) "?" else "")
+context(ImportCollector)
+private fun ParametrizedTypeRef<*, *>.renderArgsTo(appendable: Appendable) {
+    if (args.isNotEmpty()) {
+        args.values.joinTo(appendable, prefix = "<", postfix = ">") {
+            it.renderTo(appendable)
+            ""
+        }
+    }
 }

 typealias TypeParameterSubstitutionMap = Map<out TypeParameterRef, TypeRef>

-val ParametrizedTypeRef<*, *>.generics: String
-    get() = if (args.isEmpty()) "" else args.values.joinToString(prefix = "<", postfix = ">") { it.typeWithArguments }
-
 fun <Self : ParametrizedTypeRef<Self, NamedTypeParameterRef>> ParametrizedTypeRef<Self, NamedTypeParameterRef>.withArgs(
    vararg args: Pair<String, TypeRef>
 ) = copy(args.associate { (k, v) -> NamedTypeParameterRef(k) to v })
@@ -270,3 +292,5 @@ fun ClassOrElementRef.inheritanceClauseParenthesis(): String = when (this) {

 val TypeRef.nullable: Boolean
    get() = (this as? TypeRefWithNullability)?.nullable ?: false
+
+fun TypeRef.renderingIsNotSupported(): Nothing = error("Rendering is not supported for ${this::class.simpleName}")
@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.generators.tree.printer

+import org.jetbrains.kotlin.generators.tree.ImportCollector
 import org.jetbrains.kotlin.utils.SmartPrinter
 import java.io.File

@@ -23,11 +24,12 @@ fun printGeneratedType(
    packageName: String,
    typeName: String,
    fileSuppressions: List<String> = emptyList(),
-    body: SmartPrinter.() -> Unit,
+    body: context(ImportCollector) SmartPrinter.() -> Unit,
 ): GeneratedFile {
    val stringBuilder = StringBuilder()
    val file = getPathForFile(generationPath, packageName, typeName)
-    SmartPrinter(stringBuilder).body()
+    val importCollector = ImportCollector(packageName)
+    body(importCollector, SmartPrinter(stringBuilder))
    return GeneratedFile(
        file,
        buildString {
@@ -43,6 +45,7 @@ fun printGeneratedType(
            }
            appendLine("package $packageName")
            appendLine()
+            importCollector.printAllImports(this)
            append(stringBuilder)
        }
    )
@@ -7,6 +7,8 @@ package org.jetbrains.kotlin.generators.tree.printer

 import org.jetbrains.kotlin.generators.tree.AbstractElement
 import org.jetbrains.kotlin.generators.tree.ImplementationKind
+import org.jetbrains.kotlin.generators.tree.ImportCollector
+import org.jetbrains.kotlin.generators.tree.render
 import org.jetbrains.kotlin.utils.SmartPrinter

 /**
@@ -17,9 +19,10 @@ import org.jetbrains.kotlin.utils.SmartPrinter
 *
 * @param end The string to add after the closing angle bracket of the type parameter list
 */
+context(ImportCollector)
 fun AbstractElement<*, *>.typeParameters(end: String = ""): String = params.takeIf { it.isNotEmpty() }
    ?.joinToString(", ", "<", ">$end") { param ->
-        param.name + (param.bounds.singleOrNull()?.let { " : ${it.typeWithArguments}" } ?: "")
+        param.name + (param.bounds.singleOrNull()?.let { " : ${it.render()}" } ?: "")
    } ?: ""

 /**
@@ -28,12 +31,13 @@ fun AbstractElement<*, *>.typeParameters(end: String = ""): String = params.take
 *
 * Otherwise, an empty string.
 */
+context(ImportCollector)
 fun AbstractElement<*, *>.multipleUpperBoundsList(): String {
    val paramsWithMultipleUpperBounds = params.filter { it.bounds.size > 1 }.takeIf { it.isNotEmpty() } ?: return ""
    return buildString {
        append(" where ")
        paramsWithMultipleUpperBounds.joinTo(this, separator = ", ") { param ->
-            param.bounds.joinToString(", ") { bound -> "$param : ${bound.typeWithArguments}" }
+            param.bounds.joinToString(", ") { bound -> "$param : ${bound.render()}" }
        }
        append("")
    }
@@ -48,11 +52,6 @@ fun ImplementationKind?.braces(): String = when (this) {
    else -> throw IllegalStateException(this.toString())
 }

-val AbstractElement<*, *>.generics: String
-    get() = params.takeIf { it.isNotEmpty() }
-        ?.let { it.joinToString(", ", "<", ">") { it.name } }
-        ?: ""
-
 fun SmartPrinter.printKDoc(kDoc: String?) {
    if (kDoc == null) return
    println("/**")