Skip to content

Commit 81ec10a

Browse files
koperagenkoperagen
authored
Merge pull request #1051 from Kotlin/compiler-plugin-checkers
 Add checkers that report compile time schema as info warnings to observe implicit schema generation
2 parents 1a102b9 + dc858cd commit 81ec10a

File tree

12 files changed

+753
-23
lines changed

12 files changed

+753
-23
lines changed

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/FirDataFrameComponentRegistrar.kt‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ class FirDataFrameExtensionRegistrar(
6565
private val path: String?,
6666
val schemasDirectory: String?,
6767
val isTest: Boolean,
68+
val dumpSchemas: Boolean,
6869
) : FirExtensionRegistrar() {
6970
@OptIn(FirExtensionApiInternals::class)
7071
override fun ExtensionRegistrarContext.configurePlugin() {
@@ -76,7 +77,7 @@ class FirDataFrameExtensionRegistrar(
7677
+::TokenGenerator
7778
+::DataRowSchemaSupertype
7879
+{ it: FirSession ->
79-
ExpressionAnalysisAdditionalChecker(it, jsonCache(it), schemasDirectory, isTest)
80+
ExpressionAnalysisAdditionalChecker(it, jsonCache(it), schemasDirectory, isTest, dumpSchemas)
8081
}
8182
}
8283

@@ -93,7 +94,9 @@ class FirDataFrameComponentRegistrar : CompilerPluginRegistrar() {
9394
override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
9495
val schemasDirectory = configuration.get(SCHEMAS)
9596
val path = configuration.get(PATH)
96-
FirExtensionRegistrarAdapter.registerExtension(FirDataFrameExtensionRegistrar(path, schemasDirectory, isTest = false))
97+
FirExtensionRegistrarAdapter.registerExtension(
98+
FirDataFrameExtensionRegistrar(path, schemasDirectory, isTest = false, dumpSchemas = true)
99+
)
97100
IrGenerationExtension.registerExtension(IrBodyFiller(path, schemasDirectory))
98101
}
99102

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExpressionAnalysisAdditionalChecker.kt‎

Lines changed: 134 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,26 +5,38 @@
55

66
package org.jetbrains.kotlinx.dataframe.plugin.extensions
77

8+
import com.intellij.psi.PsiElement
9+
import org.jetbrains.kotlin.KtSourceElement
10+
import org.jetbrains.kotlin.diagnostics.AbstractSourceElementPositioningStrategy
11+
import org.jetbrains.kotlin.diagnostics.DiagnosticFactory1DelegateProvider
812
import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
13+
import org.jetbrains.kotlin.diagnostics.KtDiagnosticFactory1
14+
import org.jetbrains.kotlin.diagnostics.Severity
915
import org.jetbrains.kotlin.diagnostics.SourceElementPositioningStrategies
1016
import org.jetbrains.kotlin.diagnostics.error1
1117
import org.jetbrains.kotlin.diagnostics.reportOn
1218
import org.jetbrains.kotlin.diagnostics.warning1
1319
import org.jetbrains.kotlin.fir.FirSession
1420
import org.jetbrains.kotlin.fir.analysis.checkers.MppCheckerKind
1521
import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
22+
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
23+
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirPropertyChecker
24+
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirSimpleFunctionChecker
1625
import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
1726
import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirFunctionCallChecker
27+
import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirPropertyAccessExpressionChecker
1828
import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
1929
import org.jetbrains.kotlin.fir.caches.FirCache
20-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.flatten
21-
import org.jetbrains.kotlinx.dataframe.plugin.pluginDataFrameSchema
30+
import org.jetbrains.kotlin.fir.declarations.FirProperty
31+
import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
2232
import org.jetbrains.kotlin.fir.declarations.hasAnnotation
2333
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
34+
import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
2435
import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
2536
import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
2637
import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
2738
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
39+
import org.jetbrains.kotlin.fir.types.ConeKotlinType
2840
import org.jetbrains.kotlin.fir.types.FirTypeProjectionWithVariance
2941
import org.jetbrains.kotlin.fir.types.coneType
3042
import org.jetbrains.kotlin.fir.types.isSubtypeOf
@@ -38,18 +50,34 @@ import org.jetbrains.kotlin.name.FqName
3850
import org.jetbrains.kotlin.name.Name
3951
import org.jetbrains.kotlin.psi.KtElement
4052
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
53+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
4154
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
42-
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
55+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
56+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.flatten
57+
import org.jetbrains.kotlinx.dataframe.plugin.pluginDataFrameSchema
4358
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
59+
import org.jetbrains.kotlinx.dataframe.plugin.utils.isDataFrame
60+
import org.jetbrains.kotlinx.dataframe.plugin.utils.isDataRow
61+
import org.jetbrains.kotlinx.dataframe.plugin.utils.isGroupBy
4462

4563
class ExpressionAnalysisAdditionalChecker(
4664
session: FirSession,
4765
cache: FirCache<String, PluginDataFrameSchema, KotlinTypeFacade>,
4866
schemasDirectory: String?,
4967
isTest: Boolean,
68+
dumpSchemas: Boolean
5069
) : FirAdditionalCheckersExtension(session) {
5170
override val expressionCheckers: ExpressionCheckers = object : ExpressionCheckers() {
52-
override val functionCallCheckers: Set<FirFunctionCallChecker> = setOf(Checker(cache, schemasDirectory, isTest))
71+
override val functionCallCheckers: Set<FirFunctionCallChecker> = setOfNotNull(
72+
Checker(cache, schemasDirectory, isTest), FunctionCallSchemaReporter.takeIf { dumpSchemas }
73+
)
74+
override val propertyAccessExpressionCheckers: Set<FirPropertyAccessExpressionChecker> = setOfNotNull(
75+
PropertyAccessSchemaReporter.takeIf { dumpSchemas }
76+
)
77+
}
78+
override val declarationCheckers: DeclarationCheckers = object : DeclarationCheckers() {
79+
override val propertyCheckers: Set<FirPropertyChecker> = setOfNotNull(PropertySchemaReporter.takeIf { dumpSchemas })
80+
override val simpleFunctionCheckers: Set<FirSimpleFunctionChecker> = setOfNotNull(FunctionDeclarationSchemaReporter.takeIf { dumpSchemas })
5381
}
5482
}
5583

@@ -132,3 +160,105 @@ private class Checker(
132160
}
133161
}
134162
}
163+
164+
private data object PropertySchemaReporter : FirPropertyChecker(mppKind = MppCheckerKind.Common) {
165+
val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.DECLARATION_NAME)
166+
167+
override fun check(declaration: FirProperty, context: CheckerContext, reporter: DiagnosticReporter) {
168+
context.sessionContext {
169+
declaration.returnTypeRef.coneType.let { type ->
170+
reportSchema(reporter, declaration.source, SCHEMA, type, context)
171+
}
172+
}
173+
}
174+
}
175+
176+
private data object FunctionCallSchemaReporter : FirFunctionCallChecker(mppKind = MppCheckerKind.Common) {
177+
val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
178+
179+
override fun check(expression: FirFunctionCall, context: CheckerContext, reporter: DiagnosticReporter) {
180+
if (expression.calleeReference.name in setOf(Name.identifier("let"), Name.identifier("run"))) return
181+
val initializer = expression.resolvedType
182+
context.sessionContext {
183+
reportSchema(reporter, expression.source, SCHEMA, initializer, context)
184+
}
185+
}
186+
}
187+
188+
private data object PropertyAccessSchemaReporter : FirPropertyAccessExpressionChecker(mppKind = MppCheckerKind.Common) {
189+
val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
190+
191+
override fun check(
192+
expression: FirPropertyAccessExpression,
193+
context: CheckerContext,
194+
reporter: DiagnosticReporter
195+
) {
196+
val initializer = expression.resolvedType
197+
context.sessionContext {
198+
reportSchema(reporter, expression.source, SCHEMA, initializer, context)
199+
}
200+
}
201+
}
202+
203+
private data object FunctionDeclarationSchemaReporter : FirSimpleFunctionChecker(mppKind = MppCheckerKind.Common) {
204+
val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.DECLARATION_SIGNATURE)
205+
206+
override fun check(declaration: FirSimpleFunction, context: CheckerContext, reporter: DiagnosticReporter) {
207+
val type = declaration.returnTypeRef.coneType
208+
context.sessionContext {
209+
reportSchema(reporter, declaration.source, SCHEMA, type, context)
210+
}
211+
}
212+
}
213+
214+
private fun SessionContext.reportSchema(
215+
reporter: DiagnosticReporter,
216+
source: KtSourceElement?,
217+
factory: KtDiagnosticFactory1<String>,
218+
type: ConeKotlinType,
219+
context: CheckerContext,
220+
) {
221+
val expandedType = type.fullyExpandedType(session)
222+
var schema: PluginDataFrameSchema? = null
223+
when {
224+
expandedType.isDataFrame(session) -> {
225+
schema = expandedType.typeArguments.getOrNull(0)?.let {
226+
pluginDataFrameSchema(it)
227+
}
228+
}
229+
230+
expandedType.isDataRow(session) -> {
231+
schema = expandedType.typeArguments.getOrNull(0)?.let {
232+
pluginDataFrameSchema(it)
233+
}
234+
}
235+
236+
expandedType.isGroupBy(session) -> {
237+
val keys = expandedType.typeArguments.getOrNull(0)
238+
val grouped = expandedType.typeArguments.getOrNull(1)
239+
if (keys != null && grouped != null) {
240+
val keysSchema = pluginDataFrameSchema(keys)
241+
val groupedSchema = pluginDataFrameSchema(grouped)
242+
schema = PluginDataFrameSchema(
243+
listOf(
244+
SimpleColumnGroup("keys", keysSchema.columns()),
245+
SimpleFrameColumn("groups", groupedSchema.columns())
246+
)
247+
)
248+
}
249+
}
250+
}
251+
if (schema != null && source != null) {
252+
reporter.reportOn(source, factory, "\n" + schema.toString(), context)
253+
}
254+
}
255+
256+
fun CheckerContext.sessionContext(f: SessionContext.() -> Unit) {
257+
SessionContext(session).f()
258+
}
259+
260+
inline fun <reified P : PsiElement, A> info1(
261+
positioningStrategy: AbstractSourceElementPositioningStrategy = SourceElementPositioningStrategies.DEFAULT
262+
): DiagnosticFactory1DelegateProvider<A> {
263+
return DiagnosticFactory1DelegateProvider(Severity.INFO, positioningStrategy, P::class)
264+
}

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,6 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
8181
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
8282
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
8383
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBy
84-
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.createPluginDataFrameSchema
8584
import kotlin.math.abs
8685

