Add ability to set up source element when copying descriptor

compiler/frontend.java/src/org/jetbrains/kotlin/load/java/sam/SamAdapterFunctionDescriptor.java

@@ -21,6 +21,7 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor;
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor;
 import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor;
+import org.jetbrains.kotlin.descriptors.SourceElement;
 import org.jetbrains.kotlin.descriptors.annotations.Annotations;
 import org.jetbrains.kotlin.load.java.descriptors.JavaMethodDescriptor;
 import org.jetbrains.kotlin.load.java.descriptors.SamAdapterDescriptor;
@@ -58,7 +59,7 @@ import org.jetbrains.kotlin.name.Name;
             @NotNull Kind kind,
             @Nullable Name newName,
             @NotNull Annotations annotations,
-            boolean preserveSource
+            @NotNull SourceElement source
     ) {
         return new SamAdapterFunctionDescriptor(newOwner, (SimpleFunctionDescriptor) original, kind, declaration);
     }

compiler/frontend.java/src/org/jetbrains/kotlin/synthetic/SamAdapterFunctionsScope.kt

@@ -147,7 +147,7 @@ class SamAdapterFunctionsScope(storageManager: StorageManager) : SyntheticScope
                 kind: CallableMemberDescriptor.Kind,
                 newName: Name?,
                 annotations: Annotations,
