[LL FIR] fix contract violation exception from Java Resolution

When we create `SmartTypePointer` directly inside `JavaSymbolProvider`,
the Java resolution is called which is forbidden inside the
`JavaSymbolProvider`.

Instead, the lazy `JavaElementTypeSource` for JavaType is created,
which creates a `SmartTypePointer` outside `JavaSymbolProvider`.

^KT-58194

analysis/analysis-api-impl-base/src/org/jetbrains/kotlin/analysis/api/impl/base/java/source/JavaElementDelegatingTypeSourceWithSmartPointer.kt

@@ -0,0 +1,49 @@
+package org.jetbrains.kotlin.analysis.api.impl.base.java.source
+
+import com.intellij.psi.*
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementTypeSource
+
+internal abstract class JavaElementDelegatingTypeSourceWithSmartPointer<PSI : PsiElement, TYPE : PsiType> : JavaElementTypeSource<TYPE>() {
+    abstract val psiPointer: SmartPsiElementPointer<out PSI>
+    abstract fun getType(psi: PSI): TYPE
+
+    override val type: TYPE
+        get() {
+            val psi = psiPointer.element
+                ?: error("Cannot restore $psiPointer")
+            return getType(psi)
+        }
+}
+
+internal class JavaElementDelegatingVariableReturnTypeSourceWithSmartPointer<TYPE : PsiType>(
+    override val psiPointer: SmartPsiElementPointer<out PsiVariable>,
+    override val factory: JavaElementSourceFactory,
+) : JavaElementDelegatingTypeSourceWithSmartPointer<PsiVariable, TYPE>() {
+
+    override fun getType(psi: PsiVariable): TYPE {
+        @Suppress("UNCHECKED_CAST")
+        return psi.type as TYPE
+    }
+}
+
+internal class JavaElementDelegatingMethodReturnTypeSourceWithSmartPointer<TYPE : PsiType>(
+    override val psiPointer: SmartPsiElementPointer<out PsiMethod>,
+    override val factory: JavaElementSourceFactory,
+) : JavaElementDelegatingTypeSourceWithSmartPointer<PsiMethod, TYPE>() {
+
+    override fun getType(psi: PsiMethod): TYPE {
+        @Suppress("UNCHECKED_CAST")
+        return psi.returnType as TYPE
+    }
+}
+internal class JavaElementDelegatingExpressionTypeSourceWithSmartPointer<TYPE : PsiType>(
+    override val psiPointer: SmartPsiElementPointer<out PsiExpression>,
+    override val factory: JavaElementSourceFactory,
+) : JavaElementDelegatingTypeSourceWithSmartPointer<PsiExpression, TYPE>() {
+
+    override fun getType(psi: PsiExpression): TYPE {
+        @Suppress("UNCHECKED_CAST")
+        return psi.type as TYPE
+    }
+}

analysis/analysis-api-impl-base/src/org/jetbrains/kotlin/analysis/api/impl/base/java/source/JavaElementPsiSourceWithSmartPointer.kt

@@ -11,7 +11,7 @@ import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource
 import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
 
 internal class JavaElementPsiSourceWithSmartPointer<PSI : PsiElement>(
-    private val pointer: SmartPsiElementPointer<PSI>,
+    val pointer: SmartPsiElementPointer<PSI>,
     override val factory: JavaElementSourceFactory,
 ) : JavaElementPsiSource<PSI>() {
 
@@ -20,9 +20,5 @@ internal class JavaElementPsiSourceWithSmartPointer<PSI : PsiElement>(
             return pointer.element
                 ?: error("Cannot restore a PsiElement from $pointer")
         }
-
-    override fun toString(): String {
-        return pointer.element?.toString() ?: "Cannot restore a PsiElement from $pointer"
-    }
 }
 

analysis/analysis-api-impl-base/src/org/jetbrains/kotlin/analysis/api/impl/base/java/source/JavaElementSourceWithSmartPointerFactory.kt

@@ -6,15 +6,12 @@
 package org.jetbrains.kotlin.analysis.api.impl.base.java.source
 
 import com.intellij.openapi.project.Project
