[Java resolution] make it possible to provide custom source element for JavaElement

This is needed to allow using PSI pointers in IDE in JavaElement.
`JavaElement`s are reused between read actions,
so underlying PSI elements might be invalidated when using hard PSI references

^KT-58194

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirPsiJavaClassSymbol.kt

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.fir.java.modality
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.load.java.structure.JavaClass
 import org.jetbrains.kotlin.load.java.structure.impl.JavaClassImpl
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
@@ -46,7 +47,7 @@ internal class KtFirPsiJavaClassSymbol(
     /**
      * [javaClass] is used to defer some properties to the compiler's view of a Java class.
      */
-    private val javaClass: JavaClass = JavaClassImpl(psi)
+    private val javaClass: JavaClass = JavaClassImpl(JavaElementSourceFactory.getInstance(analysisSession.project).createPsiSource(psi))
 
     override val name: Name = withValidityAssertion { javaClass.name }
 

compiler/cli/cli-base/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCliJavaFileManagerImpl.kt

@@ -35,12 +35,12 @@ import org.jetbrains.kotlin.load.java.structure.impl.classFiles.BinaryClassSigna
 import org.jetbrains.kotlin.load.java.structure.impl.classFiles.BinaryJavaClass
 import org.jetbrains.kotlin.load.java.structure.impl.classFiles.ClassifierResolutionContext
 import org.jetbrains.kotlin.load.java.structure.impl.classFiles.isNotTopLevelClass
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.resolve.jvm.KotlinCliJavaFileManager
 import org.jetbrains.kotlin.util.PerformanceCounter
 import org.jetbrains.kotlin.utils.addIfNotNull
-import java.util.*
 
 // TODO: do not inherit from CoreJavaFileManager to avoid accidental usage of its methods which do not use caches/indices
 // Currently, the only relevant usage of this class as CoreJavaFileManager is at CoreJavaDirectoryService.getPackage,
@@ -136,7 +136,14 @@ class KotlinCliJavaFileManagerImpl(private val myPsiManager: PsiManager) : CoreJ
             }
         }
 
-        return virtualFile.findPsiClassInVirtualFile(classId.relativeClassName.asString())?.let(::JavaClassImpl)
+        return virtualFile.findPsiClassInVirtualFile(classId.relativeClassName.asString())
+            ?.let { createJavaClassByPsiClass(it) }
+    }
+
+    private fun createJavaClassByPsiClass(psiClass: PsiClass): JavaClassImpl {
+        val project = myPsiManager.project
+        val sourceFactory = JavaElementSourceFactory.getInstance(project)
+        return JavaClassImpl(sourceFactory.createPsiSource(psiClass))
     }
 
     // this method is called from IDEA to resolve dependencies in Java code

compiler/cli/cli-base/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt

@@ -19,7 +19,6 @@ import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.TransactionGuard
 import com.intellij.openapi.application.TransactionGuardImpl
-import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.extensions.ExtensionsArea
 import com.intellij.openapi.fileTypes.PlainTextFileType
@@ -80,6 +79,8 @@ import org.jetbrains.kotlin.extensions.internal.TypeResolutionInterceptor
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
 import org.jetbrains.kotlin.idea.KotlinFileType
 import org.jetbrains.kotlin.js.translate.extensions.JsSyntheticTranslateExtension
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaFixedElementSourceFactory
 import org.jetbrains.kotlin.load.kotlin.KotlinBinaryClassCache
 import org.jetbrains.kotlin.load.kotlin.MetadataFinderFactory
 import org.jetbrains.kotlin.load.kotlin.ModuleVisibilityManager
