[IR] Add new visitor IrTypeVisitorVoid that can traverse IR types

#KT-57812
2023-04-19 15:15:59 +02:00
parent 55cac9caa6
commit 01749a910b
6 changed files with 232 additions and 11 deletions
@@ -24,7 +24,7 @@ abstract class IrConstantObject : IrConstantValue() {

    abstract val valueArguments: MutableList<IrConstantValue>

-    abstract val typeArguments: List<IrType>
+    abstract val typeArguments: MutableList<IrType>

    override fun <R, D> accept(visitor: IrElementVisitor<R, D>, data: D): R =
        visitor.visitConstantObject(this, data)
@@ -0,0 +1,124 @@
+/*
+ * 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.
+ */
+
+// This file was generated automatically. See compiler/ir/ir.tree/tree-generator/ReadMe.md.
+// DO NOT MODIFY IT MANUALLY.
+
+package org.jetbrains.kotlin.ir.visitors
+
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrField
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrLocalDelegatedProperty
+import org.jetbrains.kotlin.ir.declarations.IrScript
+import org.jetbrains.kotlin.ir.declarations.IrTypeAlias
+import org.jetbrains.kotlin.ir.declarations.IrTypeParameter
+import org.jetbrains.kotlin.ir.declarations.IrValueParameter
+import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.expressions.IrClassReference
+import org.jetbrains.kotlin.ir.expressions.IrConstantObject
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
+import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
+import org.jetbrains.kotlin.ir.expressions.IrVararg
+import org.jetbrains.kotlin.ir.types.IrType
+
+interface IrTypeTransformerVoid<in D> : IrElementTransformer<D> {
+    fun <Type : IrType?> transformType(
+        container: IrElement,
+        type: Type,
+        data: D,
+    ): Type
+
+    override fun visitValueParameter(declaration: IrValueParameter, data: D) = run {
+        declaration.varargElementType = transformType(declaration, declaration.varargElementType,
+                data)
+        declaration.type = transformType(declaration, declaration.type, data)
+        return@run super.visitValueParameter(declaration, data)
+    }
+
+    override fun visitClass(declaration: IrClass, data: D) = run {
+        declaration.valueClassRepresentation?.mapUnderlyingType {
+            transformType(declaration, it, data)
+        }
+        declaration.superTypes = declaration.superTypes.map { transformType(declaration, it, data) }
+        return@run super.visitClass(declaration, data)
+    }
+
+    override fun visitTypeParameter(declaration: IrTypeParameter, data: D) = run {
+        declaration.superTypes = declaration.superTypes.map { transformType(declaration, it, data) }
+        return@run super.visitTypeParameter(declaration, data)
+    }
+
+    override fun visitFunction(declaration: IrFunction, data: D) = run {
+        declaration.returnType = transformType(declaration, declaration.returnType, data)
+        return@run super.visitFunction(declaration, data)
+    }
+
+    override fun visitField(declaration: IrField, data: D) = run {
+        declaration.type = transformType(declaration, declaration.type, data)
+        return@run super.visitField(declaration, data)
+    }
+
+    override fun visitLocalDelegatedProperty(declaration: IrLocalDelegatedProperty,
+            data: D) = run {
+        declaration.type = transformType(declaration, declaration.type, data)
+        return@run super.visitLocalDelegatedProperty(declaration, data)
+    }
+
+    override fun visitScript(declaration: IrScript, data: D) = run {
+        declaration.baseClass = transformType(declaration, declaration.baseClass, data)
+        return@run super.visitScript(declaration, data)
+    }
+
+    override fun visitTypeAlias(declaration: IrTypeAlias, data: D) = run {
+        declaration.expandedType = transformType(declaration, declaration.expandedType, data)
+        return@run super.visitTypeAlias(declaration, data)
+    }
+
+    override fun visitVariable(declaration: IrVariable, data: D) = run {
+        declaration.type = transformType(declaration, declaration.type, data)
+        return@run super.visitVariable(declaration, data)
+    }
+
+    override fun visitExpression(expression: IrExpression, data: D) = run {
+        expression.type = transformType(expression, expression.type, data)
+        return@run super.visitExpression(expression, data)
+    }
+
+    override fun visitMemberAccess(expression: IrMemberAccessExpression<*>, data: D) =
+            run {
+        (0 until expression.typeArgumentsCount).forEach {
+            expression.getTypeArgument(it)?.let { type ->
+                expression.putTypeArgument(it, transformType(expression, type, data))
+            }
+        }
+        return@run super.visitMemberAccess(expression, data)
+    }
+
+    override fun visitClassReference(expression: IrClassReference, data: D) = run {
+        expression.classType = transformType(expression, expression.classType, data)
+        return@run super.visitClassReference(expression, data)
+    }
+
+    override fun visitConstantObject(expression: IrConstantObject, data: D) = run {
+        for (i in 0 until expression.typeArguments.size) {
+            expression.typeArguments[i] = transformType(expression, expression.typeArguments[i],
+                    data)
+        }
+        return@run super.visitConstantObject(expression, data)
+    }
+
+    override fun visitTypeOperator(expression: IrTypeOperatorCall, data: D) = run {
+        expression.typeOperand = transformType(expression, expression.typeOperand, data)
+        return@run super.visitTypeOperator(expression, data)
+    }
+
+    override fun visitVararg(expression: IrVararg, data: D) = run {
+        expression.varargElementType = transformType(expression, expression.varargElementType, data)
+        return@run super.visitVararg(expression, data)
+    }
+}
@@ -38,11 +38,12 @@ class IrConstantObjectImpl constructor(
    override val startOffset: Int,
    override val endOffset: Int,
    override var constructor: IrConstructorSymbol,
-    initArguments: List<IrConstantValue>,
-    override val typeArguments: List<IrType>,
+    initValueArguments: List<IrConstantValue>,
+    initTypeArguments: List<IrType>,
    override var type: IrType = constructor.owner.constructedClassType,
 ) : IrConstantObject() {
-    override val valueArguments = SmartList(initArguments)
+    override val valueArguments = SmartList(initValueArguments)
+    override val typeArguments = SmartList(initTypeArguments)

    override fun contentEquals(other: IrConstantValue): Boolean =
        other is IrConstantObject &&
@@ -880,7 +880,7 @@ object IrTree : AbstractTreeBuilder() {

        +field("constructor", constructorSymbolType)
        +listField("valueArguments", constantValue, mutability = List, isChild = true)
-        +listField("typeArguments", irTypeType)
+        +listField("typeArguments", irTypeType, mutability = List)
    }
    val constantArray: ElementConfig by element(Expression) {
        visitorParent = constantValue
@@ -9,10 +9,7 @@ import org.jetbrains.kotlin.generators.util.GeneratorsFileUtil
 import org.jetbrains.kotlin.generators.util.GeneratorsFileUtil.collectPreviouslyGeneratedFiles
 import org.jetbrains.kotlin.generators.util.GeneratorsFileUtil.removeExtraFilesFromPreviousGeneration
 import org.jetbrains.kotlin.ir.generator.model.config2model
-import org.jetbrains.kotlin.ir.generator.print.printElements
-import org.jetbrains.kotlin.ir.generator.print.printTransformer
-import org.jetbrains.kotlin.ir.generator.print.printVisitor
-import org.jetbrains.kotlin.ir.generator.print.printVisitorVoid
+import org.jetbrains.kotlin.ir.generator.print.*
 import java.io.File

 const val BASE_PACKAGE = "org.jetbrains.kotlin.ir"
@@ -31,6 +28,7 @@ fun main(args: Array<String>) {
        yield(printVisitor(generationPath, model))
        yield(printVisitorVoid(generationPath, model))
        yield(printTransformer(generationPath, model))
+        yield(printTypeVisitor(generationPath, model))
        // IrElementTransformerVoid is too random to autogenerate
    }.map {
        GeneratorsFileUtil.writeFileIfContentChanged(it.file, it.newText, logNotChanged = false)
@@ -7,15 +7,17 @@ package org.jetbrains.kotlin.ir.generator.print

 import com.squareup.kotlinpoet.*
 import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
+import org.jetbrains.kotlin.ir.generator.IrTree
 import org.jetbrains.kotlin.ir.generator.VISITOR_PACKAGE
-import org.jetbrains.kotlin.ir.generator.model.Element
-import org.jetbrains.kotlin.ir.generator.model.Model
+import org.jetbrains.kotlin.ir.generator.irTypeType
+import org.jetbrains.kotlin.ir.generator.model.*
 import org.jetbrains.kotlin.ir.generator.util.GeneratedFile
 import java.io.File

 private val visitorTypeName = ClassName(VISITOR_PACKAGE, "IrElementVisitor")
 private val visitorVoidTypeName = ClassName(VISITOR_PACKAGE, "IrElementVisitorVoid")
 private val transformerTypeName = ClassName(VISITOR_PACKAGE, "IrElementTransformer")
+private val typeTransformerVoidTypeName = ClassName(VISITOR_PACKAGE, "IrTypeTransformerVoid")

 fun printVisitor(generationPath: File, model: Model): GeneratedFile {
    val visitorType = TypeSpec.interfaceBuilder(visitorTypeName).apply {
@@ -112,3 +114,99 @@ fun printTransformer(generationPath: File, model: Model): GeneratedFile {

    return printTypeCommon(generationPath, transformerTypeName.packageName, visitorType)
 }
+
+fun printTypeVisitor(generationPath: File, model: Model): GeneratedFile {
+    val transformTypeFunName = "transformType"
+
+    fun FunSpec.Builder.addVisitTypeStatement(element: Element, field: Field) {
+        val visitorParam = element.visitorParam
+        val access = "$visitorParam.${field.name}"
+        when (field) {
+            is SingleField -> addStatement("$access = $transformTypeFunName($visitorParam, $access, data)")
+            is ListField -> {
+                if (field.mutable) {
+                    addStatement("$access = $access.map { $transformTypeFunName($visitorParam, it, data) }")
+                } else {
+                    beginControlFlow("for (i in 0 until $access.size)")
+                    addStatement("$access[i] = $transformTypeFunName($visitorParam, $access[i], data)")
+                    endControlFlow()
+                }
+            }
+        }
+    }
+
+    fun Element.getFieldsWithIrTypeType(insideParent: Boolean = false): List<Field> {
+        val parentsFields = elementParents.flatMap { it.element.getFieldsWithIrTypeType(insideParent = true) }
+        if (insideParent && this.visitorParent != null) {
+            return parentsFields
+        }
+
+        val irTypeFields = this.fields
+            .filter {
+                val type = when (it) {
+                    is SingleField -> it.type
+                    is ListField -> it.elementType
+                }
+                type.toString() == irTypeType.toString()
+            }
+
+        return irTypeFields + parentsFields
+    }
+
+    val visitorType = TypeSpec.interfaceBuilder(typeTransformerVoidTypeName).apply {
+        val d = TypeVariableName("D", KModifier.IN)
+        addTypeVariable(d)
+        addSuperinterface(transformerTypeName.parameterizedBy(d))
+
+        val abstractVisitFun = FunSpec.builder(transformTypeFunName).apply {
+            val poetNullableIrType = irTypeType.toPoet().copy(nullable = true)
+            val typeVariable = TypeVariableName("Type", poetNullableIrType)
+            addTypeVariable(typeVariable)
+            addParameter("container", model.rootElement.toPoet())
+            addParameter("type", typeVariable)
+            addParameter("data", d)
+            returns(typeVariable)
+        }
+        addFunction(abstractVisitFun.addModifiers(KModifier.ABSTRACT).build())
+
+        fun buildVisitFun(element: Element) = FunSpec.builder(element.visitFunName).apply {
+            addModifiers(KModifier.OVERRIDE)
+            addParameter(element.visitorParam, element.toPoetStarParameterized())
+            addParameter("data", d)
+        }
+
+        for (element in model.elements) {
+            val irTypeFields = element.getFieldsWithIrTypeType()
+            if (irTypeFields.isEmpty()) continue
+
+            element.visitorParent?.let { _ ->
+                addFunction(buildVisitFun(element).apply {
+                    // Note: using `run` here to infer return type automatically
+                    beginControlFlow("return run")
+
+                    val visitorParam = element.visitorParam
+                    when (element.name) {
+                        IrTree.memberAccessExpression.name -> {
+                            beginControlFlow("(0 until $visitorParam.typeArgumentsCount).forEach {")
+                            beginControlFlow("$visitorParam.getTypeArgument(it)?.let { type ->")
+                            addStatement("expression.putTypeArgument(it, $transformTypeFunName($visitorParam, type, data))")
+                            endControlFlow()
+                            endControlFlow()
+                        }
+                        IrTree.`class`.name -> {
+                            beginControlFlow("$visitorParam.valueClassRepresentation?.mapUnderlyingType {")
+                            addStatement("$transformTypeFunName($visitorParam, it, data)")
+                            endControlFlow()
+                            irTypeFields.forEach { addVisitTypeStatement(element, it) }
+                        }
+                        else -> irTypeFields.forEach { addVisitTypeStatement(element, it) }
+                    }
+                    addStatement("return@run super.${element.visitFunName}($visitorParam, data)")
+                    endControlFlow()
+                }.build())
+            }
+        }
+    }.build()
+
+    return printTypeCommon(generationPath, typeTransformerVoidTypeName.packageName, visitorType)
+}