[FIR] Add plugin prototype for testing extension for injection receivers

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/AlgebraReceiverInjector.kt

@@ -0,0 +1,35 @@
+/*
+ * 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.fir.plugin
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.extensions.FirExpressionResolutionExtension
+import org.jetbrains.kotlin.fir.plugin.generators.toSimpleConeType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirTypeProjectionWithVariance
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+/*
+ * Injects Algebra<T> as implicit receiver if `injectAlgebra<T>()` was called
+ */
+class AlgebraReceiverInjector(session: FirSession) : FirExpressionResolutionExtension(session) {
+    companion object {
+        private val INJECT_ALGEBRA_NAME = Name.identifier("injectAlgebra")
+        private val ALGEBRA_CLASS_ID = ClassId.topLevel(FqName.topLevel(Name.identifier("Algebra")))
+    }
+
+    override fun addNewImplicitReceivers(functionCall: FirFunctionCall): List<ConeKotlinType> {
+        if (functionCall.calleeReference.name != INJECT_ALGEBRA_NAME) return emptyList()
+        val typeProjection = functionCall.typeArguments.firstOrNull() as? FirTypeProjectionWithVariance ?: return emptyList()
+        val argumentType = typeProjection.typeRef.coneType
+        val algebraType = ALGEBRA_CLASS_ID.toSimpleConeType(arrayOf(argumentType))
+        return listOf(algebraType)
+    }
+}

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/FirPluginPrototypeExtensionRegistrar.kt

@@ -21,6 +21,7 @@ class FirPluginPrototypeExtensionRegistrar : FirExtensionRegistrar() {
         +::SomeAdditionalSupertypeGenerator
         +::PluginAdditionalCheckers
         +::FirNumberSignAttributeExtension
+        +::AlgebraReceiverInjector
 
         // Declaration generators
         +::TopLevelDeclarationsGenerator

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/generationUtils.kt

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.ConeKotlinTypeProjection
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.name.CallableId
@@ -91,10 +92,10 @@ fun FirDeclarationGenerationExtension.buildConstructor(classId: ClassId, isInner
     }
 }
 
-fun ClassId.toSimpleConeType(): ConeClassLikeType {
+fun ClassId.toSimpleConeType(typeArguments: Array<ConeKotlinTypeProjection> = emptyArray()): ConeClassLikeType {
     return ConeClassLikeTypeImpl(
         ConeClassLikeLookupTagImpl(this),
-        emptyArray(),
+        typeArguments,
         isNullable = false
     )
 }

plugins/fir-plugin-prototype/testData/diagnostics/receivers/receiverInjection.fir.txt

@@ -0,0 +1,32 @@
+FILE: receiverInjection.kt
+    public abstract interface Algebra<T> : R|kotlin/Any| {
+        public abstract operator fun R|T|.plus(other: R|T|): R|T|
+
+    }
+    public abstract interface A : R|kotlin/Any| {
+    }
+    public abstract interface B : R|kotlin/Any| {
+    }
+    public final fun <T> injectAlgebra(): R|kotlin/Unit| {
+    }
+    public final fun test_1(a1: R|A|, a2: R|A|, b1: R|B|, b2: R|B|): R|kotlin/Unit| {
+        R|<local>/a1|.<Unresolved name: plus>#(R|<local>/a2|)
+        R|<local>/b1|.<Unresolved name: plus>#(R|<local>/b2|)
+        R|/injectAlgebra|<R|A|>()
+        (this@R|/test_1|, R|<local>/a1|).R|SubstitutionOverride</Algebra.plus: R|A|>|(R|<local>/a2|)
+        (this@R|/test_1|, R|<local>/b1|).<None of the following candidates is applicable because of receiver type mismatch: [/Algebra.plus]>#(R|<local>/b2|)
+        R|/injectAlgebra|<R|B|>()
+        (this@R|/test_1|, R|<local>/a1|).R|SubstitutionOverride</Algebra.plus: R|A|>|(R|<local>/a2|)
+        (this@R|/test_1|, R|<local>/b1|).R|SubstitutionOverride</Algebra.plus: R|B|>|(R|<local>/b2|)
+    }
+    public final fun test_2(a1: R|A|, a2: R|A|, cond: R|kotlin/Boolean|): R|kotlin/Unit| {
+        R|<local>/a1|.<Unresolved name: plus>#(R|<local>/a2|)
+        when () {
+            R|<local>/cond| ->  {
+                R|/injectAlgebra|<R|A|>()
+                (this@R|/test_2|, R|<local>/a1|).R|SubstitutionOverride</Algebra.plus: R|A|>|(R|<local>/a2|)
+            }
+        }
+
+        R|<local>/a1|.<Unresolved name: plus>#(R|<local>/a2|)
+    }

plugins/fir-plugin-prototype/testData/diagnostics/receivers/receiverInjection.kt

@@ -0,0 +1,32 @@
+interface Algebra<T> {
+    operator fun T.plus(other: T): T
+}
+
+interface A
+interface B
+
+fun <T> injectAlgebra() {}
+
+fun test_1(a1: A, a2: A, b1: B, b2: B) {
+    a1 <!UNRESOLVED_REFERENCE!>+<!> a2 // error
+    b1 <!UNRESOLVED_REFERENCE!>+<!> b2 // error
+
+    injectAlgebra<A>()
+    a1 + a2 // ok
+    b1 <!UNRESOLVED_REFERENCE_WRONG_RECEIVER!>+<!> b2 // error
+
+    injectAlgebra<B>()
+    a1 + a2 // ok
+    b1 + b2 // ok
+}
+
+fun test_2(a1: A, a2: A, cond: Boolean) {
+    a1 <!UNRESOLVED_REFERENCE!>+<!> a2 // error
+
+    if (cond) {
+        injectAlgebra<A>()
+        a1 + a2 // ok
+    }
+
+    a1 <!UNRESOLVED_REFERENCE!>+<!> a2 // error
+}

plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirPluginDiagnosticTestGenerated.java

@@ -80,6 +80,22 @@ public class FirPluginDiagnosticTestGenerated extends AbstractFirPluginDiagnosti
         }
     }
 
+    @Nested
+    @TestMetadata("plugins/fir-plugin-prototype/testData/diagnostics/receivers")
+    @TestDataPath("$PROJECT_ROOT")
+    public class Receivers {
+        @Test
+        public void testAllFilesPresentInReceivers() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("plugins/fir-plugin-prototype/testData/diagnostics/receivers"), Pattern.compile("^(.+)\\.kt$"), null, true);
+        }
+
+        @Test
+        @TestMetadata("receiverInjection.kt")
+        public void testReceiverInjection() throws Exception {
+            runTest("plugins/fir-plugin-prototype/testData/diagnostics/receivers/receiverInjection.kt");
+        }
+    }
+
     @Nested
     @TestMetadata("plugins/fir-plugin-prototype/testData/diagnostics/status")
     @TestDataPath("$PROJECT_ROOT")