Skip to content

Commit 2f62d07

Browse files
JolanrensenJolanrensen
committed
started WIP fir plugin, disabled for now
1 parent 4c17859 commit 2f62d07

File tree

11 files changed

+283
-15
lines changed

11 files changed

+283
-15
lines changed

Diff for: ‎compiler-plugin/src/main/kotlin/org/jetbrains/kotlinx/spark/api/compilerPlugin/SimplePluginRegistrar.kt

-8
This file was deleted.

Diff for: ‎compiler-plugin/src/main/kotlin/org/jetbrains/kotlinx/spark/api/compilerPlugin/SparkifyCompilerPluginRegistrar.kt

+12-1
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@ package org.jetbrains.kotlinx.spark.api.compilerPlugin
33
import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
44
import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
55
import org.jetbrains.kotlin.config.CompilerConfiguration
6+
import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
7+
import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
68
import org.jetbrains.kotlinx.spark.api.Artifacts
79
import org.jetbrains.kotlinx.spark.api.compilerPlugin.ir.SparkifyIrGenerationExtension
810

9-
open class SparkifyCompilerPluginRegistrar: CompilerPluginRegistrar() {
11+
open class SparkifyCompilerPluginRegistrar : CompilerPluginRegistrar() {
1012
init {
1113
println("SparkifyCompilerPluginRegistrar loaded")
1214
}
@@ -26,6 +28,15 @@ open class SparkifyCompilerPluginRegistrar: CompilerPluginRegistrar() {
2628
val productFqNames = // TODO: get from configuration
2729
listOf("scala.Product")
2830

31+
// Front end (FIR)
32+
// FirExtensionRegistrarAdapter.registerExtension(
33+
// SparkifyFirPluginRegistrar(
34+
// sparkifyAnnotationFqNames = sparkifyAnnotationFqNames,
35+
// productFqNames = productFqNames,
36+
// )
37+
// )
38+
39+
// Intermediate Representation IR
2940
IrGenerationExtension.registerExtension(
3041
SparkifyIrGenerationExtension(
3142
sparkifyAnnotationFqNames = sparkifyAnnotationFqNames,

Diff for: ‎compiler-plugin/src/main/kotlin/org/jetbrains/kotlinx/spark/api/compilerPlugin/SparkifyFirPluginRegistrar.kt

+23
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
package org.jetbrains.kotlinx.spark.api.compilerPlugin
2+
3+
import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
4+
import org.jetbrains.kotlinx.spark.api.compilerPlugin.fir.DataClassSparkifyFunctionsGenerator
5+
import org.jetbrains.kotlinx.spark.api.compilerPlugin.fir.DataClassSparkifySuperTypeGenerator
6+
7+
// Potential future K2 FIR hook
8+
// TODO
9+
class SparkifyFirPluginRegistrar(
10+
private val sparkifyAnnotationFqNames: List<String>,
11+
private val productFqNames: List<String>
12+
) : FirExtensionRegistrar() {
13+
override fun ExtensionRegistrarContext.configurePlugin() {
14+
+DataClassSparkifySuperTypeGenerator.builder(
15+
sparkifyAnnotationFqNames = sparkifyAnnotationFqNames,
16+
productFqNames = productFqNames,
17+
)
18+
+DataClassSparkifyFunctionsGenerator.builder(
19+
sparkifyAnnotationFqNames = sparkifyAnnotationFqNames,
20+
productFqNames = productFqNames,
21+
)
22+
}
23+
}

Diff for: ‎compiler-plugin/src/main/kotlin/org/jetbrains/kotlinx/spark/api/compilerPlugin/fir/DataClassSparkifyFunctionsGenerator.kt

+115
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
package org.jetbrains.kotlinx.spark.api.compilerPlugin.fir
2+
3+
import org.jetbrains.kotlin.GeneratedDeclarationKey
4+
import org.jetbrains.kotlin.fir.FirSession
5+
import org.jetbrains.kotlin.fir.declarations.utils.isData
6+
import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
7+
import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
8+
import org.jetbrains.kotlin.fir.plugin.createMemberFunction
9+
import org.jetbrains.kotlin.fir.render
10+
import org.jetbrains.kotlin.fir.resolve.getSuperTypes
11+
import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
12+
import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
13+
import org.jetbrains.kotlin.fir.types.toClassSymbol
14+
import org.jetbrains.kotlin.name.CallableId
15+
import org.jetbrains.kotlin.name.Name
16+
17+
class DataClassSparkifyFunctionsGenerator(
18+
session: FirSession,
19+
private val sparkifyAnnotationFqNames: List<String>,
20+
private val productFqNames: List<String>,
21+
) : FirDeclarationGenerationExtension(session) {
22+
23+
companion object {
24+
fun builder(
25+
sparkifyAnnotationFqNames: List<String>,
26+
productFqNames: List<String>
27+
): (FirSession) -> FirDeclarationGenerationExtension = {
28+
DataClassSparkifyFunctionsGenerator(
29+
session = it,
30+
sparkifyAnnotationFqNames = sparkifyAnnotationFqNames,
31+
productFqNames = productFqNames,
32+
)
33+
}
34+
35+
// functions to generate
36+
val canEqual = Name.identifier("canEqual")
37+
val productElement = Name.identifier("productElement")
38+
val productArity = Name.identifier("productArity")
39+
}
40+
41+
override fun generateFunctions(
42+
callableId: CallableId,
43+
context: MemberGenerationContext?
44+
): List<FirNamedFunctionSymbol> {
45+
val owner = context?.owner ?: return emptyList()
46+
47+
val functionName = callableId.callableName
48+
val superTypes = owner.getSuperTypes(session)
49+
val superProduct = superTypes.first {
50+
it.toString().endsWith("Product")
51+
}.toClassSymbol(session)!!
52+
val superEquals = superTypes.first {
53+
it.toString().endsWith("Equals")
54+
}.toClassSymbol(session)!!
55+
56+
val function = when (functionName) {
57+
canEqual -> {
58+
val func = createMemberFunction(
59+
owner = owner,
60+
key = Key,
61+
name = functionName,
62+
returnType = session.builtinTypes.booleanType.type,
63+
) {
64+
valueParameter(
65+
name = Name.identifier("that"),
66+
type = session.builtinTypes.nullableAnyType.type,
67+
)
68+
}
69+
// val superFunction = superEquals.declarationSymbols.first {
70+
// it is FirNamedFunctionSymbol && it.name == functionName
71+
// } as FirNamedFunctionSymbol
72+
// overrides(func, superFunction)
73+
func
74+
}
75+
76+
productElement -> {
77+
createMemberFunction(
78+
owner = owner,
79+
key = Key,
80+
name = functionName,
81+
returnType = session.builtinTypes.nullableAnyType.type,
82+
) {
83+
valueParameter(
84+
name = Name.identifier("n"),
85+
type = session.builtinTypes.intType.type,
86+
)
87+
}
88+
}
89+
90+
productArity -> {
91+
createMemberFunction(
92+
owner = owner,
93+
key = Key,
94+
name = functionName,
95+
returnType = session.builtinTypes.intType.type,
96+
)
97+
}
98+
99+
else -> {
100+
return emptyList()
101+
}
102+
}
103+
104+
return listOf(function.symbol)
105+
}
106+
107+
override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> =
108+
if (classSymbol.isData && classSymbol.annotations.any { "Sparkify" in it.render() }) {
109+
setOf(canEqual, productElement, productArity)
110+
} else {
111+
emptySet()
112+
}
113+
114+
object Key : GeneratedDeclarationKey()
115+
}

Diff for: ‎compiler-plugin/src/main/kotlin/org/jetbrains/kotlinx/spark/api/compilerPlugin/fir/DataClassSparkifySuperTypeGenerator.kt

+59
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
package org.jetbrains.kotlinx.spark.api.compilerPlugin.fir
2+
3+
import org.jetbrains.kotlin.fir.FirSession
4+
import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
5+
import org.jetbrains.kotlin.fir.declarations.utils.isData
6+
import org.jetbrains.kotlin.fir.extensions.FirSupertypeGenerationExtension
7+
import org.jetbrains.kotlin.fir.render
8+
import org.jetbrains.kotlin.fir.resolve.fqName
9+
import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
10+
import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
11+
import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
12+
import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
13+
import org.jetbrains.kotlin.name.ClassId
14+
import org.jetbrains.kotlin.name.FqName
15+
16+
/**
17+
* This class tells the FIR that all @Sparkify annotated data classes
18+
* get [scala.Product] as their super type.
19+
*/
20+
class DataClassSparkifySuperTypeGenerator(
21+
session: FirSession,
22+
private val sparkifyAnnotationFqNames: List<String>,
23+
private val productFqNames: List<String>,
24+
) : FirSupertypeGenerationExtension(session) {
25+
26+
companion object {
27+
fun builder(sparkifyAnnotationFqNames: List<String>, productFqNames: List<String>): (FirSession) -> FirSupertypeGenerationExtension = {
28+
DataClassSparkifySuperTypeGenerator(
29+
session = it,
30+
sparkifyAnnotationFqNames = sparkifyAnnotationFqNames,
31+
productFqNames = productFqNames,
32+
)
33+
}
34+
}
35+
36+
context(TypeResolveServiceContainer)
37+
override fun computeAdditionalSupertypes(
38+
classLikeDeclaration: FirClassLikeDeclaration,
39+
resolvedSupertypes: List<FirResolvedTypeRef>
40+
): List<FirResolvedTypeRef> = listOf(
41+
buildResolvedTypeRef {
42+
val scalaProduct = productFqNames.first().let {
43+
ClassId.topLevel(FqName(it))
44+
}
45+
type = ConeClassLikeTypeImpl(
46+
lookupTag = ConeClassLikeLookupTagImpl(scalaProduct),
47+
typeArguments = emptyArray(),
48+
isNullable = false,
49+
)
50+
}
51+
52+
)
53+
54+
override fun needTransformSupertypes(declaration: FirClassLikeDeclaration): Boolean =
55+
declaration.symbol.isData &&
56+
declaration.annotations.any {
57+
"Sparkify" in it.render()
58+
}
59+
}

Diff for: ‎compiler-plugin/src/main/kotlin/org/jetbrains/kotlinx/spark/api/compilerPlugin/ir/DataClassPropertyAnnotationGenerator.kt renamed to ‎compiler-plugin/src/main/kotlin/org/jetbrains/kotlinx/spark/api/compilerPlugin/ir/DataClassSparkifyGenerator.kt

+1-1
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ import org.jetbrains.kotlin.name.FqName
4646
import org.jetbrains.kotlin.name.Name
4747
import org.jetbrains.kotlin.name.SpecialNames
4848

49-
class DataClassPropertyAnnotationGenerator(
49+
class DataClassSparkifyGenerator(
5050
private val pluginContext: IrPluginContext,
5151
private val sparkifyAnnotationFqNames: List<String>,
5252
private val columnNameAnnotationFqNames: List<String>,

Diff for: ‎compiler-plugin/src/main/kotlin/org/jetbrains/kotlinx/spark/api/compilerPlugin/ir/SparkifyIrGenerationExtension.kt

+1-2
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
44
import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
55
import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
66
import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
7-
import org.jetbrains.kotlinx.spark.api.compilerPlugin.ir.DataClassPropertyAnnotationGenerator
87

98
class SparkifyIrGenerationExtension(
109
private val sparkifyAnnotationFqNames: List<String>,
@@ -13,7 +12,7 @@ class SparkifyIrGenerationExtension(
1312
) : IrGenerationExtension {
1413
override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
1514
val visitors = listOf(
16-
DataClassPropertyAnnotationGenerator(
15+
DataClassSparkifyGenerator(
1716
pluginContext = pluginContext,
1817
sparkifyAnnotationFqNames = sparkifyAnnotationFqNames,
1918
columnNameAnnotationFqNames = columnNameAnnotationFqNames,

Diff for: ‎compiler-plugin/src/test-gen/kotlin/org/jetbrains/kotlinx/spark/api/compilerPlugin/runners/DiagnosticTestGenerated.java

+28
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
2+
3+
package org.jetbrains.kotlinx.spark.api.compilerPlugin.runners;
4+
5+
import com.intellij.testFramework.TestDataPath;
6+
import org.jetbrains.kotlin.test.util.KtTestUtil;
7+
import org.jetbrains.kotlin.test.TestMetadata;
8+
import org.junit.jupiter.api.Test;
9+
10+
import java.io.File;
11+
import java.util.regex.Pattern;
12+
13+
/** This class is generated by {@link org.jetbrains.kotlinx.spark.api.compilerPlugin.GenerateTestsKt}. DO NOT MODIFY MANUALLY */
14+
@SuppressWarnings("all")
15+
@TestMetadata("/mnt/data/Projects/kotlin-spark-api/compiler-plugin/src/test/resources/testData/diagnostics")
16+
@TestDataPath("$PROJECT_ROOT")
17+
public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
18+
@Test
19+
public void testAllFilesPresentInDiagnostics() {
20+
KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("/mnt/data/Projects/kotlin-spark-api/compiler-plugin/src/test/resources/testData/diagnostics"), Pattern.compile("^(.+)\\.kt$"), null, true);
21+
}
22+
23+
@Test
24+
@TestMetadata("dataClassTest.kt")
25+
public void testDataClassTest() {
26+
runTest("/mnt/data/Projects/kotlin-spark-api/compiler-plugin/src/test/resources/testData/diagnostics/dataClassTest.kt");
27+
}
28+
}

Diff for: ‎compiler-plugin/src/test/kotlin/org/jetbrains/kotlinx/spark/api/compilerPlugin/GenerateTests.kt

+4-3
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,17 @@ package org.jetbrains.kotlinx.spark.api.compilerPlugin
33
import org.jetbrains.kotlin.generators.generateTestGroupSuiteWithJUnit5
44
import org.jetbrains.kotlinx.spark.api.Artifacts
55
import org.jetbrains.kotlinx.spark.api.compilerPlugin.runners.AbstractBoxTest
6+
import org.jetbrains.kotlinx.spark.api.compilerPlugin.runners.AbstractDiagnosticTest
67

78
fun main() {
89
generateTestGroupSuiteWithJUnit5 {
910
testGroup(
1011
testDataRoot = "${Artifacts.projectRoot}/${Artifacts.compilerPluginArtifactId}/src/test/resources/testData",
1112
testsRoot = "${Artifacts.projectRoot}/${Artifacts.compilerPluginArtifactId}/src/test-gen/kotlin",
1213
) {
13-
// testClass<AbstractDiagnosticTest> {
14-
// model("diagnostics")
15-
// }
14+
testClass<AbstractDiagnosticTest> {
15+
model("diagnostics")
16+
}
1617

1718
testClass<AbstractBoxTest> {
1819
model("box")

Diff for: ‎compiler-plugin/src/test/kotlin/org/jetbrains/kotlinx/spark/api/compilerPlugin/services/ExtensionRegistrarConfigurator.kt

+12
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,11 @@ package org.jetbrains.kotlinx.spark.api.compilerPlugin.services
33
import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
44
import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
55
import org.jetbrains.kotlin.config.CompilerConfiguration
6+
import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
67
import org.jetbrains.kotlin.test.model.TestModule
78
import org.jetbrains.kotlin.test.services.EnvironmentConfigurator
89
import org.jetbrains.kotlin.test.services.TestServices
10+
import org.jetbrains.kotlinx.spark.api.compilerPlugin.SparkifyFirPluginRegistrar
911
import org.jetbrains.kotlinx.spark.api.compilerPlugin.ir.SparkifyIrGenerationExtension
1012

1113
class ExtensionRegistrarConfigurator(testServices: TestServices) : EnvironmentConfigurator(testServices) {
@@ -16,6 +18,16 @@ class ExtensionRegistrarConfigurator(testServices: TestServices) : EnvironmentCo
1618
val sparkifyAnnotationFqNames = listOf("foo.bar.Sparkify")
1719
val columnNameAnnotationFqNames = listOf("foo.bar.ColumnName")
1820
val productFqNames = listOf("foo.bar.Product")
21+
22+
// Front end (FIR)
23+
// FirExtensionRegistrarAdapter.registerExtension(
24+
// SparkifyFirPluginRegistrar(
25+
// sparkifyAnnotationFqNames = sparkifyAnnotationFqNames,
26+
// productFqNames = productFqNames,
27+
// )
28+
// )
29+
30+
// Intermediate Representation IR
1931
IrGenerationExtension.registerExtension(
2032
SparkifyIrGenerationExtension(
2133
sparkifyAnnotationFqNames = sparkifyAnnotationFqNames,

Diff for: ‎compiler-plugin/src/test/resources/testData/diagnostics/dataClassTest.kt

+28
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
package foo.bar
2+
3+
annotation class Sparkify
4+
annotation class ColumnName(val name: String)
5+
6+
// Fake Equals
7+
interface Equals {
8+
fun canEqual(that: Any?): Boolean
9+
}
10+
11+
// Fake Product
12+
interface Product: Equals {
13+
fun productElement(n: Int): Any
14+
fun productArity(): Int
15+
}
16+
17+
fun test() {
18+
val user = User()
19+
user.productArity() // should not be an error
20+
}
21+
22+
@Sparkify
23+
data <!ABSTRACT_MEMBER_NOT_IMPLEMENTED!>class User<!>(
24+
val name: String = "John Doe",
25+
val age: Int = 25,
26+
@ColumnName("a") val test: Double = 1.0,
27+
@get:ColumnName("b") val test2: Double = 2.0,
28+
)

0 commit comments

Comments
 (0)