8786
@OptIn(FirExtensionApiInternals::class)
@@ -321,14 +320,14 @@ class FunctionCallTransformer(
321320
val receiverType = explicitReceiver?.resolvedType
322321
val returnType = call.resolvedType
323322
val scopeFunction = if (explicitReceiver != null) findLet() else findRun()
323+
val originalSource = call.calleeReference.source
324324

325325
// original call is inserted later
326326
call.transformCalleeReference(object : FirTransformer<Nothing?>() {
327327
override fun <E : FirElement> transformElement(element: E, data: Nothing?): E {
328328
return if (element is FirResolvedNamedReference) {
329329
@Suppress("UNCHECKED_CAST")
330330
buildResolvedNamedReference {
331-
source = call.calleeReference.source
332331
this.name = element.name
333332
resolvedSymbol = originalSymbol
334333
} as E
@@ -430,7 +429,8 @@ class FunctionCallTransformer(
430429
}
431430

432431
val newCall1 = buildFunctionCall {
433-
source = call.source
432+
// source = call.source makes IDE navigate to `let` declaration
433+
source = null
434434
this.coneTypeOrNull = returnType
435435
if (receiverType != null) {
436436
typeArguments += buildTypeProjectionWithVariance {
@@ -455,7 +455,7 @@ class FunctionCallTransformer(
455455
linkedMapOf(argument to scopeFunction.valueParameterSymbols[0].fir)
456456
)
457457
calleeReference = buildResolvedNamedReference {
458-
source = call.calleeReference.source
458+
source = originalSource
459459
this.name = scopeFunction.name
460460
resolvedSymbol = scopeFunction
461461
}

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/KotlinTypeFacade.kt‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,7 @@ import kotlin.reflect.KType
2727
import kotlin.reflect.KTypeProjection
2828
import kotlin.reflect.KVariance
2929

30-
interface KotlinTypeFacade {
31-
val session: FirSession
30+
interface KotlinTypeFacade : SessionContext {
3231
val resolutionPath: String? get() = null
3332
val cache: FirCache<String, PluginDataFrameSchema, KotlinTypeFacade>
3433
val schemasDirectory: String?
@@ -99,6 +98,14 @@ interface KotlinTypeFacade {
9998
}
10099
}
101100

101+
interface SessionContext {
102+
val session: FirSession
103+
}
104+
105+
fun SessionContext(session: FirSession) = object : SessionContext {
106+
override val session: FirSession = session
107+
}
108+
102109
private val List = "List".collectionsId()
103110

104111
private fun ConeKotlinType.isBuiltinType(classId: ClassId, isNullable: Boolean?): Boolean {

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/SimpleCol.kt‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
1-
@file:Suppress("INVISIBLE_REFERENCE", "CANNOT_OVERRIDE_INVISIBLE_MEMBER")
2-
31
package org.jetbrains.kotlinx.dataframe.plugin.impl
42

53
import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId
64
import org.jetbrains.kotlin.fir.types.ConeKotlinType
75
import org.jetbrains.kotlin.fir.types.ConeNullability
86
import org.jetbrains.kotlin.fir.types.isNullable
7+
import org.jetbrains.kotlin.fir.types.renderReadable
98
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
109
import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap
1110
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
@@ -18,6 +17,7 @@ data class PluginDataFrameSchema(
1817
companion object {
1918
val EMPTY = PluginDataFrameSchema(emptyList())
2019
}
20+
2121
fun columns(): List<SimpleCol> {
2222
return columns
2323
}
@@ -32,16 +32,19 @@ fun PluginDataFrameSchema.add(name: String, type: ConeKotlinType, context: Kotli
3232
}
3333

3434
private fun List<SimpleCol>.asString(indent: String = ""): String {
35+
if (isEmpty()) return "$indent<empty compile time schema>"
3536
return joinToString("\n") {
3637
val col = when (it) {
3738
is SimpleFrameColumn -> {
38-
"${it.name}*\n" + it.columns().asString("$indent ")
39+
"${it.name}: *\n" + it.columns().asString("$indent ")
3940
}
41+
4042
is SimpleColumnGroup -> {
41-
"${it.name}\n" + it.columns().asString("$indent ")
43+
"${it.name}:\n" + it.columns().asString("$indent ")
4244
}
45+
4346
is SimpleDataColumn -> {
44-
"${it.name}: ${it.type}"
47+
"${it.name}: ${it.type.type.renderReadable()}"
4548
}
4649
}
4750
"$indent$col"
@@ -127,6 +130,7 @@ private fun KotlinTypeFacade.makeNullable(column: SimpleCol): SimpleCol {
127130
is SimpleColumnGroup -> {
128131
SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
129132
}
133+
130134
is SimpleFrameColumn -> column
131135
is SimpleDataColumn -> SimpleDataColumn(column.name, column.type.changeNullability { true })
132136
}

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/interpret.kt‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ import org.jetbrains.kotlin.name.FqName
6969
import org.jetbrains.kotlin.name.Name
7070
import org.jetbrains.kotlin.utils.addToStdlib.runIf
7171
import org.jetbrains.kotlinx.dataframe.annotations.*
72+
import org.jetbrains.kotlinx.dataframe.plugin.extensions.SessionContext
7273
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
7374
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
7475
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.DataFrameCallableId
@@ -333,7 +334,7 @@ interface InterpretationErrorReporter {
333334
}
334335
}
335336

336-
fun KotlinTypeFacade.pluginDataFrameSchema(schemaTypeArg: ConeTypeProjection): PluginDataFrameSchema {
337+
fun SessionContext.pluginDataFrameSchema(schemaTypeArg: ConeTypeProjection): PluginDataFrameSchema {
337338
val schema = if (schemaTypeArg.isStarProjection) {
338339
PluginDataFrameSchema.EMPTY
339340
} else {
@@ -343,7 +344,7 @@ fun KotlinTypeFacade.pluginDataFrameSchema(schemaTypeArg: ConeTypeProjection): P
343344
return schema
344345
}
345346

346-
fun KotlinTypeFacade.pluginDataFrameSchema(coneClassLikeType: ConeClassLikeType): PluginDataFrameSchema {
347+
fun SessionContext.pluginDataFrameSchema(coneClassLikeType: ConeClassLikeType): PluginDataFrameSchema {
347348
val symbol = coneClassLikeType.toSymbol(session) as? FirRegularClassSymbol ?: return PluginDataFrameSchema.EMPTY
348349
val declarationSymbols = if (symbol.isLocal && symbol.resolvedSuperTypes.firstOrNull() != session.builtinTypes.anyType.type) {
349350
val rootSchemaSymbol = symbol.resolvedSuperTypes.first().toSymbol(session) as? FirRegularClassSymbol
@@ -407,7 +408,7 @@ private fun KotlinTypeFacade.columnWithPathApproximations(result: FirPropertyAcc
407408
}
408409
}
409410

410-
private fun KotlinTypeFacade.columnOf(it: FirPropertySymbol, mapping: Map<FirTypeParameterSymbol, ConeTypeProjection>): SimpleCol? {
411+
private fun SessionContext.columnOf(it: FirPropertySymbol, mapping: Map<FirTypeParameterSymbol, ConeTypeProjection>): SimpleCol? {
411412
val annotation = it.getAnnotationByClassId(Names.COLUMN_NAME_ANNOTATION, session)
412413
val columnName = (annotation?.argumentMapping?.mapping?.get(Names.COLUMN_NAME_ARGUMENT) as? FirLiteralExpression)?.value as? String
413414
val name = columnName ?: it.name.identifier
@@ -456,14 +457,14 @@ private fun KotlinTypeFacade.columnOf(it: FirPropertySymbol, mapping: Map<FirTyp
456457
}
457458
}
458459

459-
private fun KotlinTypeFacade.shouldBeConvertedToColumnGroup(it: FirPropertySymbol) =
460+
private fun SessionContext.shouldBeConvertedToColumnGroup(it: FirPropertySymbol) =
460461
isDataRow(it) ||
461462
it.resolvedReturnType.toRegularClassSymbol(session)?.hasAnnotation(Names.DATA_SCHEMA_CLASS_ID, session) == true
462463

463464
private fun isDataRow(it: FirPropertySymbol) =
464465
it.resolvedReturnType.classId == Names.DATA_ROW_CLASS_ID
465466

466-
private fun KotlinTypeFacade.shouldBeConvertedToFrameColumn(it: FirPropertySymbol) =
467+
private fun SessionContext.shouldBeConvertedToFrameColumn(it: FirPropertySymbol) =
467468
isDataFrame(it) ||
468469
(it.resolvedReturnType.classId == Names.LIST &&
469470
it.resolvedReturnType.typeArguments[0].type?.toRegularClassSymbol(session)?.hasAnnotation(Names.DATA_SCHEMA_CLASS_ID, session) == true)

0 commit comments

Comments
 (0)