@@ -747,6 +748,7 @@ class KotlinCoreEnvironment private constructor(
         @JvmStatic
         fun registerProjectServices(project: MockProject) {
             with(project) {
+                registerService(JavaElementSourceFactory::class.java, JavaFixedElementSourceFactory::class.java)
                 registerService(KotlinJavaPsiFacade::class.java, KotlinJavaPsiFacade(this))
                 registerService(ModuleAnnotationsResolver::class.java, CliModuleAnnotationsResolver())
             }

compiler/frontend.java/src/org/jetbrains/kotlin/load/java/JavaClassFinderImpl.kt

@@ -17,10 +17,12 @@
 package org.jetbrains.kotlin.load.java
 
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiPackage
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.load.java.structure.JavaClass
 import org.jetbrains.kotlin.load.java.structure.JavaPackage
 import org.jetbrains.kotlin.load.java.structure.impl.JavaPackageImpl
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.resolve.jvm.KotlinJavaPsiFacade
 import javax.inject.Inject
@@ -49,7 +51,17 @@ class JavaClassFinderImpl : AbstractJavaClassFinder() {
     }
 
     override fun findPackage(fqName: FqName, mayHaveAnnotations: Boolean): JavaPackage? {
-        return javaFacade.findPackage(fqName.asString(), javaSearchScope)?.let { JavaPackageImpl(it, javaSearchScope, mayHaveAnnotations) }
+        return javaFacade.findPackage(fqName.asString(), javaSearchScope)
+            ?.let { createJavaPackage(it, mayHaveAnnotations) }
+    }
+
+    private fun createJavaPackage(
+        psiPackage: PsiPackage,
+        mayHaveAnnotations: Boolean,
+    ): JavaPackageImpl {
+        val project = javaFacade.project
+        val sourceFactory = JavaElementSourceFactory.getInstance(project)
+        return JavaPackageImpl(sourceFactory.createPsiSource(psiPackage), javaSearchScope, mayHaveAnnotations)
     }
 
     override fun knownClassNamesInPackage(packageFqName: FqName): Set<String>? {

compiler/frontend.java/src/org/jetbrains/kotlin/resolve/jvm/modules/JavaModuleInfo.kt

@@ -27,10 +27,11 @@ import org.jetbrains.kotlin.load.java.structure.impl.classFiles.BinaryClassSigna
 import org.jetbrains.kotlin.load.java.structure.impl.classFiles.BinaryJavaAnnotation
 import org.jetbrains.kotlin.load.java.structure.impl.classFiles.ClassifierResolutionContext
 import org.jetbrains.kotlin.load.java.structure.impl.convert
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.resolve.jvm.KotlinCliJavaFileManager
-import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.kotlin.utils.compact
+import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.org.objectweb.asm.Opcodes.ACC_TRANSITIVE
 import java.io.IOException
 
@@ -59,7 +60,11 @@ class JavaModuleInfo(
                     Exports(FqName(packageName), statement.moduleNames)
                 }
             },
-            psiJavaModule.annotations.convert(::JavaAnnotationImpl)
+            psiJavaModule.annotations.convert {
+                JavaAnnotationImpl(
+                    JavaElementSourceFactory.getInstance(psiJavaModule.project).createPsiSource(it)
+                )
+            }
         )
 
         fun read(file: VirtualFile, javaFileManager: KotlinCliJavaFileManager, searchScope: GlobalSearchScope): JavaModuleInfo? {

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

@@ -26,6 +26,7 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.load.java.structure.JavaAnnotation;
 import org.jetbrains.kotlin.load.java.structure.JavaAnnotationArgument;
 import org.jetbrains.kotlin.load.java.structure.JavaClass;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource;
 import org.jetbrains.kotlin.name.ClassId;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.name.Name;
@@ -36,14 +37,14 @@ import java.util.Objects;
 import static org.jetbrains.kotlin.load.java.structure.impl.JavaElementCollectionFromPsiArrayUtil.namedAnnotationArguments;
 
 public class JavaAnnotationImpl extends JavaElementImpl<PsiAnnotation> implements JavaAnnotation {
-    public JavaAnnotationImpl(@NotNull PsiAnnotation psiAnnotation) {
+    public JavaAnnotationImpl(@NotNull JavaElementPsiSource<PsiAnnotation> psiAnnotation) {
         super(psiAnnotation);
     }
 
     @Override
     @NotNull
     public Collection<JavaAnnotationArgument> getArguments() {
-        return namedAnnotationArguments(getPsi().getParameterList().getAttributes());
+        return namedAnnotationArguments(getPsi().getParameterList().getAttributes(), getSourceFactory());
     }
 
     @Override
@@ -64,7 +65,7 @@ public class JavaAnnotationImpl extends JavaElementImpl<PsiAnnotation> implement
     @Override
     public JavaClass resolve() {
         PsiClass resolved = resolvePsi();
-        return resolved == null ? null : new JavaClassImpl(resolved);
+        return resolved == null ? null : new JavaClassImpl(createPsiSource(resolved));
     }
 
     @Nullable

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

@@ -19,15 +19,16 @@ package org.jetbrains.kotlin.load.java.structure.impl;
 import com.intellij.psi.PsiArrayType;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.load.java.structure.JavaArrayType;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementTypeSource;
 
 public class JavaArrayTypeImpl extends JavaTypeImpl<PsiArrayType> implements JavaArrayType {
-    public JavaArrayTypeImpl(@NotNull PsiArrayType psiArrayType) {
+    public JavaArrayTypeImpl(@NotNull JavaElementTypeSource<PsiArrayType> psiArrayTypeSource) {
-        super(psiArrayType);
+        super(psiArrayTypeSource);
     }
 
     @Override
     @NotNull
     public JavaTypeImpl<?> getComponentType() {
-        return JavaTypeImpl.create(getPsi().getComponentType());
+        return JavaTypeImpl.create(createTypeSource(getPsi().getComponentType()));
     }
 }

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

@@ -25,20 +25,21 @@ import org.jetbrains.kotlin.asJava.KtLightClassMarker
 import org.jetbrains.kotlin.asJava.isSyntheticValuesOrValueOfMethod
 import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.load.java.structure.*
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtPsiUtil
 
-class JavaClassImpl(psiClass: PsiClass) : JavaClassifierImpl<PsiClass>(psiClass), VirtualFileBoundJavaClass, JavaAnnotationOwnerImpl, JavaModifierListOwnerImpl {
+class JavaClassImpl(psiClassSource: JavaElementPsiSource<PsiClass>) : JavaClassifierImpl<PsiClass>(psiClassSource), VirtualFileBoundJavaClass, JavaAnnotationOwnerImpl, JavaModifierListOwnerImpl {
     init {
-        assert(psiClass !is PsiTypeParameter) { "PsiTypeParameter should be wrapped in JavaTypeParameter, not JavaClass: use JavaClassifier.create()" }
+        assert(psiClassSource.psi !is PsiTypeParameter) { "PsiTypeParameter should be wrapped in JavaTypeParameter, not JavaClass: use JavaClassifier.create()" }
     }
 
     override val innerClassNames: Collection<Name>
         get() = psi.innerClasses.mapNotNull { it.name?.takeIf(Name::isValidIdentifier)?.let(Name::identifier) }
 
     override fun findInnerClass(name: Name): JavaClass? {
-        return psi.findInnerClassByName(name.asString(), false)?.let(::JavaClassImpl)
+        return psi.findInnerClassByName(name.asString(), false)?.let { JavaClassImpl(psiElementSource.factory.createPsiSource(it)) }
     }
 
     override val fqName: FqName?
@@ -63,7 +64,7 @@ class JavaClassImpl(psiClass: PsiClass) : JavaClassifierImpl<PsiClass>(psiClass)
         get() = JavaElementUtil.isSealed(this)
 
     override val permittedTypes: Collection<JavaClassifierType>
-        get() = classifierTypes(psi.permitsListTypes)
+        get() = classifierTypes(psi.permitsListTypes, sourceFactory)
 
     override val isRecord: Boolean
         get() = psi.isRecord
@@ -71,14 +72,14 @@ class JavaClassImpl(psiClass: PsiClass) : JavaClassifierImpl<PsiClass>(psiClass)
     override val outerClass: JavaClassImpl?
         get() {
             val outer = psi.containingClass
-            return if (outer == null) null else JavaClassImpl(outer)
+            return if (outer == null) null else JavaClassImpl(psiElementSource.factory.createPsiSource(outer))
         }
 
     override val typeParameters: List<JavaTypeParameter>
-        get() = typeParameters(psi.typeParameters)
+        get() = typeParameters(psi.typeParameters, sourceFactory)
 
     override val supertypes: Collection<JavaClassifierType>
-        get() = classifierTypes(psi.superTypes)
+        get() = classifierTypes(psi.superTypes, sourceFactory)
 
     override val methods: Collection<JavaMethod>
         get() {
@@ -88,7 +89,8 @@ class JavaClassImpl(psiClass: PsiClass) : JavaClassifierImpl<PsiClass>(psiClass)
             return methods(
                 psi.methods.filter { method ->
                     !method.isConstructor && method.returnType != null && !(isEnum && isSyntheticValuesOrValueOfMethod(method))
-                }
+                },
+                sourceFactory,
             ).distinct()
         }
 
@@ -98,7 +100,7 @@ class JavaClassImpl(psiClass: PsiClass) : JavaClassifierImpl<PsiClass>(psiClass)
             return fields(psi.fields.filter {
                 // ex. Android plugin generates LightFields for resources started from '.' (.DS_Store file etc)
                 Name.isValidIdentifier(it.name)
-            })
+            }, sourceFactory)
         }
 
     override val constructors: Collection<JavaConstructor>
@@ -106,14 +108,14 @@ class JavaClassImpl(psiClass: PsiClass) : JavaClassifierImpl<PsiClass>(psiClass)
             assertNotLightClass()
             // See for example org.jetbrains.plugins.scala.lang.psi.light.ScFunctionWrapper,
             // which is present in getConstructors(), but its isConstructor() returns false
-            return constructors(psi.constructors.filter { method -> method.isConstructor })
+            return constructors(psi.constructors.filter { method -> method.isConstructor }, sourceFactory)
         }
 
     override val recordComponents: Collection<JavaRecordComponent>
         get() {
             assertNotLightClass()
 
-            return psi.recordComponents.convert(::JavaRecordComponentImpl)
+            return psi.recordComponents.convert { JavaRecordComponentImpl(psiElementSource.factory.createPsiSource(it)) }
         }
 
     override fun hasDefaultConstructor() = !isInterface && constructors.isEmpty()

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

@@ -22,22 +22,24 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.load.java.structure.JavaAnnotation;
 import org.jetbrains.kotlin.load.java.structure.JavaClassifier;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory;
 import org.jetbrains.kotlin.name.FqName;
 
 import java.util.Collection;
 
 public abstract class JavaClassifierImpl<Psi extends PsiClass> extends JavaElementImpl<Psi> implements JavaClassifier, JavaAnnotationOwnerImpl {
-    protected JavaClassifierImpl(@NotNull Psi psiClass) {
+    protected JavaClassifierImpl(@NotNull JavaElementPsiSource<Psi> psi) {
-        super(psiClass);
+        super(psi);
     }
 
     @NotNull
-    /* package */ static JavaClassifierImpl<?> create(@NotNull PsiClass psiClass) {
+    /* package */ static JavaClassifierImpl<?> create(@NotNull PsiClass psiClass, JavaElementSourceFactory sourceFactory) {
         if (psiClass instanceof PsiTypeParameter) {
-            return new JavaTypeParameterImpl((PsiTypeParameter) psiClass);
+            return new JavaTypeParameterImpl(sourceFactory.createPsiSource((PsiTypeParameter) psiClass));
         }
         else {
-            return new JavaClassImpl(psiClass);
+            return new JavaClassImpl(sourceFactory.createPsiSource(psiClass));
         }
     }
 
@@ -49,13 +51,13 @@ public abstract class JavaClassifierImpl<Psi extends PsiClass> extends JavaEleme
     @NotNull
     @Override
     public Collection<JavaAnnotation> getAnnotations() {
-        return JavaElementUtil.getAnnotations(this);
+        return JavaElementUtil.getAnnotations(this, getSourceFactory());
     }
 
     @Nullable
     @Override
     public JavaAnnotation findAnnotation(@NotNull FqName fqName) {
-        return JavaElementUtil.findAnnotation(this, fqName);
+        return JavaElementUtil.findAnnotation(this, fqName, getSourceFactory());
     }
 
     @Override

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

@@ -19,10 +19,11 @@ package org.jetbrains.kotlin.load.java.structure.impl
 import com.intellij.psi.*
 import org.jetbrains.kotlin.load.java.structure.JavaClassifierType
 import org.jetbrains.kotlin.load.java.structure.JavaType
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementTypeSource
 
-import java.util.ArrayList
+class JavaClassifierTypeImpl(
-
+    psiClassTypeSource: JavaElementTypeSource<PsiClassType>,
-class JavaClassifierTypeImpl(psiClassType: PsiClassType) : JavaTypeImpl<PsiClassType>(psiClassType), JavaClassifierType {
+) : JavaTypeImpl<PsiClassType>(psiClassTypeSource), JavaClassifierType {
 
     private var resolutionResult: ResolutionResult? = null
 
@@ -52,7 +53,7 @@ class JavaClassifierTypeImpl(psiClassType: PsiClassType) : JavaTypeImpl<PsiClass
             val result = ArrayList<JavaType?>(parameters.size)
             for (typeParameter in parameters) {
                 val substitutedType = substitutor.substitute(typeParameter)
-                result.add(substitutedType?.let { JavaTypeImpl.create(it) })
+                result.add(substitutedType?.let { JavaTypeImpl.create(createTypeSource(it)) })
             }
 
             return result
@@ -70,7 +71,7 @@ class JavaClassifierTypeImpl(psiClassType: PsiClassType) : JavaTypeImpl<PsiClass
             val psiClass = result.element
             val substitutor = result.substitutor
             ResolutionResult(
-                psiClass?.let { JavaClassifierImpl.create(it) }, substitutor, PsiClassType.isRaw(result)
+                psiClass?.let { JavaClassifierImpl.create(it, sourceFactory) }, substitutor, PsiClassType.isRaw(result)
             ).apply {
                 resolutionResult = this
             }

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

@@ -21,6 +21,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.load.java.structure.JavaConstructor;
 import org.jetbrains.kotlin.load.java.structure.JavaTypeParameter;
 import org.jetbrains.kotlin.load.java.structure.JavaValueParameter;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource;
 
 import java.util.List;
 
@@ -28,8 +29,9 @@ import static org.jetbrains.kotlin.load.java.structure.impl.JavaElementCollectio
 import static org.jetbrains.kotlin.load.java.structure.impl.JavaElementCollectionFromPsiArrayUtil.valueParameters;
 
 public class JavaConstructorImpl extends JavaMemberImpl<PsiMethod> implements JavaConstructor {
-    public JavaConstructorImpl(@NotNull PsiMethod psiMethod) {
+    public JavaConstructorImpl(@NotNull JavaElementPsiSource<PsiMethod> psiConstructorSource) {
-        super(psiMethod);
+        super(psiConstructorSource);
+        PsiMethod psiMethod = psiConstructorSource.getPsi();
         assert psiMethod.isConstructor() :
                 "PsiMethod which is not a constructor should not be wrapped in JavaConstructorImpl: " +
                 psiMethod.getName() + " " + psiMethod.getClass().getName();
@@ -38,12 +40,12 @@ public class JavaConstructorImpl extends JavaMemberImpl<PsiMethod> implements Ja
     @NotNull
     @Override
     public List<JavaValueParameter> getValueParameters() {
-        return valueParameters(getPsi().getParameterList().getParameters());
+        return valueParameters(getPsi().getParameterList().getParameters(), getSourceFactory());
     }
 
     @NotNull
     @Override
     public List<JavaTypeParameter> getTypeParameters() {
-        return typeParameters(getPsi().getTypeParameters());
+        return typeParameters(getPsi().getTypeParameters(), getSourceFactory());
     }
 }

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

@@ -22,6 +22,7 @@ import com.intellij.psi.*
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.load.java.NULLABILITY_ANNOTATIONS
 import org.jetbrains.kotlin.load.java.structure.*
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
 import org.jetbrains.kotlin.name.Name
 
 inline fun <Psi, Java> Array<Psi>.convert(factory: (Psi) -> Java): List<Java> =
@@ -38,47 +39,47 @@ fun <Psi, Java> Collection<Psi>.convert(factory: (Psi) -> Java): List<Java> =
             else -> map(factory)
         }
 
-internal fun classes(classes: Array<PsiClass>): Collection<JavaClass> =
+internal fun classes(classes: Array<PsiClass>, sourceFactory: JavaElementSourceFactory): Collection<JavaClass> =
-        classes.convert(::JavaClassImpl)
+        classes.convert { JavaClassImpl(sourceFactory.createPsiSource(it)) }
 
-internal fun classes(classes: Collection<PsiClass>): Collection<JavaClass> =
+internal fun classes(classes: Collection<PsiClass>, sourceFactory: JavaElementSourceFactory): Collection<JavaClass> =
-        classes.convert(::JavaClassImpl)
+        classes.convert { JavaClassImpl(sourceFactory.createPsiSource(it)) }
 
-internal fun packages(packages: Array<PsiPackage>, scope: GlobalSearchScope): Collection<JavaPackage> =
+internal fun packages(packages: Array<PsiPackage>, scope: GlobalSearchScope, sourceFactory: JavaElementSourceFactory): Collection<JavaPackage> =
-        packages.convert { psi -> JavaPackageImpl(psi, scope) }
+        packages.convert { psi -> JavaPackageImpl(sourceFactory.createPsiSource(psi), scope) }
 
-internal fun methods(methods: Collection<PsiMethod>): Collection<JavaMethod> =
+internal fun methods(methods: Collection<PsiMethod>, sourceFactory: JavaElementSourceFactory): Collection<JavaMethod> =
-        methods.convert(::JavaMethodImpl)
+        methods.convert { JavaMethodImpl(sourceFactory.createPsiSource(it)) }
 
-internal fun constructors(methods: Collection<PsiMethod>): Collection<JavaConstructor> =
+internal fun constructors(methods: Collection<PsiMethod>, sourceFactory: JavaElementSourceFactory): Collection<JavaConstructor> =
-        methods.convert(::JavaConstructorImpl)
+        methods.convert { JavaConstructorImpl(sourceFactory.createPsiSource(it)) }
 
-internal fun fields(fields: Collection<PsiField>): Collection<JavaField> =
+internal fun fields(fields: Collection<PsiField>, sourceFactory: JavaElementSourceFactory): Collection<JavaField> =
-        fields.convert(::JavaFieldImpl)
+        fields.convert { JavaFieldImpl(sourceFactory.createPsiSource(it)) }
 
-internal fun valueParameters(parameters: Array<PsiParameter>): List<JavaValueParameter> =
+internal fun valueParameters(parameters: Array<PsiParameter>, sourceFactory: JavaElementSourceFactory): List<JavaValueParameter> =
-        parameters.convert(::JavaValueParameterImpl)
+        parameters.convert { JavaValueParameterImpl(sourceFactory.createPsiSource(it)) }
 
-internal fun typeParameters(typeParameters: Array<PsiTypeParameter>): List<JavaTypeParameter> =
+internal fun typeParameters(typeParameters: Array<PsiTypeParameter>, sourceFactory: JavaElementSourceFactory): List<JavaTypeParameter> =
-        typeParameters.convert(::JavaTypeParameterImpl)
+        typeParameters.convert { JavaTypeParameterImpl(sourceFactory.createPsiSource(it)) }
 
-internal fun classifierTypes(classTypes: Array<PsiClassType>): Collection<JavaClassifierType> =
+internal fun classifierTypes(classTypes: Array<PsiClassType>, sourceFactory: JavaElementSourceFactory): Collection<JavaClassifierType> =
-        classTypes.convert(::JavaClassifierTypeImpl)
+        classTypes.convert { JavaClassifierTypeImpl(sourceFactory.createTypeSource(it)) }
 
-internal fun annotations(annotations: Array<out PsiAnnotation>): Collection<JavaAnnotation> =
+internal fun annotations(annotations: Array<out PsiAnnotation>, sourceFactory: JavaElementSourceFactory): Collection<JavaAnnotation> =
-        annotations.convert(::JavaAnnotationImpl)
+        annotations.convert { JavaAnnotationImpl(sourceFactory.createPsiSource(it)) }
 
-internal fun nullabilityAnnotations(annotations: Array<out PsiAnnotation>): Collection<JavaAnnotation> =
+internal fun nullabilityAnnotations(annotations: Array<out PsiAnnotation>, sourceFactory: JavaElementSourceFactory): Collection<JavaAnnotation> =
-        annotations.convert(::JavaAnnotationImpl)
+        annotations.convert { JavaAnnotationImpl(sourceFactory.createPsiSource(it)) }
                 .filter { annotation ->
                     val fqName = annotation.classId?.asSingleFqName() ?: return@filter false
                     fqName in NULLABILITY_ANNOTATIONS
                 }
 
 
-internal fun namedAnnotationArguments(nameValuePairs: Array<PsiNameValuePair>): Collection<JavaAnnotationArgument> =
+internal fun namedAnnotationArguments(nameValuePairs: Array<PsiNameValuePair>, sourceFactory: JavaElementSourceFactory): Collection<JavaAnnotationArgument> =
         nameValuePairs.convert { psi ->
             val name = psi.name?.let(Name::identifier)
             val value = psi.value ?: error("Annotation argument value cannot be null: $name")
-            JavaAnnotationArgumentImpl.create(value, name)
+            JavaAnnotationArgumentImpl.create(value, name, sourceFactory)
         }

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

@@ -19,17 +19,26 @@ package org.jetbrains.kotlin.load.java.structure.impl;
 import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.load.java.structure.JavaElement;
+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.JavaSourceFactoryOwner;
 
-public abstract class JavaElementImpl<Psi extends PsiElement> implements JavaElement {
+public abstract class JavaElementImpl<Psi extends PsiElement> implements JavaElement, JavaSourceFactoryOwner {
-    private final Psi psiElement;
+    protected final JavaElementPsiSource<Psi> psiElementSource;
 
-    protected JavaElementImpl(@NotNull Psi psiElement) {
+    @Override
-        this.psiElement = psiElement;
+    @NotNull
+    public JavaElementSourceFactory getSourceFactory() {
+        return psiElementSource.getFactory();
+    }
+
+    protected JavaElementImpl(@NotNull JavaElementPsiSource<Psi> psiElementSource) {
+        this.psiElementSource = psiElementSource;
     }
 
     @NotNull
     public Psi getPsi() {
-        return psiElement;
+        return psiElementSource.getPsi();
     }
 
     @Override
@@ -44,6 +53,6 @@ public abstract class JavaElementImpl<Psi extends PsiElement> implements JavaEle
 
     @Override
     public String toString() {
-        return getClass().getSimpleName() + ": " + getPsi();
+        return getClass().getSimpleName() + ": " + psiElementSource;
     }
 }

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

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.descriptors.Visibilities;
 import org.jetbrains.kotlin.descriptors.Visibility;
 import org.jetbrains.kotlin.descriptors.java.JavaVisibilities;
 import org.jetbrains.kotlin.load.java.structure.JavaAnnotation;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory;
 import org.jetbrains.kotlin.name.FqName;
 
 import java.util.ArrayList;
@@ -72,10 +73,10 @@ import static org.jetbrains.kotlin.load.java.structure.impl.JavaElementCollectio
     }
 
     @NotNull
-    public static Collection<JavaAnnotation> getAnnotations(@NotNull JavaAnnotationOwnerImpl owner) {
+    public static Collection<JavaAnnotation> getAnnotations(@NotNull JavaAnnotationOwnerImpl owner, JavaElementSourceFactory sourceFactory) {
         PsiAnnotationOwner annotationOwnerPsi = owner.getAnnotationOwnerPsi();
         if (annotationOwnerPsi != null) {
-            return annotations(annotationOwnerPsi.getAnnotations());
+            return annotations(annotationOwnerPsi.getAnnotations(), sourceFactory);
         }
         return Collections.emptyList();
     }
@@ -90,23 +91,23 @@ import static org.jetbrains.kotlin.load.java.structure.impl.JavaElementCollectio
 
     @NotNull
     static <T extends JavaAnnotationOwnerImpl & JavaModifierListOwnerImpl>
-    Collection<JavaAnnotation> getRegularAndExternalAnnotations(@NotNull T owner) {
+    Collection<JavaAnnotation> getRegularAndExternalAnnotations(@NotNull T owner, JavaElementSourceFactory sourceFactory) {
         PsiAnnotation[] externalAnnotations = getExternalAnnotations(owner);
         if (externalAnnotations == null) {
-            return getAnnotations(owner);
+            return getAnnotations(owner, sourceFactory);
         }
-        Collection<JavaAnnotation> annotations = new ArrayList<>(getAnnotations(owner));
+        Collection<JavaAnnotation> annotations = new ArrayList<>(getAnnotations(owner, sourceFactory));
-        annotations.addAll(nullabilityAnnotations(externalAnnotations));
+        annotations.addAll(nullabilityAnnotations(externalAnnotations, sourceFactory));
         return annotations;
     }
 
 
     @Nullable
-    public static JavaAnnotation findAnnotation(@NotNull JavaAnnotationOwnerImpl owner, @NotNull FqName fqName) {
+    public static JavaAnnotation findAnnotation(@NotNull JavaAnnotationOwnerImpl owner, @NotNull FqName fqName, JavaElementSourceFactory sourceFactory) {
         PsiAnnotationOwner annotationOwnerPsi = owner.getAnnotationOwnerPsi();
         if (annotationOwnerPsi != null) {
             PsiAnnotation psiAnnotation = annotationOwnerPsi.findAnnotation(fqName.asString());
-            return psiAnnotation == null ? null : new JavaAnnotationImpl(psiAnnotation);
+            return psiAnnotation == null ? null : new JavaAnnotationImpl(sourceFactory.createPsiSource(psiAnnotation));
         }
         return null;
     }

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

@@ -24,10 +24,11 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.load.java.structure.JavaField;
 import org.jetbrains.kotlin.load.java.structure.JavaType;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource;
 
 public class JavaFieldImpl extends JavaMemberImpl<PsiField> implements JavaField {
-    public JavaFieldImpl(@NotNull PsiField psiField) {
+    public JavaFieldImpl(@NotNull JavaElementPsiSource<PsiField> psiFieldSource) {
-        super(psiField);
+        super(psiFieldSource);
     }
 
     @Override
@@ -38,7 +39,7 @@ public class JavaFieldImpl extends JavaMemberImpl<PsiField> implements JavaField
     @Override
     @NotNull
     public JavaType getType() {
-        return JavaTypeImpl.create(getPsi().getType());
+        return JavaTypeImpl.create(psiElementSource.getPsi().getType(), createVariableReturnTypeSource(psiElementSource));
     }
 
     @Nullable

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

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.descriptors.Visibility;
 import org.jetbrains.kotlin.load.java.structure.JavaAnnotation;
 import org.jetbrains.kotlin.load.java.structure.JavaClass;
 import org.jetbrains.kotlin.load.java.structure.JavaMember;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.name.Name;
 
@@ -33,7 +34,7 @@ import java.util.Collection;
 
 public abstract class JavaMemberImpl<Psi extends PsiMember> extends JavaElementImpl<Psi>
         implements JavaMember, JavaAnnotationOwnerImpl, JavaModifierListOwnerImpl {
-    protected JavaMemberImpl(@NotNull Psi psiMember) {
+    protected JavaMemberImpl(JavaElementPsiSource<Psi> psiMember) {
         super(psiMember);
     }
 
@@ -61,7 +62,7 @@ public abstract class JavaMemberImpl<Psi extends PsiMember> extends JavaElementI
     public JavaClass getContainingClass() {
         PsiClass psiClass = getPsi().getContainingClass();
         assert psiClass != null : "Member must have a containing class: " + getPsi();
-        return new JavaClassImpl(psiClass);
+        return new JavaClassImpl(createPsiSource(psiClass));
     }
 
     @Override
@@ -88,13 +89,13 @@ public abstract class JavaMemberImpl<Psi extends PsiMember> extends JavaElementI
     @NotNull
     @Override
     public Collection<JavaAnnotation> getAnnotations() {
-        return JavaElementUtil.getRegularAndExternalAnnotations(this);
+        return JavaElementUtil.getRegularAndExternalAnnotations(this, getSourceFactory());
     }
 
     @Nullable
     @Override
     public JavaAnnotation findAnnotation(@NotNull FqName fqName) {
-        return JavaElementUtil.findAnnotation(this, fqName);
+        return JavaElementUtil.findAnnotation(this, fqName, getSourceFactory());
     }
 
     @Override

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

@@ -23,6 +23,7 @@ import com.intellij.psi.PsiType;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.load.java.structure.*;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource;
 import org.jetbrains.kotlin.name.Name;
 
 import java.util.List;
@@ -31,10 +32,11 @@ import static org.jetbrains.kotlin.load.java.structure.impl.JavaElementCollectio
 import static org.jetbrains.kotlin.load.java.structure.impl.JavaElementCollectionFromPsiArrayUtil.valueParameters;
 
 public class JavaMethodImpl extends JavaMemberImpl<PsiMethod> implements JavaMethod {
-    public JavaMethodImpl(@NotNull PsiMethod psiMethod) {
+    public JavaMethodImpl(@NotNull JavaElementPsiSource<PsiMethod> psiMethodSource) {
-        super(psiMethod);
+        super(psiMethodSource);
-        assert !psiMethod.isConstructor() :
+        PsiMethod method = psiMethodSource.getPsi();
-                "PsiMethod which is a constructor should be wrapped in JavaConstructorImpl: " + psiMethod.getName();
+        assert !method.isConstructor() :
+                "PsiMethod which is a constructor should be wrapped in JavaConstructorImpl: " + method.getName();
     }
 
     @NotNull
@@ -46,13 +48,13 @@ public class JavaMethodImpl extends JavaMemberImpl<PsiMethod> implements JavaMet
     @NotNull
     @Override
     public List<JavaTypeParameter> getTypeParameters() {
-        return typeParameters(getPsi().getTypeParameters());
+        return typeParameters(getPsi().getTypeParameters(), getSourceFactory());
     }
 
     @Override
     @NotNull
     public List<JavaValueParameter> getValueParameters() {
-        return valueParameters(getPsi().getParameterList().getParameters());
+        return valueParameters(getPsi().getParameterList().getParameters(), getSourceFactory());
     }
 
     @Override
@@ -62,7 +64,7 @@ public class JavaMethodImpl extends JavaMemberImpl<PsiMethod> implements JavaMet
         if (psiMethod instanceof PsiAnnotationMethod) {
             PsiAnnotationMemberValue defaultValue = ((PsiAnnotationMethod) psiMethod).getDefaultValue();
             if (defaultValue != null) {
-                return JavaAnnotationArgumentImpl.Factory.create(defaultValue, null);
+                return JavaAnnotationArgumentImpl.Factory.create(defaultValue, null, getSourceFactory());
             }
         }
 
@@ -80,6 +82,6 @@ public class JavaMethodImpl extends JavaMemberImpl<PsiMethod> implements JavaMet
     public JavaType getReturnType() {
         PsiType psiType = getPsi().getReturnType();
         assert psiType != null : "Method is not a constructor and has no return type: " + getName();
-        return JavaTypeImpl.create(psiType);
+        return JavaTypeImpl.create(createTypeSource(psiType));
     }
 }

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

@@ -19,24 +19,25 @@ package org.jetbrains.kotlin.load.java.structure.impl
 import com.intellij.psi.PsiPackage
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.load.java.structure.*
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
 class JavaPackageImpl(
-    psiPackage: PsiPackage, private val scope: GlobalSearchScope,
+    psiPackageSource: JavaElementPsiSource<PsiPackage>, private val scope: GlobalSearchScope,
     private val mayHaveAnnotations: Boolean = true,
-) : JavaElementImpl<PsiPackage>(psiPackage), JavaPackage, MapBasedJavaAnnotationOwner {
+) : JavaElementImpl<PsiPackage>(psiPackageSource), JavaPackage, MapBasedJavaAnnotationOwner {
 
     override fun getClasses(nameFilter: (Name) -> Boolean): Collection<JavaClass> {
         val psiClasses = psi.getClasses(scope).filter {
             val name = it.name
             name != null && nameFilter(Name.identifier(name))
         }
-        return classes(psiClasses)
+        return classes(psiClasses, sourceFactory)
     }
 
     override val subPackages: Collection<JavaPackage>
-        get() = packages(psi.getSubPackages(scope), scope)
+        get() = packages(psi.getSubPackages(scope), scope, sourceFactory)
 
     override val fqName: FqName
         get() = FqName(psi.qualifiedName)
@@ -44,7 +45,7 @@ class JavaPackageImpl(
     override val annotations: Collection<JavaAnnotation>
         get() =
             if (mayHaveAnnotations)
-                org.jetbrains.kotlin.load.java.structure.impl.annotations(psi.annotationList?.annotations.orEmpty())
+                org.jetbrains.kotlin.load.java.structure.impl.annotations(psi.annotationList?.annotations.orEmpty(), sourceFactory)
             else
                 emptyList()
 

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

@@ -21,11 +21,12 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.PrimitiveType;
 import org.jetbrains.kotlin.load.java.structure.JavaPrimitiveType;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementTypeSource;
 import org.jetbrains.kotlin.resolve.jvm.JvmPrimitiveType;
 
 public class JavaPrimitiveTypeImpl extends JavaTypeImpl<PsiPrimitiveType> implements JavaPrimitiveType {
-    public JavaPrimitiveTypeImpl(@NotNull PsiPrimitiveType psiPrimitiveType) {
+    public JavaPrimitiveTypeImpl(@NotNull JavaElementTypeSource<PsiPrimitiveType> psiPrimitiveTypeSource) {
-        super(psiPrimitiveType);
+        super(psiPrimitiveTypeSource);
     }
 
     @Override

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

@@ -8,10 +8,13 @@ package org.jetbrains.kotlin.load.java.structure.impl
 import com.intellij.psi.PsiRecordComponent
 import org.jetbrains.kotlin.load.java.structure.JavaRecordComponent
 import org.jetbrains.kotlin.load.java.structure.JavaType
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource
 
-class JavaRecordComponentImpl(psiRecordComponent: PsiRecordComponent) :  JavaMemberImpl<PsiRecordComponent>(psiRecordComponent), JavaRecordComponent {
+class JavaRecordComponentImpl(
+    psiRecordComponentSource: JavaElementPsiSource<PsiRecordComponent>
+) : JavaMemberImpl<PsiRecordComponent>(psiRecordComponentSource), JavaRecordComponent {
     override val type: JavaType
-        get() = JavaTypeImpl.create(psi.type)
+        get() = JavaTypeImpl.create(sourceFactory.createTypeSource(psiElementSource.psi.type))
 
     override val isVararg: Boolean
         get() = psi.isVarArgs

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

@@ -21,20 +21,30 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.load.java.structure.JavaAnnotation;
 import org.jetbrains.kotlin.load.java.structure.JavaType;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementTypeSource;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaSourceFactoryOwner;
 import org.jetbrains.kotlin.name.FqName;
 
 import java.util.Collection;
 
-public abstract class JavaTypeImpl<Psi extends PsiType> implements JavaType, JavaAnnotationOwnerImpl {
+public abstract class JavaTypeImpl<Psi extends PsiType> implements JavaType, JavaAnnotationOwnerImpl, JavaSourceFactoryOwner {
-    private final Psi psiType;
+    private final JavaElementTypeSource<Psi> psiType;
 
-    public JavaTypeImpl(@NotNull Psi psiType) {
+    public JavaTypeImpl(@NotNull JavaElementTypeSource<Psi> psiTypeSource) {
-        this.psiType = psiType;
+        this.psiType = psiTypeSource;
     }
 
+    @Override
+    @NotNull
+    public JavaElementSourceFactory getSourceFactory() {
+        return psiType.getFactory();
+    }
+
+
     @NotNull
     public Psi getPsi() {
-        return psiType;
+        return psiType.getType();
     }
 
     @Nullable
@@ -44,8 +54,10 @@ public abstract class JavaTypeImpl<Psi extends PsiType> implements JavaType, Jav
     }
 
     @NotNull
-    public static JavaTypeImpl<?> create(@NotNull PsiType psiType) {
+    public static JavaTypeImpl<?> create(@NotNull JavaElementTypeSource<PsiType> psiTypeSource) {
-        return psiType.accept(new PsiTypeVisitor<JavaTypeImpl<?>>() {
+        JavaElementSourceFactory sourceFactory = psiTypeSource.getFactory();
+        return psiTypeSource.getType().accept(new PsiTypeVisitor<JavaTypeImpl<?>>() {
+
             @Nullable
             @Override
             public JavaTypeImpl<?> visitType(@NotNull PsiType type) {
@@ -55,25 +67,25 @@ public abstract class JavaTypeImpl<Psi extends PsiType> implements JavaType, Jav
             @Nullable
             @Override
             public JavaTypeImpl<?> visitPrimitiveType(@NotNull PsiPrimitiveType primitiveType) {
-                return new JavaPrimitiveTypeImpl(primitiveType);
+                return new JavaPrimitiveTypeImpl(sourceFactory.createTypeSource(primitiveType));
             }
 
             @Nullable
             @Override
             public JavaTypeImpl<?> visitArrayType(@NotNull PsiArrayType arrayType) {
-                return new JavaArrayTypeImpl(arrayType);
+                return new JavaArrayTypeImpl(sourceFactory.createTypeSource(arrayType));
             }
 
             @Nullable
             @Override
             public JavaTypeImpl<?> visitClassType(@NotNull PsiClassType classType) {
-                return new JavaClassifierTypeImpl(classType);
+                return new JavaClassifierTypeImpl(sourceFactory.createTypeSource(classType));
             }
 
             @Nullable
             @Override
             public JavaTypeImpl<?> visitWildcardType(@NotNull PsiWildcardType wildcardType) {
-                return new JavaWildcardTypeImpl(wildcardType);
+                return new JavaWildcardTypeImpl(sourceFactory.createTypeSource(wildcardType));
             }
         });
     }
@@ -81,13 +93,13 @@ public abstract class JavaTypeImpl<Psi extends PsiType> implements JavaType, Jav
     @NotNull
     @Override
     public Collection<JavaAnnotation> getAnnotations() {
-        return JavaElementUtil.getAnnotations(this);
+        return JavaElementUtil.getAnnotations(this, getSourceFactory());
     }
 
     @Nullable
     @Override
     public JavaAnnotation findAnnotation(@NotNull FqName fqName) {
-        return JavaElementUtil.findAnnotation(this, fqName);
+        return JavaElementUtil.findAnnotation(this, fqName, getSourceFactory());
     }
 
     @Override

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

@@ -22,6 +22,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.load.java.structure.JavaClassifierType;
 import org.jetbrains.kotlin.load.java.structure.JavaTypeParameter;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource;
 import org.jetbrains.kotlin.name.Name;
 import org.jetbrains.kotlin.name.SpecialNames;
 
@@ -30,8 +31,8 @@ import java.util.Collection;
 import static org.jetbrains.kotlin.load.java.structure.impl.JavaElementCollectionFromPsiArrayUtil.classifierTypes;
 
 public class JavaTypeParameterImpl extends JavaClassifierImpl<PsiTypeParameter> implements JavaTypeParameter {
-    public JavaTypeParameterImpl(@NotNull PsiTypeParameter psiTypeParameter) {
+    public JavaTypeParameterImpl(@NotNull JavaElementPsiSource<PsiTypeParameter> psiTypeParameterSource) {
-        super(psiTypeParameter);
+        super(psiTypeParameterSource);
     }
 
     @NotNull
@@ -43,7 +44,7 @@ public class JavaTypeParameterImpl extends JavaClassifierImpl<PsiTypeParameter>
     @Override
     @NotNull
     public Collection<JavaClassifierType> getUpperBounds() {
-        return classifierTypes(getPsi().getExtendsList().getReferencedTypes());
+        return classifierTypes(getPsi().getExtendsList().getReferencedTypes(), getSourceFactory());
     }
 
     @Nullable

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

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.descriptors.Visibility;
 import org.jetbrains.kotlin.load.java.structure.JavaAnnotation;
 import org.jetbrains.kotlin.load.java.structure.JavaType;
 import org.jetbrains.kotlin.load.java.structure.JavaValueParameter;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.name.Name;
 
@@ -33,8 +34,8 @@ import java.util.Collection;
 
 public class JavaValueParameterImpl extends JavaElementImpl<PsiParameter>
         implements JavaValueParameter, JavaAnnotationOwnerImpl, JavaModifierListOwnerImpl {
-    public JavaValueParameterImpl(@NotNull PsiParameter psiParameter) {
+    public JavaValueParameterImpl(@NotNull JavaElementPsiSource<PsiParameter> psiParameterSource) {
-        super(psiParameter);
+        super(psiParameterSource);
     }
 
     @Nullable
@@ -72,13 +73,13 @@ public class JavaValueParameterImpl extends JavaElementImpl<PsiParameter>
     @NotNull
     @Override
     public Collection<JavaAnnotation> getAnnotations() {
-        return JavaElementUtil.getRegularAndExternalAnnotations(this);
+        return JavaElementUtil.getRegularAndExternalAnnotations(this, getSourceFactory());
     }
 
     @Nullable
     @Override
     public JavaAnnotation findAnnotation(@NotNull FqName fqName) {
-        return JavaElementUtil.findAnnotation(this, fqName);
+        return JavaElementUtil.findAnnotation(this, fqName, getSourceFactory());
     }
 
     @Override
@@ -101,7 +102,7 @@ public class JavaValueParameterImpl extends JavaElementImpl<PsiParameter>
     @Override
     @NotNull
     public JavaType getType() {
-        return JavaTypeImpl.create(getPsi().getType());
+        return JavaTypeImpl.create(createTypeSource(getPsi().getType()));
     }
 
     @Override

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

@@ -21,17 +21,18 @@ import com.intellij.psi.PsiWildcardType;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.load.java.structure.JavaWildcardType;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementTypeSource;
 
 public class JavaWildcardTypeImpl extends JavaTypeImpl<PsiWildcardType> implements JavaWildcardType {
-    public JavaWildcardTypeImpl(@NotNull PsiWildcardType psiWildcardType) {
+    public JavaWildcardTypeImpl(@NotNull JavaElementTypeSource<PsiWildcardType> psiWildcardTypeSource) {
-        super(psiWildcardType);
+        super(psiWildcardTypeSource);
     }
 
     @Override
     @Nullable
     public JavaTypeImpl<?> getBound() {
         PsiType bound = getPsi().getBound();
-        return bound == null ? null : create(bound);
+        return bound == null ? null : create(createTypeSource(bound));
     }
 
     @Override

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

@@ -18,6 +18,8 @@ package org.jetbrains.kotlin.load.java.structure.impl
 
 import com.intellij.psi.*
 import org.jetbrains.kotlin.load.java.structure.*
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -26,14 +28,18 @@ sealed class JavaAnnotationArgumentImpl(
     override val name: Name?
 ) : JavaAnnotationArgument {
     companion object Factory {
-        fun create(argument: PsiAnnotationMemberValue, name: Name?): JavaAnnotationArgument {
+        fun create(
+            argument: PsiAnnotationMemberValue,
+            name: Name?,
+            sourceFactory: JavaElementSourceFactory
+        ): JavaAnnotationArgument {
             if (argument is PsiClassObjectAccessExpression) {
-                return JavaClassObjectAnnotationArgumentImpl(argument, name)
+                return JavaClassObjectAnnotationArgumentImpl(sourceFactory.createPsiSource(argument), name)
             }
 
             val value = JavaPsiFacade.getInstance(argument.project).constantEvaluationHelper.computeConstantExpression(argument)
             if (value is Enum<*>) {
-                return JavaEnumValueAnnotationArgumentImpl(argument as PsiReferenceExpression, name)
+                return JavaEnumValueAnnotationArgumentImpl(sourceFactory.createPsiSource(argument as PsiReferenceExpression), name)
             }
 
             if (value != null || argument is PsiLiteralExpression) {
@@ -41,9 +47,9 @@ sealed class JavaAnnotationArgumentImpl(
             }
 
             return when (argument) {
-                is PsiReferenceExpression -> JavaEnumValueAnnotationArgumentImpl(argument, name)
+                is PsiReferenceExpression -> JavaEnumValueAnnotationArgumentImpl(sourceFactory.createPsiSource(argument), name)
-                is PsiArrayInitializerMemberValue -> JavaArrayAnnotationArgumentImpl(argument, name)
+                is PsiArrayInitializerMemberValue -> JavaArrayAnnotationArgumentImpl(sourceFactory.createPsiSource(argument), name)
-                is PsiAnnotation -> JavaAnnotationAsAnnotationArgumentImpl(argument, name)
+                is PsiAnnotation -> JavaAnnotationAsAnnotationArgumentImpl(sourceFactory.createPsiSource(argument), name)
                 else -> JavaUnknownAnnotationArgumentImpl(name)
             }
         }
@@ -56,44 +62,47 @@ class JavaLiteralAnnotationArgumentImpl(
 ) : JavaLiteralAnnotationArgument
 
 class JavaArrayAnnotationArgumentImpl(
-    private val psiValue: PsiArrayInitializerMemberValue,
+    private val psiValueSource: JavaElementPsiSource<PsiArrayInitializerMemberValue>,
-    name: Name?
+    name: Name?,
 ) : JavaAnnotationArgumentImpl(name), JavaArrayAnnotationArgument {
-    override fun getElements() = psiValue.initializers.map { create(it, null) }
+    override fun getElements() = psiValueSource.psi.initializers.map { create(it, null, psiValueSource.factory) }
 }
 
 class JavaEnumValueAnnotationArgumentImpl(
-    private val psiReference: PsiReferenceExpression,
+    private val psiReferenceSource: JavaElementPsiSource<PsiReferenceExpression>,
     name: Name?
 ) : JavaAnnotationArgumentImpl(name), JavaEnumValueAnnotationArgument {
     override val enumClassId: ClassId?
         get() {
-            val element = psiReference.resolve()
+            val element = psiReferenceSource.psi.resolve()
             if (element is PsiEnumConstant) {
-                return JavaFieldImpl(element).containingClass.classId
+                return JavaFieldImpl(psiReferenceSource.factory.createPsiSource(element)).containingClass.classId
             }
 
-            val fqName = (psiReference.qualifier as? PsiReferenceExpression)?.qualifiedName ?: return null
+            val fqName = ( psiReferenceSource.psi.qualifier as? PsiReferenceExpression)?.qualifiedName ?: return null
             // TODO: find a way to construct a correct name (with nested classes) for unresolved enums
             return ClassId.topLevel(FqName(fqName))
         }
 
     override val entryName: Name?
-        get() = psiReference.referenceName?.let(Name::identifier)
+        get() = psiReferenceSource.psi.referenceName?.let(Name::identifier)
 }
 
 class JavaClassObjectAnnotationArgumentImpl(
-    private val psiExpression: PsiClassObjectAccessExpression,
+    private val psiExpressionSource: JavaElementPsiSource<PsiClassObjectAccessExpression>,
     name: Name?
 ) : JavaAnnotationArgumentImpl(name), JavaClassObjectAnnotationArgument {
-    override fun getReferencedType() = JavaTypeImpl.create(psiExpression.operand.type)
+    override fun getReferencedType(): JavaTypeImpl<*> {
+        val operand = psiExpressionSource.psi.operand
+        return JavaTypeImpl.create(operand.type, psiExpressionSource.factory.createTypeSource(operand.type))
+    }
 }
 
 class JavaAnnotationAsAnnotationArgumentImpl(
-    private val psiAnnotation: PsiAnnotation,
+    private val psiAnnotationSource: JavaElementPsiSource<PsiAnnotation>,
-    name: Name?
+    name: Name?,
 ) : JavaAnnotationArgumentImpl(name), JavaAnnotationAsAnnotationArgument {
-    override fun getAnnotation() = JavaAnnotationImpl(psiAnnotation)
+    override fun getAnnotation() = JavaAnnotationImpl(psiAnnotationSource)
 }
 
 class JavaUnknownAnnotationArgumentImpl(name: Name?) : JavaAnnotationArgumentImpl(name), JavaUnknownAnnotationArgument

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

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.load.java.structure.impl.source
+
+import com.intellij.psi.PsiElement
+
+abstract class JavaElementPsiSource<PSI : PsiElement> {
+    abstract val psi: PSI
+    abstract val factory: JavaElementSourceFactory
+}
+
+class JavaElementPsiSourceWithFixedPsi<PSI : PsiElement>(
+    override val psi: PSI
+) : JavaElementPsiSource<PSI>() {
+    override val factory: JavaElementSourceFactory
+        get() = JavaElementSourceFactory.getInstance(psi.project)
+
+    override fun toString(): String {
+        return psi.toString()
+    }
+}

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.load.java.structure.impl.source
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+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>
+
+    companion object {
+        @JvmStatic
+        fun getInstance(project: Project): JavaElementSourceFactory {
+            return project.getService(JavaElementSourceFactory::class.java)
+        }
+    }
+}
+
+class JavaFixedElementSourceFactory : JavaElementSourceFactory() {
+    override fun <PSI : PsiElement> createPsiSource(psi: PSI): JavaElementPsiSource<PSI> {
+        return JavaElementPsiSourceWithFixedPsi(psi)
+    }
+
+    override fun <TYPE : PsiType> createTypeSource(type: TYPE): JavaElementTypeSource<TYPE> {
+        return JavaElementTypeSourceWithFixedType(type, this)
+    }
+}

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.load.java.structure.impl.source
+
+import com.intellij.psi.PsiType
+
+abstract class JavaElementTypeSource<TYPE : PsiType> {
+    abstract val type: TYPE
+    abstract val factory: JavaElementSourceFactory
+}
+
+class JavaElementTypeSourceWithFixedType<TYPE : PsiType>(
+    override val type: TYPE,
+    override val factory: JavaElementSourceFactory,
+) : JavaElementTypeSource<TYPE>() {
+
+    override fun toString(): String {
+        return type.toString()
+    }
+}

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

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.load.java.structure.impl.source
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiType
+
+interface JavaSourceFactoryOwner {
+    val sourceFactory: JavaElementSourceFactory
+
+    fun <PSI : PsiElement> createPsiSource(psi: PSI): JavaElementPsiSource<PSI> {
+        return sourceFactory.createPsiSource(psi)
+    }
+
+    fun <TYPE : PsiType> createTypeSource(type: TYPE): JavaElementTypeSource<TYPE> {
+        return sourceFactory.createTypeSource(type)
+    }
+}

compiler/resolution.common.jvm/src/org/jetbrains/kotlin/resolve/jvm/KotlinJavaPsiFacade.java

@@ -48,6 +48,7 @@ import org.jetbrains.kotlin.idea.KotlinLanguage;
 import org.jetbrains.kotlin.load.java.JavaClassFinder;
 import org.jetbrains.kotlin.load.java.structure.JavaClass;
 import org.jetbrains.kotlin.load.java.structure.impl.JavaClassImpl;
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory;
 import org.jetbrains.kotlin.name.ClassId;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus;
@@ -222,8 +223,8 @@ public class KotlinJavaPsiFacade implements Disposable {
     }
 
     @NotNull
-    private static JavaClass createJavaClass(@NotNull ClassId classId, @NotNull PsiClass psiClass) {
+    private JavaClass createJavaClass(@NotNull ClassId classId, @NotNull PsiClass psiClass) {
-        JavaClassImpl javaClass = new JavaClassImpl(psiClass);
+        JavaClassImpl javaClass = new JavaClassImpl(JavaElementSourceFactory.getInstance(project).createPsiSource(psiClass));
         FqName fqName = classId.asSingleFqName();
         if (!fqName.equals(javaClass.getFqName())) {
             throw new IllegalStateException("Requested " + fqName + ", got " + javaClass.getFqName());

compiler/tests/org/jetbrains/kotlin/jvm/compiler/KotlinClassFinderTest.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.asJava.classes.KtLightClass
 import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.load.java.structure.impl.JavaClassImpl
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
 import org.jetbrains.kotlin.load.kotlin.VirtualFileFinder
 import org.jetbrains.kotlin.load.kotlin.findKotlinClass
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmMetadataVersion
@@ -62,7 +63,7 @@ class KotlinClassFinderTest : KotlinTestWithEnvironmentManagement() {
         assertTrue(psiClass !is KtLightClass, "Kotlin light classes are not not expected")
 
         val binaryClass = VirtualFileFinder.SERVICE.getInstance(project).findKotlinClass(
-            JavaClassImpl(psiClass), JvmMetadataVersion.INSTANCE
+            JavaClassImpl(JavaElementSourceFactory.getInstance(project).createPsiSource(psiClass)), JvmMetadataVersion.INSTANCE
         )
         assertNotNull(binaryClass, "No binary class for $className")