[Assign plugin] Add a compiler plugin for overloading assign ('=') operator

2022-09-01 07:51:38 +02:00
parent 6052962f50
commit 09d6dfc8bf
62 changed files with 2822 additions and 6 deletions
@@ -0,0 +1,29 @@
+description = "Kotlin Assignment Compiler Plugin (K2)"
+
+plugins {
+    kotlin("jvm")
+    id("jps-compatible")
+}
+
+dependencies {
+    implementation(project(":kotlin-assignment-compiler-plugin.common"))
+
+    compileOnly(project(":compiler:fir:cones"))
+    compileOnly(project(":compiler:fir:tree"))
+    compileOnly(project(":compiler:fir:resolve"))
+    compileOnly(project(":compiler:fir:checkers"))
+    compileOnly(project(":compiler:ir.backend.common"))
+    compileOnly(project(":compiler:fir:entrypoint"))
+
+    compileOnly(intellijCore())
+    runtimeOnly(kotlinStdlib())
+}
+
+sourceSets {
+    "main" { projectDefault() }
+    "test" { none() }
+}
+
+runtimeJar()
+sourcesJar()
+javadocJar()
@@ -0,0 +1,51 @@
+/*
+ * Copyright 2010-2022 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.assignment.plugin.k2
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.caches.getValue
+import org.jetbrains.kotlin.fir.expressions.classId
+import org.jetbrains.kotlin.fir.extensions.FirExtensionSessionComponent
+import org.jetbrains.kotlin.fir.extensions.FirExtensionSessionComponent.Factory
+import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
+import org.jetbrains.kotlin.name.ClassId
+
+internal class FirAssignAnnotationMatchingService(
+    session: FirSession,
+    private val annotationClassIds: List<ClassId>
+) : FirExtensionSessionComponent(session) {
+
+    companion object {
+        fun getFactory(annotations: List<String>): Factory {
+            return Factory { session -> FirAssignAnnotationMatchingService(session, annotations.map { ClassId.fromString(it) }) }
+        }
+    }
+
+    private val cache: FirCache<FirRegularClassSymbol, Boolean, Nothing?> = session.firCachesFactory.createCache { symbol, _ ->
+        symbol.annotated()
+    }
+
+    fun isAnnotated(symbol: FirRegularClassSymbol?): Boolean {
+        if (symbol == null) {
+            return false
+        }
+        return cache.getValue(symbol)
+    }
+
+    private fun FirRegularClassSymbol.annotated(): Boolean {
+        if (this.annotations.any { it.classId in annotationClassIds }) return true
+        return resolvedSuperTypeRefs.any {
+            val symbol = it.type.fullyExpandedType(session).toRegularClassSymbol(session) ?: return@any false
+            symbol.annotations.any { it.classId in annotationClassIds }
+        }
+    }
+}
+
+internal val FirSession.annotationMatchingService: FirAssignAnnotationMatchingService by FirSession.sessionComponentAccessor()
@@ -0,0 +1,75 @@
+/*
+ * Copyright 2010-2022 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.assignment.plugin.k2
+
+import org.jetbrains.kotlin.KtFakeSourceElementKind
+import org.jetbrains.kotlin.assignment.plugin.AssignmentPluginNames.ASSIGN_METHOD
+import org.jetbrains.kotlin.fakeElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.builder.buildFunctionCall
+import org.jetbrains.kotlin.fir.expressions.builder.buildPropertyAccessExpression
+import org.jetbrains.kotlin.fir.extensions.FirAssignExpressionAltererExtension
+import org.jetbrains.kotlin.fir.references.builder.buildSimpleNamedReference
+import org.jetbrains.kotlin.fir.resolvedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirBackingFieldSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirFieldSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
+import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
+import org.jetbrains.kotlin.fir.types.upperBoundIfFlexible
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
+
+class FirAssignmentPluginAssignAltererExtension(
+    session: FirSession
+) : FirAssignExpressionAltererExtension(session) {
+
+    override fun transformVariableAssignment(variableAssignment: FirVariableAssignment): FirStatement? {
+        return runIf(variableAssignment.supportsTransformVariableAssignment()) {
+            buildFunctionCall(variableAssignment)
+        }
+    }
+
+    private fun FirVariableAssignment.supportsTransformVariableAssignment(): Boolean {
+        return when (val lSymbol = lValue.resolvedSymbol as? FirVariableSymbol<*>) {
+            is FirPropertySymbol -> lSymbol.isVal && !lSymbol.isLocal && lSymbol.hasSpecialAnnotation()
+            is FirBackingFieldSymbol -> lSymbol.isVal && lSymbol.hasSpecialAnnotation()
+            is FirFieldSymbol -> lSymbol.isVal && lSymbol.hasSpecialAnnotation()
+            else -> false
+        }
+    }
+
+    private fun FirVariableSymbol<*>.hasSpecialAnnotation(): Boolean =
+        session.annotationMatchingService.isAnnotated(resolvedReturnType.upperBoundIfFlexible().toRegularClassSymbol(session))
+
+    private fun buildFunctionCall(variableAssignment: FirVariableAssignment): FirFunctionCall {
+        val leftArgument = variableAssignment.lValue
+        val leftSymbol = leftArgument.resolvedSymbol as FirVariableSymbol<*>
+        val leftResolvedType = leftSymbol.resolvedReturnTypeRef
+        val rightArgument = variableAssignment.rValue
+        return buildFunctionCall {
+            source = variableAssignment.source?.fakeElement(KtFakeSourceElementKind.DesugaredCompoundAssignment)
+            explicitReceiver = buildPropertyAccessExpression {
+                source = leftArgument.source
+                typeRef = leftResolvedType
+                calleeReference = variableAssignment.calleeReference
+                typeArguments += variableAssignment.typeArguments
+                annotations += variableAssignment.annotations
+                explicitReceiver = variableAssignment.explicitReceiver
+                dispatchReceiver = variableAssignment.dispatchReceiver
+                extensionReceiver = variableAssignment.extensionReceiver
+                contextReceiverArguments += variableAssignment.contextReceiverArguments
+            }
+            argumentList = buildUnaryArgumentList(rightArgument)
+            calleeReference = buildSimpleNamedReference {
+                source = variableAssignment.source
+                name = ASSIGN_METHOD
+                candidateSymbol = null
+            }
+            origin = FirFunctionCallOrigin.Regular
+        }
+    }
+}
@@ -0,0 +1,30 @@
+/*
+ * Copyright 2010-2022 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.assignment.plugin.k2
+
+import org.jetbrains.kotlin.assignment.plugin.k2.diagnostics.FirAssignmentPluginFunctionCallChecker
+import org.jetbrains.kotlin.assignment.plugin.k2.diagnostics.FirAssignmentPluginFunctionChecker
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirSimpleFunctionChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirFunctionCallChecker
+import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
+
+class FirAssignmentPluginCheckersExtension(
+    session: FirSession
+) : FirAdditionalCheckersExtension(session) {
+
+    override val declarationCheckers: DeclarationCheckers = object : DeclarationCheckers() {
+        override val simpleFunctionCheckers: Set<FirSimpleFunctionChecker>
+            get() = setOf(FirAssignmentPluginFunctionChecker)
+    }
+
+    override val expressionCheckers: ExpressionCheckers = object : ExpressionCheckers() {
+        override val functionCallCheckers: Set<FirFunctionCallChecker>
+            get() = setOf(FirAssignmentPluginFunctionCallChecker)
+    }
+}
@@ -0,0 +1,16 @@
+/*
+ * Copyright 2010-2022 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.assignment.plugin.k2
+
+import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+
+class FirAssignmentPluginExtensionRegistrar(private val annotations: List<String>) : FirExtensionRegistrar() {
+    override fun ExtensionRegistrarContext.configurePlugin() {
+        +::FirAssignmentPluginAssignAltererExtension
+        +::FirAssignmentPluginCheckersExtension
+        +FirAssignAnnotationMatchingService.getFactory(annotations)
+    }
+}
@@ -0,0 +1,64 @@
+/*
+ * Copyright 2010-2022 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.assignment.plugin.k2.diagnostics
+
+import org.jetbrains.kotlin.KtFakeSourceElementKind
+import org.jetbrains.kotlin.assignment.plugin.AssignmentPluginNames.ASSIGN_METHOD
+import org.jetbrains.kotlin.assignment.plugin.k2.annotationMatchingService
+import org.jetbrains.kotlin.assignment.plugin.k2.diagnostics.FirErrorsAssignmentPlugin.CALL_ERROR_ASSIGN_METHOD_SHOULD_RETURN_UNIT
+import org.jetbrains.kotlin.assignment.plugin.k2.diagnostics.FirErrorsAssignmentPlugin.NO_APPLICABLE_ASSIGN_METHOD
+import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirFunctionCallChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClassSymbol
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.arguments
+import org.jetbrains.kotlin.fir.expressions.toResolvedCallableSymbol
+import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
+import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeAmbiguityError
+import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeDiagnosticWithSingleCandidate
+import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedNameError
+import org.jetbrains.kotlin.fir.types.isUnit
+
+object FirAssignmentPluginFunctionCallChecker : FirFunctionCallChecker() {
+
+    override fun check(expression: FirFunctionCall, context: CheckerContext, reporter: DiagnosticReporter) {
+        if (!expression.isOverloadAssignCallCandidate()) return
+
+        if (expression.isFunctionResolveError()) {
+            if (expression.isOverloadedAssignCallError(context.session)) {
+                reporter.reportOn(expression.source, NO_APPLICABLE_ASSIGN_METHOD, context)
+            }
+        } else if (expression.isOverloadedAssignCall(context.session) && !expression.isReturnTypeUnit()) {
+            reporter.reportOn(expression.source, CALL_ERROR_ASSIGN_METHOD_SHOULD_RETURN_UNIT, context)
+        }
+    }
+
+    private fun FirFunctionCall.isOverloadAssignCallCandidate() =
+        arguments.size == 1 && source?.kind == KtFakeSourceElementKind.DesugaredCompoundAssignment
+
+    private fun FirFunctionCall.isFunctionResolveError() = calleeReference is FirErrorNamedReference
+
+    private fun FirFunctionCall.isOverloadedAssignCallError(session: FirSession): Boolean {
+        val functionName = when (val diagnostic = (calleeReference as? FirErrorNamedReference)?.diagnostic) {
+            is ConeAmbiguityError -> diagnostic.name
+            is ConeDiagnosticWithSingleCandidate -> diagnostic.candidate.callInfo.name
+            is ConeUnresolvedNameError -> diagnostic.name
+            else -> calleeReference.name
+        }
+        return functionName == ASSIGN_METHOD && isAnnotated(session)
+    }
+
+    private fun FirFunctionCall.isOverloadedAssignCall(session: FirSession) =
+        calleeReference.name == ASSIGN_METHOD && isAnnotated(session)
+
+    private fun FirFunctionCall.isAnnotated(session: FirSession): Boolean =
+        session.annotationMatchingService.isAnnotated(explicitReceiver?.typeRef?.toRegularClassSymbol(session))
+
+    private fun FirFunctionCall.isReturnTypeUnit() = toResolvedCallableSymbol()?.resolvedReturnType?.isUnit ?: false
+}
@@ -0,0 +1,43 @@
+/*
+ * Copyright 2010-2022 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.assignment.plugin.k2.diagnostics
+
+import org.jetbrains.kotlin.assignment.plugin.AssignmentPluginNames.ASSIGN_METHOD
+import org.jetbrains.kotlin.assignment.plugin.k2.annotationMatchingService
+import org.jetbrains.kotlin.assignment.plugin.k2.diagnostics.FirErrorsAssignmentPlugin.DECLARATION_ERROR_ASSIGN_METHOD_SHOULD_RETURN_UNIT
+import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirSimpleFunctionChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClassSymbol
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.symbols.impl.isExtension
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.isUnit
+import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
+
+object FirAssignmentPluginFunctionChecker : FirSimpleFunctionChecker() {
+
+    override fun check(declaration: FirSimpleFunction, context: CheckerContext, reporter: DiagnosticReporter) {
+        if (declaration.origin != FirDeclarationOrigin.Source) return
+        if (!declaration.isAssignMethod()) return
+
+        val receiverClassSymbol = if (declaration.symbol.isExtension) {
+            declaration.symbol.resolvedReceiverTypeRef?.toRegularClassSymbol(context.session)
+        } else {
+            declaration.dispatchReceiverType?.toRegularClassSymbol(context.session)
+        }
+        if (!context.session.annotationMatchingService.isAnnotated(receiverClassSymbol)) return
+        if (!declaration.returnTypeRef.coneType.isUnit) {
+            reporter.reportOn(declaration.source, DECLARATION_ERROR_ASSIGN_METHOD_SHOULD_RETURN_UNIT, context)
+        }
+    }
+
+    private fun FirSimpleFunction.isAssignMethod(): Boolean {
+        return valueParameters.size == 1 && this.name == ASSIGN_METHOD
+    }
+}
@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2022 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.assignment.plugin.k2.diagnostics
+
+import org.jetbrains.kotlin.assignment.plugin.k2.diagnostics.FirErrorsAssignmentPlugin.CALL_ERROR_ASSIGN_METHOD_SHOULD_RETURN_UNIT
+import org.jetbrains.kotlin.assignment.plugin.k2.diagnostics.FirErrorsAssignmentPlugin.DECLARATION_ERROR_ASSIGN_METHOD_SHOULD_RETURN_UNIT
+import org.jetbrains.kotlin.assignment.plugin.k2.diagnostics.FirErrorsAssignmentPlugin.NO_APPLICABLE_ASSIGN_METHOD
+import org.jetbrains.kotlin.diagnostics.KtDiagnostic
+import org.jetbrains.kotlin.diagnostics.KtDiagnosticFactoryToRendererMap
+import org.jetbrains.kotlin.diagnostics.KtDiagnosticRenderer
+
+object FirDefaultErrorMessagesAssignmentPlugin {
+    fun getRendererForDiagnostic(diagnostic: KtDiagnostic): KtDiagnosticRenderer {
+        val factory = diagnostic.factory
+        return MAP[factory] ?: factory.ktRenderer
+    }
+
+    val MAP = KtDiagnosticFactoryToRendererMap("ValueContainerAssignment").also { map ->
+        map.put(
+            DECLARATION_ERROR_ASSIGN_METHOD_SHOULD_RETURN_UNIT,
+            "Function 'assign' used for '=' overload should return 'Unit'"
+        )
+
+        map.put(
+            CALL_ERROR_ASSIGN_METHOD_SHOULD_RETURN_UNIT,
+            "Function 'assign' used for '=' overload should return 'Unit'"
+        )
+
+        map.put(
+            NO_APPLICABLE_ASSIGN_METHOD,
+            "No applicable 'assign' function found for '=' overload"
+        )
+    }
+}
@@ -0,0 +1,17 @@
+/*
+ * Copyright 2010-2022 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.assignment.plugin.k2.diagnostics
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.diagnostics.SourceElementPositioningStrategies.DECLARATION_RETURN_TYPE
+import org.jetbrains.kotlin.diagnostics.SourceElementPositioningStrategies.OPERATOR
+import org.jetbrains.kotlin.diagnostics.error0
+
+object FirErrorsAssignmentPlugin {
+    val DECLARATION_ERROR_ASSIGN_METHOD_SHOULD_RETURN_UNIT by error0<PsiElement>(DECLARATION_RETURN_TYPE)
+    val CALL_ERROR_ASSIGN_METHOD_SHOULD_RETURN_UNIT by error0<PsiElement>(OPERATOR)
+    val NO_APPLICABLE_ASSIGN_METHOD by error0<PsiElement>(OPERATOR)
+}