-import com.intellij.psi.PsiElement
+import com.intellij.psi.*
-import com.intellij.psi.PsiType
-import com.intellij.psi.SmartPointerManager
-import com.intellij.psi.SmartTypePointerManager
 import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource
 import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
 import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementTypeSource
 
- class JavaElementSourceWithSmartPointerFactory(project: Project) : JavaElementSourceFactory() {
+class JavaElementSourceWithSmartPointerFactory(project: Project) : JavaElementSourceFactory() {
     private val smartTypePointerManager = SmartTypePointerManager.getInstance(project)
     private val smartPsiPointerManager = SmartPointerManager.getInstance(project)
 
@@ -25,4 +22,19 @@ import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementTypeSourc
     override fun <TYPE : PsiType> createTypeSource(type: TYPE): JavaElementTypeSource<TYPE> {
         return JavaElementTypeSourceWithSmartPointer(smartTypePointerManager.createSmartTypePointer(type), this)
     }
+
+    override fun <TYPE : PsiType> createVariableReturnTypeSource(psiVariableSource: JavaElementPsiSource<out PsiVariable>): JavaElementTypeSource<TYPE> {
+        require(psiVariableSource is JavaElementPsiSourceWithSmartPointer)
+        return JavaElementDelegatingVariableReturnTypeSourceWithSmartPointer(psiVariableSource.pointer, psiVariableSource.factory)
+    }
+
+    override fun <TYPE : PsiType> createMethodReturnTypeSource(psiMethodSource: JavaElementPsiSource<out PsiMethod>): JavaElementTypeSource<TYPE> {
+        require(psiMethodSource is JavaElementPsiSourceWithSmartPointer)
+        return JavaElementDelegatingMethodReturnTypeSourceWithSmartPointer(psiMethodSource.pointer, psiMethodSource.factory)
+    }
+
+    override fun <TYPE : PsiType> createExpressionTypeSource(psiExpressionSource: JavaElementPsiSource<out PsiExpression>): JavaElementTypeSource<TYPE> {
+        require(psiExpressionSource is JavaElementPsiSourceWithSmartPointer)
+        return JavaElementDelegatingExpressionTypeSourceWithSmartPointer(psiExpressionSource.pointer, psiExpressionSource.factory)
+    }
 }

compiler/resolution.common.jvm/src/org/jetbrains/kotlin/load/java/structure/impl/JavaMethodImpl.java

@@ -19,7 +19,6 @@ package org.jetbrains.kotlin.load.java.structure.impl;
 import com.intellij.psi.PsiAnnotationMemberValue;
 import com.intellij.psi.PsiAnnotationMethod;
 import com.intellij.psi.PsiMethod;
-import com.intellij.psi.PsiType;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.load.java.structure.*;
@@ -27,6 +26,7 @@ import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource
 import org.jetbrains.kotlin.name.Name;
 
 import java.util.List;
+import java.util.Objects;
 
 import static org.jetbrains.kotlin.load.java.structure.impl.JavaElementCollectionFromPsiArrayUtil.typeParameters;
 import static org.jetbrains.kotlin.load.java.structure.impl.JavaElementCollectionFromPsiArrayUtil.valueParameters;
@@ -80,8 +80,6 @@ public class JavaMethodImpl extends JavaMemberImpl<PsiMethod> implements JavaMet
     @Override
     @NotNull
     public JavaType getReturnType() {
-        PsiType psiType = getPsi().getReturnType();
+        return JavaTypeImpl.create(Objects.requireNonNull(getPsi().getReturnType()), createMethodReturnTypeSource(psiElementSource));
-        assert psiType != null : "Method is not a constructor and has no return type: " + getName();
-        return JavaTypeImpl.create(createTypeSource(psiType));
     }
 }

compiler/resolution.common.jvm/src/org/jetbrains/kotlin/load/java/structure/impl/JavaRecordComponentImpl.kt

