New J2K: Assignment expression conversion improvements

j2k/newSrc/org/jetbrains/kotlin/j2k/JKSymbolProvider.kt

@@ -12,10 +12,12 @@ import org.jetbrains.kotlin.j2k.tree.JKField
 import org.jetbrains.kotlin.j2k.tree.JKMethod
 import org.jetbrains.kotlin.j2k.tree.impl.*
 import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtNamedFunction
 
 
 class JKSymbolProvider {
-    val symbols = mutableMapOf<PsiElement, JKSymbol>()
+    val symbolsByPsi = mutableMapOf<PsiElement, JKSymbol>()
+    val symbolsByJK = mutableMapOf<JKDeclaration, JKSymbol>()
     private val elementVisitor = ElementVisitor()
 
     fun preBuildTree(files: List<PsiJavaFile>) {
@@ -23,12 +25,13 @@ class JKSymbolProvider {
     }
 
     fun provideDirectSymbol(psi: PsiElement): JKSymbol {
-        return symbols.getOrPut(psi) {
+        return symbolsByPsi.getOrPut(psi) {
             when (psi) {
                 is PsiClass -> JKMultiverseClassSymbol(psi)
                 is KtClassOrObject -> JKMultiverseKtClassSymbol(psi)
                 is PsiMethod -> JKMultiverseMethodSymbol(psi)
                 is PsiField -> JKMultiverseFieldSymbol(psi)
+                is KtNamedFunction -> JKMultiverseFunctionSymbol(psi)
                 else -> TODO(psi::class.toString())
             }
         }
@@ -40,22 +43,36 @@ class JKSymbolProvider {
         return JKUnresolvedField(reference).let { if (it is T) it else JKUnresolvedMethod(reference) as T }
     }
 
-    fun provideUniverseSymbol(psi: PsiElement, jk: JKDeclaration? = null): JKSymbol = symbols.getOrPut(psi) {
+    fun provideUniverseSymbol(psi: PsiElement, jk: JKDeclaration): JKSymbol = provideUniverseSymbol(psi).also {
+        when (it) {
+            is JKUniverseClassSymbol -> it.target = jk as JKClass
+            is JKUniverseFieldSymbol -> it.target = jk as JKField
+            is JKUniverseMethodSymbol -> it.target = jk as JKMethod
+        }
+        symbolsByJK[jk] = it
+    }
+
+    fun provideUniverseSymbol(psi: PsiElement): JKSymbol = symbolsByPsi.getOrPut(psi) {
         when (psi) {
             is PsiField, is PsiParameter, is PsiLocalVariable -> JKUniverseFieldSymbol()
             is PsiMethod -> JKUniverseMethodSymbol()
             is PsiClass -> JKUniverseClassSymbol()
             else -> TODO()
         }
-    }.also {
-        if (jk != null)
-            when (it) {
-                is JKUniverseClassSymbol -> it.target = jk as JKClass
-                is JKUniverseFieldSymbol -> it.target = jk as JKField
-                is JKUniverseMethodSymbol -> it.target = jk as JKMethod
-            }
     }
 
+    fun provideUniverseSymbol(jk: JKClass): JKClassSymbol = symbolsByJK.getOrPut(jk) {
+        JKUniverseClassSymbol().also { it.target = jk }
+    } as JKClassSymbol
+
+    fun provideUniverseSymbol(jk: JKField): JKFieldSymbol = symbolsByJK.getOrPut(jk) {
+        JKUniverseFieldSymbol().also { it.target = jk }
+    } as JKFieldSymbol
+
+    fun provideUniverseSymbol(jk: JKMethod): JKMethodSymbol = symbolsByJK.getOrPut(jk) {
+        JKUniverseMethodSymbol().also { it.target = jk }
+    } as JKMethodSymbol
+
     private inner class ElementVisitor : JavaElementVisitor() {
         override fun visitClass(aClass: PsiClass) {
             provideUniverseSymbol(aClass)

j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/AssignmentAsExpressionToAlsoConversion.kt

@@ -6,28 +6,31 @@
 package org.jetbrains.kotlin.j2k.conversions
 
 import org.jetbrains.kotlin.j2k.ConversionContext
-import org.jetbrains.kotlin.j2k.tree.JKExpressionStatement
-import org.jetbrains.kotlin.j2k.tree.JKJavaAssignmentExpression
-import org.jetbrains.kotlin.j2k.tree.JKTreeElement
+import org.jetbrains.kotlin.j2k.tree.*
 import org.jetbrains.kotlin.j2k.tree.impl.*
 import org.jetbrains.kotlin.name.ClassId
 
 class AssignmentAsExpressionToAlsoConversion(val context: ConversionContext) : RecursiveApplicableConversionBase() {
-
     override fun applyToElement(element: JKTreeElement): JKTreeElement {
-        if (element !is JKJavaAssignmentExpression || element.parent !is JKExpressionStatement) return recurse(element)
+        if (element !is JKJavaAssignmentExpression) return recurse(element)
 
+        val alsoSymbol = context.symbolProvider.provideDirectSymbol(
+            resolveFqName(ClassId.fromString("kotlin/also"), element.getParentOfType<JKClass>() ?: element, context)!!
+        ) as JKMethodSymbol
 
-        val alsoElement = resolveFqName(ClassId.fromString("kotlin/also"), element, context) ?: return recurse(element)
-        val alsoSymbol = context.symbolProvider.provideDirectSymbol(alsoElement) as? JKMethodSymbol ?: return recurse(element)
+        val alsoExpression = JKKtAlsoCallExpressionImpl(
+            JKKtAssignmentStatementImpl(element.field, JKStubExpressionImpl(), element.operator), alsoSymbol
+        ).also {
+            (it.statement as JKKtAssignmentStatement).expression =
+                    JKKtFieldAccessExpressionImpl(context.symbolProvider.provideUniverseSymbol(it.parameter))
+        }
         element.invalidate()
 
         return JKQualifiedExpressionImpl(
             element.expression,
             JKKtQualifierImpl.DOT,
-            JKKtCallExpressionImpl(alsoSymbol, JKExpressionListImpl(
-                // TODO: Lambda expression here
-            ))
+            alsoExpression
         )
+
     }
 }

j2k/newSrc/org/jetbrains/kotlin/j2k/tree/impl/k.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.j2k.tree.*
 import org.jetbrains.kotlin.j2k.tree.visitors.JKVisitor
 import org.jetbrains.kotlin.lexer.KtSingleValueToken
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.utils.addToStdlib.cast
 
 class JKKtPropertyImpl(
     modifierList: JKModifierList,
@@ -129,13 +130,11 @@ class JKKtModifierImpl(override val type: JKKtModifier.KtModifierType) : JKKtMod
 }
 
 class JKKtAlsoCallExpressionImpl(
-    override val statement: JKStatement,
-    override val parameterName: String = "it"
+    override val statement: JKStatement, override val identifier: JKMethodSymbol, override val parameterName: String = "it"
 ) : JKKtAlsoCallExpression, JKBranchElementBase() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtAlsoCallExpression(this, data)
-
-    override val identifier
-        get() = TODO()
+    val parameter
+        get() = arguments.expressions.first().cast<JKLambdaExpression>().parameters.first()
     override var arguments: JKExpressionList by child(
         JKExpressionListImpl(
             listOf(
@@ -151,4 +150,12 @@ class JKKtAlsoCallExpressionImpl(
             )
         )
     )
+}
+
+class JKKtAssignmentStatementImpl(
+    override var field: JKAssignableExpression, expression: JKExpression, override var operator: JKOperator
+) : JKKtAssignmentStatement, JKBranchElementBase() {
+    override var expression by child(expression)
+    override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtAssignmentStatement(this, data)
+
 }

j2k/newSrc/org/jetbrains/kotlin/j2k/tree/impl/symbols.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.j2k.tree.JKClass
 import org.jetbrains.kotlin.j2k.tree.JKField
 import org.jetbrains.kotlin.j2k.tree.JKMethod
 import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtNamedFunction
 
 interface JKSymbol {
     val target: Any
@@ -70,11 +71,14 @@ class JKMultiverseMethodSymbol(override val target: PsiMethod) : JKMethodSymbol
 
 }
 
-class JKUniverseFieldSymbol() : JKFieldSymbol {
-    constructor(target: JKField) : this() {
-        this.target = target
-    }
+class JKMultiverseFunctionSymbol(override val target: KtNamedFunction) : JKMethodSymbol {
+    override val declaredIn: JKSymbol
+        get() = TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.
+    override val fqName: String
+        get() = target.name!!
+}
 
+class JKUniverseFieldSymbol : JKFieldSymbol {
     override lateinit var target: JKField
     override val declaredIn: JKSymbol
         get() = TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.

j2k/newSrc/org/jetbrains/kotlin/j2k/tree/markers.kt

@@ -64,4 +64,13 @@ interface JKJavaPrimitiveType : JKType {
 
 interface JKJavaArrayType : JKType {
     val type: JKType
-}
+}
+
+inline fun <reified T> JKElement.getParentOfType(): T? {
+    var p = parent
+    while (true) {
+        if (p is T || p == null)
+            return p as? T
+        p = p.parent
+    }
+}