[psi2ir] Emit source info for annotation parameters

We need this in the JS IR backend for annotations like
@JsFun and @JsPolyfill

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/descriptors/LazyAnnotations.kt

@@ -27,7 +27,6 @@ import org.jetbrains.kotlin.resolve.AnnotationResolver
 import org.jetbrains.kotlin.resolve.AnnotationResolverImpl
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingTrace
-import org.jetbrains.kotlin.resolve.constants.ConstantValue
 import org.jetbrains.kotlin.resolve.lazy.ForceResolveUtil
 import org.jetbrains.kotlin.resolve.lazy.LazyEntity
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
@@ -108,20 +107,27 @@ class LazyAnnotationDescriptor(
         LexicalScope.Base(c.scope, FileDescriptorForVisibilityChecks(source, it))
     } ?: c.scope
 
-    override val allValueArguments by c.storageManager.createLazyValue {
+    private val valueArgumentsWithSourceInfo by c.storageManager.createLazyValue {
         val resolutionResults = c.annotationResolver.resolveAnnotationCall(annotationEntry, scope, c.trace)
         AnnotationResolverImpl.checkAnnotationType(annotationEntry, c.trace, resolutionResults)
 
-        if (!resolutionResults.isSingleResult) return@createLazyValue emptyMap<Name, ConstantValue<*>>()
+        if (!resolutionResults.isSingleResult) return@createLazyValue emptyMap()
 
         resolutionResults.resultingCall.valueArguments.mapNotNull { (valueParameter, resolvedArgument) ->
             if (resolvedArgument == null) null
             else c.annotationResolver.getAnnotationArgumentValue(c.trace, valueParameter, resolvedArgument)?.let { value ->
-                valueParameter.name to value
+                valueParameter.name to (value to resolvedArgument.arguments.firstOrNull()?.getArgumentExpression().toSourceElement())
             }
         }.toMap()
     }
 
+    override val allValueArguments by c.storageManager.createLazyValue {
+        valueArgumentsWithSourceInfo.mapValues { it.value.first }
+    }
+
+    override fun getSourceForArgument(name: Name): SourceElement =
+        valueArgumentsWithSourceInfo[name]?.second ?: SourceElement.NO_SOURCE
+
     override fun forceResolveAllContents() {
         ForceResolveUtil.forceResolveAllContents(type)
         allValueArguments

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ConstantValueGeneratorImpl.kt

@@ -6,23 +6,30 @@
 package org.jetbrains.kotlin.psi2ir.generators
 
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.descriptors.SourceElement
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.util.ConstantValueGenerator
 import org.jetbrains.kotlin.ir.util.ReferenceSymbolTable
 import org.jetbrains.kotlin.ir.util.TypeTranslator
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
-import org.jetbrains.kotlin.resolve.source.PsiSourceElement
+import org.jetbrains.kotlin.resolve.source.getPsi
-import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class ConstantValueGeneratorImpl(
     moduleDescriptor: ModuleDescriptor,
     symbolTable: ReferenceSymbolTable,
     typeTranslator: TypeTranslator,
 ) : ConstantValueGenerator(moduleDescriptor, symbolTable, typeTranslator) {
-    override fun extractAnnotationOffsets(annotationDescriptor: AnnotationDescriptor): Pair<Int, Int> {
+    override fun extractAnnotationOffsets(annotationDescriptor: AnnotationDescriptor): Pair<Int, Int> =
-        val psi = annotationDescriptor.source.safeAs<PsiSourceElement>()?.psi
+        extractOffsets(annotationDescriptor.source)
+
+    override fun extractAnnotationParameterOffsets(annotationDescriptor: AnnotationDescriptor, argumentName: Name): Pair<Int, Int> =
+        extractOffsets(annotationDescriptor.getSourceForArgument(argumentName))
+
+    private fun extractOffsets(sourceElement: SourceElement): Pair<Int, Int> {
+        val psi = sourceElement.getPsi()
         if (psi == null || psi.containingFile.fileType.isBinary) return UNDEFINED_OFFSET to UNDEFINED_OFFSET
         return Pair(psi.startOffset, psi.endOffset)
     }

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/ConstantValueGenerator.kt

@@ -13,10 +13,10 @@ import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
-import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.constants.*
 import org.jetbrains.kotlin.types.KotlinType
@@ -31,6 +31,8 @@ abstract class ConstantValueGenerator(
 ) {
     protected abstract fun extractAnnotationOffsets(annotationDescriptor: AnnotationDescriptor): Pair<Int, Int>
 
+    protected abstract fun extractAnnotationParameterOffsets(annotationDescriptor: AnnotationDescriptor, argumentName: Name): Pair<Int, Int>
+
     private fun KotlinType.toIrType() = typeTranslator.translateType(this)
 
     fun generateConstantValueAsExpression(
@@ -192,7 +194,8 @@ abstract class ConstantValueGenerator(
             val argumentIndex = valueParameter.index
             val argumentValue = annotationDescriptor.allValueArguments[valueParameter.name] ?: continue
             val adjustedValue = adjustAnnotationArgumentValue(argumentValue, valueParameter)
-            val irArgument = generateAnnotationValueAsExpression(UNDEFINED_OFFSET, UNDEFINED_OFFSET, adjustedValue, valueParameter)
+            val (parameterStartOffset, parameterEndOffset) = extractAnnotationParameterOffsets(annotationDescriptor, valueParameter.name)
+            val irArgument = generateAnnotationValueAsExpression(parameterStartOffset, parameterEndOffset, adjustedValue, valueParameter)
             irCall.putValueArgument(argumentIndex, irArgument)
         }
 

core/descriptors/src/org/jetbrains/kotlin/descriptors/annotations/AnnotationDescriptor.kt

@@ -36,6 +36,8 @@ interface AnnotationDescriptor : AnnotationMarker {
     val allValueArguments: Map<Name, ConstantValue<*>>
 
     val source: SourceElement
+
+    fun getSourceForArgument(name: Name): SourceElement = SourceElement.NO_SOURCE
 }
 
 val AnnotationDescriptor.abbreviationFqName: FqName?

core/descriptors/src/org/jetbrains/kotlin/descriptors/annotations/AnnotationDescriptorImpl.java

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.renderer.DescriptorRenderer;
 import org.jetbrains.kotlin.resolve.constants.ConstantValue;
 import org.jetbrains.kotlin.types.KotlinType;
 
+import java.util.Collections;
 import java.util.Map;
 
 public class AnnotationDescriptorImpl implements AnnotationDescriptor {
@@ -66,6 +67,12 @@ public class AnnotationDescriptorImpl implements AnnotationDescriptor {
         return source;
     }
 
+    @NotNull
+    @Override
+    public SourceElement getSourceForArgument(@NotNull Name name) {
+        return SourceElement.NO_SOURCE;
+    }
+
     @Override
     public String toString() {
         return DescriptorRenderer.FQ_NAMES_IN_TYPES.renderAnnotation(this, null);