@@ -14,7 +14,7 @@ class JavaRecordComponentImpl(
     psiRecordComponentSource: JavaElementPsiSource<PsiRecordComponent>
 ) : JavaMemberImpl<PsiRecordComponent>(psiRecordComponentSource), JavaRecordComponent {
     override val type: JavaType
-        get() = JavaTypeImpl.create(sourceFactory.createTypeSource(psiElementSource.psi.type))
+        get() = JavaTypeImpl.create(psi.type, sourceFactory.createVariableReturnTypeSource(psiElementSource))
 
     override val isVararg: Boolean
         get() = psi.isVarArgs

compiler/resolution.common.jvm/src/org/jetbrains/kotlin/load/java/structure/impl/JavaTypeImpl.java

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.load.java.structure.impl.source.JavaSourceFactoryOwn
 import org.jetbrains.kotlin.name.FqName;
 
 import java.util.Collection;
+import java.util.function.Function;
 
 public abstract class JavaTypeImpl<Psi extends PsiType> implements JavaType, JavaAnnotationOwnerImpl, JavaSourceFactoryOwner {
     private final JavaElementTypeSource<Psi> psiType;
@@ -54,9 +55,13 @@ public abstract class JavaTypeImpl<Psi extends PsiType> implements JavaType, Jav
     }
 
     @NotNull
-    public static JavaTypeImpl<?> create(@NotNull JavaElementTypeSource<PsiType> psiTypeSource) {
+    public static JavaTypeImpl<?> create(JavaElementTypeSource<? extends PsiType> psiTypeSource) {
-        JavaElementSourceFactory sourceFactory = psiTypeSource.getFactory();
+        return create(psiTypeSource.getType(), psiTypeSource);
-        return psiTypeSource.getType().accept(new PsiTypeVisitor<JavaTypeImpl<?>>() {
+    }
+
+    @NotNull
+    public static JavaTypeImpl<?> create(@NotNull PsiType psiType, JavaElementTypeSource<? extends PsiType> psiTypeSource) {
+        return psiType.accept(new PsiTypeVisitor<JavaTypeImpl<?>>() {
 
             @Nullable
             @Override
@@ -66,26 +71,30 @@ public abstract class JavaTypeImpl<Psi extends PsiType> implements JavaType, Jav
 
             @Nullable
             @Override
+            @SuppressWarnings("unchecked")
             public JavaTypeImpl<?> visitPrimitiveType(@NotNull PsiPrimitiveType primitiveType) {
-                return new JavaPrimitiveTypeImpl(sourceFactory.createTypeSource(primitiveType));
+                return new JavaPrimitiveTypeImpl((JavaElementTypeSource<PsiPrimitiveType>)psiTypeSource );
             }
 
             @Nullable
             @Override
+            @SuppressWarnings("unchecked")
             public JavaTypeImpl<?> visitArrayType(@NotNull PsiArrayType arrayType) {
-                return new JavaArrayTypeImpl(sourceFactory.createTypeSource(arrayType));
+                return new JavaArrayTypeImpl((JavaElementTypeSource<PsiArrayType>) psiTypeSource);
             }
 
             @Nullable
             @Override
+            @SuppressWarnings("unchecked")
             public JavaTypeImpl<?> visitClassType(@NotNull PsiClassType classType) {
-                return new JavaClassifierTypeImpl(sourceFactory.createTypeSource(classType));
+                return new JavaClassifierTypeImpl((JavaElementTypeSource<PsiClassType>) psiTypeSource);
             }
 
             @Nullable
             @Override
+            @SuppressWarnings("unchecked")
             public JavaTypeImpl<?> visitWildcardType(@NotNull PsiWildcardType wildcardType) {
-                return new JavaWildcardTypeImpl(sourceFactory.createTypeSource(wildcardType));
+                return new JavaWildcardTypeImpl((JavaElementTypeSource<PsiWildcardType>) psiTypeSource);
             }
         });
     }

compiler/resolution.common.jvm/src/org/jetbrains/kotlin/load/java/structure/impl/JavaValueParameterImpl.java

@@ -102,7 +102,7 @@ public class JavaValueParameterImpl extends JavaElementImpl<PsiParameter>
     @Override
     @NotNull
     public JavaType getType() {
-        return JavaTypeImpl.create(createTypeSource(getPsi().getType()));
+        return JavaTypeImpl.create(getPsi().getType(), createVariableReturnTypeSource(psiElementSource));
     }
 
     @Override

compiler/resolution.common.jvm/src/org/jetbrains/kotlin/load/java/structure/impl/source/JavaElementSourceFactory.kt

@@ -6,12 +6,15 @@
 package org.jetbrains.kotlin.load.java.structure.impl.source
 
 import com.intellij.openapi.project.Project
-import com.intellij.psi.PsiElement
+import com.intellij.psi.*
-import com.intellij.psi.PsiType
 
 abstract class JavaElementSourceFactory {
     abstract fun <PSI : PsiElement> createPsiSource(psi: PSI): JavaElementPsiSource<PSI>
     abstract fun <TYPE : PsiType> createTypeSource(type: TYPE): JavaElementTypeSource<TYPE>
+    abstract fun <TYPE : PsiType> createVariableReturnTypeSource(psiVariableSource: JavaElementPsiSource<out PsiVariable>): JavaElementTypeSource<TYPE>
+    abstract fun <TYPE : PsiType> createMethodReturnTypeSource(psiMethodSource: JavaElementPsiSource<out PsiMethod>): JavaElementTypeSource<TYPE>
+
+    abstract fun <TYPE : PsiType> createExpressionTypeSource(psiExpressionSource: JavaElementPsiSource<out PsiExpression>): JavaElementTypeSource<TYPE>
 
     companion object {
         @JvmStatic
@@ -29,4 +32,22 @@ class JavaFixedElementSourceFactory : JavaElementSourceFactory() {
     override fun <TYPE : PsiType> createTypeSource(type: TYPE): JavaElementTypeSource<TYPE> {
         return JavaElementTypeSourceWithFixedType(type, this)
     }
+
+    override fun <TYPE : PsiType> createVariableReturnTypeSource(psiVariableSource: JavaElementPsiSource<out PsiVariable>): JavaElementTypeSource<TYPE> {
+        @Suppress("UNCHECKED_CAST")
+        return createTypeSource(psiVariableSource.psi.type as TYPE)
+    }
+
+
+    override fun <TYPE : PsiType> createExpressionTypeSource(psiExpressionSource: JavaElementPsiSource<out PsiExpression>): JavaElementTypeSource<TYPE> {
+        @Suppress("UNCHECKED_CAST")
+        return createTypeSource(psiExpressionSource.psi.type as TYPE)
+    }
+
+    override fun <TYPE : PsiType> createMethodReturnTypeSource(psiMethodSource: JavaElementPsiSource<out PsiMethod>): JavaElementTypeSource<TYPE> {
+        val psiType: PsiType = psiMethodSource.psi.returnType
+            ?: error("Method is not a constructor and has no return type: " + psiMethodSource.psi.name)
+        @Suppress("UNCHECKED_CAST")
+        return createTypeSource(psiType as TYPE)
+    }
 }

compiler/resolution.common.jvm/src/org/jetbrains/kotlin/load/java/structure/impl/source/JavaSourceFactoryOwner.kt

@@ -6,7 +6,9 @@
 package org.jetbrains.kotlin.load.java.structure.impl.source
 
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiType
+import com.intellij.psi.PsiVariable
 
 interface JavaSourceFactoryOwner {
     val sourceFactory: JavaElementSourceFactory
@@ -18,4 +20,13 @@ interface JavaSourceFactoryOwner {
     fun <TYPE : PsiType> createTypeSource(type: TYPE): JavaElementTypeSource<TYPE> {
         return sourceFactory.createTypeSource(type)
     }
+
+    fun <TYPE : PsiType> createVariableReturnTypeSource(psiVariableSource: JavaElementPsiSource<out PsiVariable>): JavaElementTypeSource<TYPE> {
+        return sourceFactory.createVariableReturnTypeSource(psiVariableSource)
+    }
+
+    fun <TYPE : PsiType> createMethodReturnTypeSource(psiMethodSource: JavaElementPsiSource<out PsiMethod>): JavaElementTypeSource<TYPE> {
+        return sourceFactory.createMethodReturnTypeSource(psiMethodSource)
+    }
+
 }