From 106753122449f5d3da1c76d6cb757afd8193fedb Mon Sep 17 00:00:00 2001
From: Denis Zharkov <denis.zharkov@jetbrains.com>
Date: Wed, 1 Jun 2016 12:50:59 +0300
Subject: [PATCH] Introduce light-weight user data parts into
 FunctionDescriptor

---
 .../descriptors/FunctionDescriptor.java       |  9 +++++
 .../impl/FunctionDescriptorImpl.java          | 38 +++++++++++++++++++
 .../impl/PropertyAccessorDescriptorImpl.java  |  6 +++
 .../ErrorSimpleFunctionDescriptorImpl.java    | 14 +++++++
 4 files changed, 67 insertions(+)

diff --git a/core/descriptors/src/org/jetbrains/kotlin/descriptors/FunctionDescriptor.java b/core/descriptors/src/org/jetbrains/kotlin/descriptors/FunctionDescriptor.java
index dd0cf664aca..c09bc57a52f 100644
--- a/core/descriptors/src/org/jetbrains/kotlin/descriptors/FunctionDescriptor.java
+++ b/core/descriptors/src/org/jetbrains/kotlin/descriptors/FunctionDescriptor.java
@@ -82,6 +82,12 @@ public interface FunctionDescriptor extends CallableMemberDescriptor {
 
     boolean isSuspend();
 
+    interface UserDataKey<V> {}
+
+    // TODO: pull up userdata related members to DeclarationDescriptor and use more efficient implementation (e.g. THashMap)
+    @Nullable
+    <V> V getUserData(UserDataKey<V> key);
+
     @NotNull
     CopyBuilder<? extends FunctionDescriptor> newCopyBuilder();
 
@@ -146,6 +152,9 @@ public interface FunctionDescriptor extends CallableMemberDescriptor {
         @NotNull
         CopyBuilder<D> setSubstitution(@NotNull TypeSubstitution substitution);
 
+        @NotNull
+        <V> CopyBuilder<D> putUserData(@NotNull UserDataKey<V> userDataKey, V value);
+
         @Nullable
         D build();
     }
diff --git a/core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/FunctionDescriptorImpl.java b/core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/FunctionDescriptorImpl.java
index e58c43b988b..7bc0a86e4da 100644
--- a/core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/FunctionDescriptorImpl.java
+++ b/core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/FunctionDescriptorImpl.java
@@ -58,6 +58,8 @@ public abstract class FunctionDescriptorImpl extends DeclarationDescriptorNonRoo
     @Nullable
     private FunctionDescriptor initialSignatureDescriptor = null;
 
+    private Map<UserDataKey<?>, Object> userDataMap = null;
+
     protected FunctionDescriptorImpl(
             @NotNull DeclarationDescriptor containingDeclaration,
             @Nullable FunctionDescriptor original,
@@ -245,6 +247,13 @@ public abstract class FunctionDescriptorImpl extends DeclarationDescriptorNonRoo
         return isSuspend;
     }
 
+    @Override
+    public <V> V getUserData(UserDataKey<V> key) {
+        if (userDataMap == null) return null;
+        //noinspection unchecked
+        return (V) userDataMap.get(key);
+    }
+
     @Override
     public boolean isHiddenToOvercomeSignatureClash() {
         return isHiddenToOvercomeSignatureClash;
@@ -341,6 +350,7 @@ public abstract class FunctionDescriptorImpl extends DeclarationDescriptorNonRoo
         private Annotations additionalAnnotations = null;
         private boolean isHiddenForResolutionEverywhereBesideSupercalls = isHiddenForResolutionEverywhereBesideSupercalls();
         private SourceElement sourceElement;
+        private Map<UserDataKey<?>, Object> userDataMap = new LinkedHashMap<UserDataKey<?>, Object>();
 
         public CopyConfiguration(
                 @NotNull TypeSubstitution substitution,
@@ -504,6 +514,13 @@ public abstract class FunctionDescriptorImpl extends DeclarationDescriptorNonRoo
             return this;
         }
 
+        @NotNull
+        @Override
+        public <V> CopyBuilder<FunctionDescriptor> putUserData(@NotNull UserDataKey<V> userDataKey, V value) {
+            userDataMap.put(userDataKey, value);
+            return this;
+        }
+
         @Override
         @Nullable
         public FunctionDescriptor build() {
@@ -620,6 +637,27 @@ public abstract class FunctionDescriptorImpl extends DeclarationDescriptorNonRoo
         substitutedDescriptor.setHiddenToOvercomeSignatureClash(configuration.isHiddenToOvercomeSignatureClash);
         substitutedDescriptor.setHiddenForResolutionEverywhereBesideSupercalls(configuration.isHiddenForResolutionEverywhereBesideSupercalls);
 
+        if (!configuration.userDataMap.isEmpty() || userDataMap != null) {
+            Map<UserDataKey<?>, Object> newMap = configuration.userDataMap;
+
+            if (userDataMap != null) {
+                for (Map.Entry<UserDataKey<?>, Object> entry : userDataMap.entrySet()) {
+                    if (!newMap.containsKey(entry.getKey())) {
+                        newMap.put(entry.getKey(), entry.getValue());
+                    }
+                }
+            }
+
+            if (newMap.size() == 1) {
+                substitutedDescriptor.userDataMap =
+                        Collections.<UserDataKey<?>, Object>singletonMap(
+                                newMap.keySet().iterator().next(), newMap.values().iterator().next());
+            }
+            else {
+                substitutedDescriptor.userDataMap = newMap;
+            }
+        }
+
         if (configuration.signatureChange || getInitialSignatureDescriptor() != null) {
             FunctionDescriptor initialSignature = (getInitialSignatureDescriptor() != null ? getInitialSignatureDescriptor() : this);
             FunctionDescriptor initialSignatureSubstituted = initialSignature.substitute(substitutor);
diff --git a/core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/PropertyAccessorDescriptorImpl.java b/core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/PropertyAccessorDescriptorImpl.java
index 91affbb6117..97a17355ca8 100644
--- a/core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/PropertyAccessorDescriptorImpl.java
+++ b/core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/PropertyAccessorDescriptorImpl.java
@@ -219,4 +219,10 @@ public abstract class PropertyAccessorDescriptorImpl extends DeclarationDescript
     public boolean isHiddenForResolutionEverywhereBesideSupercalls() {
         return false;
     }
+
+    @Nullable
+    @Override
+    public <V> V getUserData(UserDataKey<V> key) {
+        return null;
+    }
 }
diff --git a/core/descriptors/src/org/jetbrains/kotlin/types/error/ErrorSimpleFunctionDescriptorImpl.java b/core/descriptors/src/org/jetbrains/kotlin/types/error/ErrorSimpleFunctionDescriptorImpl.java
index 9e14ebf0cc1..e2fedeee101 100644
--- a/core/descriptors/src/org/jetbrains/kotlin/types/error/ErrorSimpleFunctionDescriptorImpl.java
+++ b/core/descriptors/src/org/jetbrains/kotlin/types/error/ErrorSimpleFunctionDescriptorImpl.java
@@ -111,6 +111,15 @@ public class ErrorSimpleFunctionDescriptorImpl extends SimpleFunctionDescriptorI
                 return this;
             }
 
+            @NotNull
+            @Override
+            public <V> CopyBuilder<SimpleFunctionDescriptor> putUserData(
+                    @NotNull UserDataKey<V> userDataKey,
+                    V value
+            ) {
+                return this;
+            }
+
             @NotNull
             @Override
             public CopyBuilder<SimpleFunctionDescriptor> setTypeParameters(@NotNull List<TypeParameterDescriptor> parameters) {
@@ -196,6 +205,11 @@ public class ErrorSimpleFunctionDescriptorImpl extends SimpleFunctionDescriptorI
         return false;
     }
 
+    @Override
+    public <V> V getUserData(UserDataKey<V> key) {
+        return null;
+    }
+
     @Override
     public void setOverriddenDescriptors(@NotNull Collection<? extends CallableMemberDescriptor> overriddenDescriptors) {
         // nop