Skip to content

Commit f3e5738

Browse files
koperagenkoperagen
committed
[Compiler plugin] refactor SimpleCol to closed hierarchy
1 parent 430e2fa commit f3e5738

File tree

19 files changed

+104
-156
lines changed

19 files changed

+104
-156
lines changed

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ import org.jetbrains.kotlin.name.FqName
3636
import org.jetbrains.kotlin.name.Name
3737
import org.jetbrains.kotlin.psi.KtElement
3838
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
39+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
40+
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
3941

4042
class ExpressionAnalysisAdditionalChecker(
4143
session: FirSession,
@@ -96,6 +98,7 @@ private class Checker(val cache: FirCache<String, PluginDataFrameSchema, KotlinT
9698
for (target in targetColumns) {
9799
val source = sourceMap[target.path.path]
98100
val present = if (source != null) {
101+
if (source !is SimpleDataColumn || target.column !is SimpleDataColumn) { continue }
99102
if (source.type.type().isSubtypeOf(target.column.type.type(), session)) {
100103
true
101104
} else {

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ import org.jetbrains.kotlin.text
7272
import org.jetbrains.kotlin.types.Variance
7373
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
7474
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
75+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
7576
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
7677
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
7778
import kotlin.math.abs
@@ -278,14 +279,12 @@ class FunctionCallTransformer(
278279
)
279280
}
280281

281-
is SimpleCol -> SchemaProperty(
282+
is SimpleDataColumn -> SchemaProperty(
282283
marker = schema.defaultType(),
283284
name = it.name,
284285
dataRowReturnType = it.type.type(),
285286
columnContainerReturnType = it.type.type().toFirResolvedTypeRef().projectOverDataColumnType()
286287
)
287-
288-
else -> TODO("shouldn't happen")
289288
}
290289
}
291290
schema.callShapeData = CallShapeData.Schema(properties)

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

Lines changed: 29 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -34,135 +34,74 @@ private fun List<SimpleCol>.asString(indent: String = ""): String {
3434
is SimpleColumnGroup -> {
3535
"${it.name}\n" + it.columns().asString("$indent ")
3636
}
37-
is SimpleCol -> {
38-
// val type = (it.type as TypeApproximationImpl).let {
39-
// val nullability = if (it.nullable) "?" else ""
40-
// "${it.fqName}$nullability"
41-
// }
37+
is SimpleDataColumn -> {
4238
"${it.name}: ${it.type}"
4339
}
44-
else -> TODO()
4540
}
4641
"$indent$col"
4742
}
4843
}
4944

50-
open class SimpleCol(
51-
val name: String,
52-
open val type: TypeApproximation
53-
) : GenericColumn {
45+
sealed interface SimpleCol : GenericColumn {
46+
val name: String
5447

5548
override fun name(): String {
5649
return name
5750
}
5851

59-
open fun rename(s: String): SimpleCol {
60-
return SimpleCol(s, type)
61-
}
62-
63-
open fun changeType(type: TypeApproximation): SimpleCol {
64-
return SimpleCol(name, type)
65-
}
66-
67-
override fun equals(other: Any?): Boolean {
68-
if (this === other) return true
69-
if (javaClass != other?.javaClass) return false
70-
71-
other as SimpleCol
72-
73-
if (name != other.name) return false
74-
if (type != other.type) return false
52+
fun rename(s: String): SimpleCol
53+
}
7554

76-
return true
77-
}
55+
data class SimpleDataColumn(
56+
override val name: String,
57+
val type: TypeApproximation
58+
) : GenericColumn, SimpleCol {
7859

79-
override fun hashCode(): Int {
80-
var result = name.hashCode()
81-
result = 31 * result + type.hashCode()
82-
return result
60+
override fun name(): String {
61+
return name
8362
}
8463

85-
override fun toString(): String {
86-
return "SimpleCol(name='$name', type=$type)"
64+
override fun rename(s: String): SimpleDataColumn {
65+
return SimpleDataColumn(s, type)
8766
}
8867

89-
open fun kind(): SimpleColumnKind {
90-
return SimpleColumnKind.VALUE
68+
fun changeType(type: TypeApproximation): SimpleDataColumn {
69+
return SimpleDataColumn(name, type)
9170
}
92-
}
9371

94-
enum class SimpleColumnKind {
95-
VALUE, GROUP, FRAME
9672
}
9773

9874
data class SimpleFrameColumn(
99-
private val name1: String,
100-
private val columns: List<SimpleCol>,
101-
// probably shouldn't be called at all?
102-
// exists only because SimpleCol has it
103-
// but in fact it's for `materialize` to decide what should be the type of the property / accessors
104-
val anyFrameType: TypeApproximation,
105-
) : GenericColumnGroup<SimpleCol>, SimpleCol(name1, anyFrameType) {
75+
override val name: String,
76+
private val columns: List<SimpleCol>
77+
) : GenericColumnGroup<SimpleCol>, SimpleCol {
10678
override fun columns(): List<SimpleCol> {
10779
return columns
10880
}
10981

11082
override fun rename(s: String): SimpleFrameColumn {
111-
return SimpleFrameColumn(name1, columns, anyFrameType)
112-
}
113-
114-
override fun kind(): SimpleColumnKind {
115-
return SimpleColumnKind.FRAME
83+
return SimpleFrameColumn(s, columns)
11684
}
11785
}
11886

119-
class SimpleColumnGroup(
120-
name: String,
121-
private val columns: List<SimpleCol>,
122-
columnGroupType: TypeApproximation
123-
) : GenericColumnGroup<SimpleCol>, SimpleCol(name, columnGroupType) {
87+
data class SimpleColumnGroup(
88+
override val name: String,
89+
private val columns: List<SimpleCol>
90+
) : GenericColumnGroup<SimpleCol>, SimpleCol {
12491

12592
override fun columns(): List<SimpleCol> {
12693
return columns
12794
}
12895

12996
override fun rename(s: String): SimpleColumnGroup {
130-
return SimpleColumnGroup(s, columns, type)
131-
}
132-
133-
override fun changeType(type: TypeApproximation): SimpleCol {
134-
return TODO()
135-
}
136-
137-
override fun equals(other: Any?): Boolean {
138-
if (this === other) return true
139-
if (javaClass != other?.javaClass) return false
140-
if (!super.equals(other)) return false
141-
142-
other as SimpleColumnGroup
143-
144-
if (name != other.name) return false
145-
if (columns != other.columns) return false
146-
147-
return true
148-
}
149-
150-
override fun hashCode(): Int {
151-
var result = super.hashCode()
152-
result = 31 * result + name.hashCode()
153-
result = 31 * result + columns.hashCode()
154-
return result
155-
}
156-
157-
override fun kind(): SimpleColumnKind {
158-
return SimpleColumnKind.GROUP
97+
return SimpleColumnGroup(s, columns)
15998
}
16099
}
161100

162101
fun KotlinTypeFacade.simpleColumnOf(name: String, type: ConeKotlinType): SimpleCol {
163102
return if (type.classId == Names.DATA_ROW_CLASS_ID) {
164103
val schema = pluginDataFrameSchema(type)
165-
val group = SimpleColumnGroup(name, schema.columns(), anyRow)
104+
val group = SimpleColumnGroup(name, schema.columns())
166105
val column = if (type.isNullable) {
167106
makeNullable(group)
168107
} else {
@@ -171,18 +110,18 @@ fun KotlinTypeFacade.simpleColumnOf(name: String, type: ConeKotlinType): SimpleC
171110
column
172111
} else if (type.classId == Names.DF_CLASS_ID && type.nullability == ConeNullability.NOT_NULL) {
173112
val schema = pluginDataFrameSchema(type)
174-
SimpleFrameColumn(name, schema.columns(), anyDataFrame)
113+
SimpleFrameColumn(name, schema.columns())
175114
} else {
176-
SimpleCol(name, type.wrap())
115+
SimpleDataColumn(name, type.wrap())
177116
}
178117
}
179118

180119
private fun KotlinTypeFacade.makeNullable(column: SimpleCol): SimpleCol {
181120
return when (column) {
182121
is SimpleColumnGroup -> {
183-
SimpleColumnGroup(column.name, column.columns().map { makeNullable(column) }, anyRow)
122+
SimpleColumnGroup(column.name, column.columns().map { makeNullable(column) })
184123
}
185124
is SimpleFrameColumn -> column
186-
else -> SimpleCol(column.name, column.type.changeNullability { true })
125+
is SimpleDataColumn -> SimpleDataColumn(column.name, column.type.changeNullability { true })
187126
}
188127
}

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
77
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
88
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
99
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
10+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
1011
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
1112
import org.jetbrains.kotlinx.dataframe.plugin.impl.string
1213
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
@@ -19,7 +20,7 @@ class Add : AbstractSchemaModificationInterpreter() {
1920
val Arguments.type: TypeApproximation by type(name("expression"))
2021

2122
override fun Arguments.interpret(): PluginDataFrameSchema {
22-
return PluginDataFrameSchema(receiver.columns() + SimpleCol(name, type))
23+
return PluginDataFrameSchema(receiver.columns() + SimpleDataColumn(name, type))
2324
}
2425
}
2526

@@ -29,7 +30,7 @@ class From : AbstractInterpreter<Unit>() {
2930
val Arguments.type: TypeApproximation by type(name("expression"))
3031

3132
override fun Arguments.interpret() {
32-
dsl.columns += SimpleCol(receiver, type)
33+
dsl.columns += SimpleDataColumn(receiver, type)
3334
}
3435
}
3536

@@ -39,7 +40,7 @@ class Into : AbstractInterpreter<Unit>() {
3940
val Arguments.name: String by string()
4041

4142
override fun Arguments.interpret() {
42-
dsl.columns += SimpleCol(name, receiver)
43+
dsl.columns += SimpleDataColumn(name, receiver)
4344
}
4445
}
4546

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

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
1313
import org.jetbrains.kotlinx.dataframe.api.Infer
1414
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
1515
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
16+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
1617
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
17-
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnKind
1818
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
1919
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
2020
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
@@ -73,14 +73,14 @@ internal fun KotlinTypeFacade.convertImpl(
7373
type: TypeApproximation
7474
): PluginDataFrameSchema {
7575
return pluginDataFrameSchema.map(columns.toSet()) { path, column ->
76-
require(column.kind() == SimpleColumnKind.VALUE) {
77-
"$path must be ${SimpleColumnKind.VALUE}, but was ${column.kind()}"
76+
require(column is SimpleDataColumn) {
77+
"$path must be ${SimpleDataColumn::class}, but was ${column::class}"
7878
}
7979
val unwrappedType = type.type
8080
// TODO: AnyRow
8181
if (unwrappedType is ConeClassLikeType && unwrappedType.classId == Names.DF_CLASS_ID && !unwrappedType.isNullable) {
8282
val f = unwrappedType.typeArguments.single()
83-
SimpleFrameColumn(column.name, pluginDataFrameSchema(f).columns(), anyDataFrame)
83+
SimpleFrameColumn(column.name, pluginDataFrameSchema(f).columns())
8484
} else {
8585
column.changeType(type)
8686
}
@@ -111,7 +111,7 @@ internal fun f(columns: List<SimpleCol>, transform: ColumnMapper, selected: Colu
111111
} else {
112112
it.map(transform, selected, fullPath)
113113
}
114-
else -> if (fullPath in selected) {
114+
is SimpleDataColumn -> if (fullPath in selected) {
115115
transform(path, it)
116116
} else {
117117
it
@@ -123,16 +123,14 @@ internal fun f(columns: List<SimpleCol>, transform: ColumnMapper, selected: Colu
123123
internal fun SimpleColumnGroup.map(transform: ColumnMapper, selected: ColumnsSet, path: List<String>): SimpleColumnGroup {
124124
return SimpleColumnGroup(
125125
name,
126-
f(columns(), transform, selected, path),
127-
type
126+
f(columns(), transform, selected, path)
128127
)
129128
}
130129

131130
internal fun SimpleFrameColumn.map(transform: ColumnMapper, selected: ColumnsSet, path: List<String>): SimpleFrameColumn {
132131
return SimpleFrameColumn(
133132
name,
134-
f(columns(), transform, selected, path),
135-
anyFrameType
133+
f(columns(), transform, selected, path)
136134
)
137135
}
138136

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
1111
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
1212
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
1313
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
14-
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnKind
14+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
1515
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
1616
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
1717

@@ -31,7 +31,7 @@ fun KotlinTypeFacade.fillNullsImpl(
3131
): List<SimpleCol> {
3232
return columns.map {
3333
// else report?
34-
if (p + it.name() in paths && it.kind() == SimpleColumnKind.VALUE) {
34+
if (p + it.name() in paths && it is SimpleDataColumn) {
3535
val coneType = it.type.type as? ConeSimpleKotlinType
3636
if (coneType != null) {
3737
val type = coneType.withNullability(ConeNullability.NOT_NULL, session.typeContext)
@@ -43,7 +43,7 @@ fun KotlinTypeFacade.fillNullsImpl(
4343
} else {
4444
if (it is SimpleColumnGroup) {
4545
val updatedColumns = fillNullsImpl(it.columns(), paths, p + it.name())
46-
SimpleColumnGroup(it.name(), updatedColumns, anyRow)
46+
SimpleColumnGroup(it.name(), updatedColumns)
4747
} else {
4848
it
4949
}

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

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
66
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
77
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
88
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
9+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
910
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
1011
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
1112
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
@@ -32,9 +33,9 @@ val KotlinTypeFacade.explodeImpl: PluginDataFrameSchema.(dropEmpty: Boolean, sel
3233

3334
fun makeNullable(column: SimpleCol): SimpleCol {
3435
return when (column) {
35-
is SimpleColumnGroup -> SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) }, column.type)
36+
is SimpleColumnGroup -> SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
3637
is SimpleFrameColumn -> column
37-
else -> {
38+
is SimpleDataColumn -> {
3839
// val nullable = if (dropEmpty) (column.type as TypeApproximationImpl).nullable else true
3940

4041
column.changeType(type = column.type.changeNullability { nullable -> if (dropEmpty) nullable else true })
@@ -46,21 +47,21 @@ val KotlinTypeFacade.explodeImpl: PluginDataFrameSchema.(dropEmpty: Boolean, sel
4647
val fullPath = path + listOf(column.name)
4748
return when (column) {
4849
is SimpleColumnGroup -> {
49-
SimpleColumnGroup(column.name, column.columns().map { explode(it, fullPath) }, column.type)
50+
SimpleColumnGroup(column.name, column.columns().map { explode(it, fullPath) })
5051
}
5152
is SimpleFrameColumn -> {
5253
if (fullPath in selected) {
53-
SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) }, anyDataFrame)
54+
SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
5455
} else {
5556
column
5657
}
5758
}
58-
else -> if (fullPath in selected) {
59+
is SimpleDataColumn -> if (fullPath in selected) {
5960
val newType = when {
6061
column.type.isList() -> column.type.typeArgument()
6162
else -> column.type
6263
}
63-
SimpleCol(column.name, newType)
64+
SimpleDataColumn(column.name, newType)
6465
} else {
6566
column
6667
}

0 commit comments

Comments
 (0)