-                preserveSource: Boolean
+                source: SourceElement
         ): MyFunctionDescriptor {
             return MyFunctionDescriptor(
                     containingDeclaration, original as SimpleFunctionDescriptor?, annotations, newName ?: name, kind, source

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/descriptors/LazyClassDescriptor.java

@@ -283,7 +283,7 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
                     @NotNull Kind kind,
                     @Nullable Name newName,
                     @NotNull Annotations annotations,
-                    boolean preserveSource
+                    @NotNull SourceElement source
             ) {
                 throw new UnsupportedOperationException();
             }

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/descriptors/JavaConstructorDescriptor.java

@@ -79,7 +79,7 @@ public class JavaConstructorDescriptor extends ConstructorDescriptorImpl impleme
             @NotNull Kind kind,
             @Nullable Name newName,
             @NotNull Annotations annotations,
-            boolean preserveSource
+            @NotNull SourceElement source
     ) {
         if (kind != Kind.DECLARATION && kind != Kind.SYNTHESIZED) {
             throw new IllegalStateException(
@@ -92,8 +92,8 @@ public class JavaConstructorDescriptor extends ConstructorDescriptorImpl impleme
 
         assert newName == null : "Attempt to rename constructor: " + this;
 
-        JavaConstructorDescriptor result = createDescriptor((ClassDescriptor) newOwner, (JavaConstructorDescriptor) original, kind,
+        JavaConstructorDescriptor result =
-                                                            getSourceToUseForCopy(preserveSource, original), annotations);
+                createDescriptor((ClassDescriptor) newOwner, (JavaConstructorDescriptor) original, kind, source, annotations);
         result.setHasStableParameterNames(hasStableParameterNames());
         result.setHasSynthesizedParameterNames(hasSynthesizedParameterNames());
         return result;
@@ -121,7 +121,7 @@ public class JavaConstructorDescriptor extends ConstructorDescriptorImpl impleme
             @NotNull KotlinType enhancedReturnType
     ) {
         JavaConstructorDescriptor enhanced = createSubstitutedCopy(
-                getContainingDeclaration(), /* original = */ null, getKind(), null, getAnnotations(), true);
+                getContainingDeclaration(), /* original = */ null, getKind(), null, getAnnotations(), getSource());
         // We do not use doSubstitute here as in JavaMethodDescriptor.enhance because type parameters of constructor belongs to class
         enhanced.initialize(
                 enhancedReceiverType,

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/descriptors/JavaMethodDescriptor.java

@@ -115,7 +115,7 @@ public class JavaMethodDescriptor extends SimpleFunctionDescriptorImpl implement
             @NotNull Kind kind,
             @Nullable Name newName,
             @NotNull Annotations annotations,
-            boolean preserveSource
+            @NotNull SourceElement source
     ) {
         JavaMethodDescriptor result = new JavaMethodDescriptor(
                 newOwner,
@@ -123,7 +123,7 @@ public class JavaMethodDescriptor extends SimpleFunctionDescriptorImpl implement
                 annotations,
                 newName != null ? newName : getName(),
                 kind,
-                getSourceToUseForCopy(preserveSource, original)
+                source
         );
         result.setParameterNamesStatus(hasStableParameterNames(), hasSynthesizedParameterNames());
         return result;

core/descriptors/src/org/jetbrains/kotlin/builtins/functions/FunctionInvokeDescriptor.kt

@@ -47,7 +47,7 @@ class FunctionInvokeDescriptor private constructor(
             kind: CallableMemberDescriptor.Kind,
             newName: Name?,
             annotations: Annotations,
-            preserveSource: Boolean
+            source: SourceElement
     ): FunctionDescriptorImpl {
         return FunctionInvokeDescriptor(newOwner, original as FunctionInvokeDescriptor?, kind)
     }

core/descriptors/src/org/jetbrains/kotlin/descriptors/FunctionDescriptor.java

@@ -128,6 +128,9 @@ public interface FunctionDescriptor extends CallableMemberDescriptor {
         @NotNull
         CopyBuilder<D> setPreserveSourceElement();
 
+        @NotNull
+        CopyBuilder<D> setSource(@NotNull SourceElement source);
+
         @NotNull
         CopyBuilder<D> setDropOriginalInContainingParts();
 

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/ConstructorDescriptorImpl.java

@@ -135,7 +135,7 @@ public class ConstructorDescriptorImpl extends FunctionDescriptorImpl implements
             @NotNull Kind kind,
             @Nullable Name newName,
             @NotNull Annotations annotations,
-            boolean preserveSource
+            @NotNull SourceElement source
     ) {
         if (kind != Kind.DECLARATION && kind != Kind.SYNTHESIZED) {
             throw new IllegalStateException("Attempt at creating a constructor that is not a declaration: \n" +
@@ -150,7 +150,7 @@ public class ConstructorDescriptorImpl extends FunctionDescriptorImpl implements
                 annotations,
                 isPrimary,
                 Kind.DECLARATION,
-                getSourceToUseForCopy(preserveSource, original)
+                source
         );
     }
 

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/FunctionDescriptorImpl.java

@@ -340,6 +340,7 @@ public abstract class FunctionDescriptorImpl extends DeclarationDescriptorNonRoo
         private List<TypeParameterDescriptor> newTypeParameters = null;
         private Annotations additionalAnnotations = null;
         private boolean isHiddenForResolutionEverywhereBesideSupercalls;
+        private SourceElement sourceElement;
 
         public CopyConfiguration(
                 @NotNull TypeSubstitution substitution,
@@ -464,6 +465,13 @@ public abstract class FunctionDescriptorImpl extends DeclarationDescriptorNonRoo
             return this;
         }
 
+        @NotNull
+        @Override
+        public CopyBuilder<FunctionDescriptor> setSource(@NotNull SourceElement source) {
+            this.sourceElement = source;
+            return this;
+        }
+
         @Override
         @NotNull
         public CopyConfiguration setDropOriginalInContainingParts() {
@@ -539,7 +547,7 @@ public abstract class FunctionDescriptorImpl extends DeclarationDescriptorNonRoo
 
         FunctionDescriptorImpl substitutedDescriptor = createSubstitutedCopy(
                 configuration.newOwner, configuration.original, configuration.kind, configuration.name, resultAnnotations,
-                configuration.preserveSourceElement);
+                getSourceToUseForCopy(configuration.preserveSourceElement, configuration.original, configuration.sourceElement));
 
         List<TypeParameterDescriptor> substitutedTypeParameters;
         final TypeSubstitutor substitutor;
@@ -655,11 +663,16 @@ public abstract class FunctionDescriptorImpl extends DeclarationDescriptorNonRoo
             @NotNull Kind kind,
             @Nullable Name newName,
             @NotNull Annotations annotations,
-            boolean preserveSource
+            @NotNull SourceElement source
     );
 
     @NotNull
-    protected SourceElement getSourceToUseForCopy(boolean preserveSource, @Nullable FunctionDescriptor original) {
+    private SourceElement getSourceToUseForCopy(
+            boolean preserveSource,
+            @Nullable FunctionDescriptor original,
+            @Nullable SourceElement sourceElement
+    ) {
+        if (sourceElement != null) return sourceElement;
         return preserveSource
                ? (original != null ? original.getSource() : getOriginal().getSource())
                : SourceElement.NO_SOURCE;

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/SimpleFunctionDescriptorImpl.java

@@ -78,7 +78,7 @@ public class SimpleFunctionDescriptorImpl extends FunctionDescriptorImpl impleme
             @NotNull Kind kind,
             @Nullable Name newName,
             @NotNull Annotations annotations,
-            boolean preserveSource
+            @NotNull SourceElement source
     ) {
         return new SimpleFunctionDescriptorImpl(
                 newOwner,
@@ -86,7 +86,7 @@ public class SimpleFunctionDescriptorImpl extends FunctionDescriptorImpl impleme
                 annotations,
                 newName != null ? newName : getName(),
                 kind,
-                getSourceToUseForCopy(preserveSource, original)
+                source
         );
     }
 

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/TypeAliasConstructorDescriptor.kt

@@ -63,7 +63,7 @@ class TypeAliasConstructorDescriptorImpl private constructor(
             kind: Kind,
             newName: Name?,
             annotations: Annotations,
-            preserveSource: Boolean
+            source: SourceElement
     ): TypeAliasConstructorDescriptorImpl {
         assert(kind == Kind.DECLARATION || kind == Kind.SYNTHESIZED) {
             "Creating a type alias constructor that is not a declaration: \ncopy from: ${this}\nnewOwner: $newOwner\nkind: $kind"
@@ -73,7 +73,7 @@ class TypeAliasConstructorDescriptorImpl private constructor(
                 typeAliasDescriptor,
                 newOwner as ClassDescriptor,
                 this, annotations, isPrimary, Kind.DECLARATION,
-                getSourceToUseForCopy(preserveSource, original))
+                source)
     }
 
     companion object {

core/descriptors/src/org/jetbrains/kotlin/types/error/ErrorSimpleFunctionDescriptorImpl.java

@@ -48,7 +48,7 @@ public class ErrorSimpleFunctionDescriptorImpl extends SimpleFunctionDescriptorI
             @NotNull Kind kind,
             @Nullable Name newName,
             @NotNull Annotations annotations,
-            boolean preserveSource
+            @NotNull SourceElement source
     ) {
         return this;
     }
@@ -153,6 +153,12 @@ public class ErrorSimpleFunctionDescriptorImpl extends SimpleFunctionDescriptorI
                 return this;
             }
 
+            @NotNull
+            @Override
+            public CopyBuilder<SimpleFunctionDescriptor> setSource(@NotNull SourceElement source) {
+                return this;
+            }
+
             @NotNull
             @Override
             public CopyBuilder<SimpleFunctionDescriptor> setDropOriginalInContainingParts() {

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/descriptors/DeserializedMemberDescriptor.kt

@@ -50,9 +50,12 @@ class DeserializedSimpleFunctionDescriptor(
         override val proto: ProtoBuf.Function,
         override val nameResolver: NameResolver,
         override val typeTable: TypeTable,
-        override val containerSource: SourceElement?
+        override val containerSource: SourceElement?,
+        source: SourceElement? = null
 ) : DeserializedCallableMemberDescriptor,
-        SimpleFunctionDescriptorImpl(containingDeclaration, original, annotations, name, kind, SourceElement.NO_SOURCE) {
+    SimpleFunctionDescriptorImpl(
+                containingDeclaration, original, annotations, name, kind,
+                source ?: org.jetbrains.kotlin.descriptors.SourceElement.NO_SOURCE) {
 
     override fun createSubstitutedCopy(
             newOwner: DeclarationDescriptor,
@@ -60,11 +63,11 @@ class DeserializedSimpleFunctionDescriptor(
             kind: CallableMemberDescriptor.Kind,
             newName: Name?,
             annotations: Annotations,
-            preserveSource: Boolean
+            source: SourceElement
     ): FunctionDescriptorImpl {
         return DeserializedSimpleFunctionDescriptor(
                 newOwner, original as SimpleFunctionDescriptor?, annotations, newName ?: name, kind,
-                proto, nameResolver, typeTable, containerSource
+                proto, nameResolver, typeTable, containerSource, source
         )
     }
 }
@@ -121,11 +124,11 @@ class DeserializedConstructorDescriptor(
             kind: CallableMemberDescriptor.Kind,
             newName: Name?,
             annotations: Annotations,
-            preserveSource: Boolean
+            source: SourceElement
     ): DeserializedConstructorDescriptor {
         return DeserializedConstructorDescriptor(
                 newOwner as ClassDescriptor, original as ConstructorDescriptor?, annotations, isPrimary, kind,
-                proto, nameResolver, typeTable, containerSource
+                proto, nameResolver, typeTable, source
         )
     }