Weakly retained memoized functions moved out of util.runtime

compiler/frontend/serialization/src/org/jetbrains/jet/descriptors/serialization/AbstractDescriptorFinder.java

@@ -25,8 +25,6 @@ import org.jetbrains.jet.lang.descriptors.ClassDescriptor;
 import org.jetbrains.jet.storage.MemoizedFunctionToNullable;
 import org.jetbrains.jet.storage.StorageManager;
 
-import static org.jetbrains.jet.storage.StorageManager.ReferenceKind.STRONG;
-
 public abstract class AbstractDescriptorFinder implements DescriptorFinder {
 
     private final MemoizedFunctionToNullable<ClassId, ClassDescriptor> findClass;
@@ -52,7 +50,7 @@ public abstract class AbstractDescriptorFinder implements DescriptorFinder {
                 classDescriptorCreated(classDescriptor);
                 return classDescriptor;
             }
-        }, STRONG);
+        });
     }
 
     @Nullable

compiler/frontend/serialization/src/org/jetbrains/jet/descriptors/serialization/TypeDeserializer.java

@@ -37,8 +37,6 @@ import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
-import static org.jetbrains.jet.storage.StorageManager.ReferenceKind.STRONG;
-
 public class TypeDeserializer {
 
     public interface TypeParameterResolver {
@@ -99,7 +97,7 @@ public class TypeDeserializer {
             public ClassDescriptor invoke(Integer fqNameIndex) {
                 return computeClassDescriptor(fqNameIndex);
             }
-        }, STRONG);
+        });
     }
 
     /* package */ DescriptorFinder getDescriptorFinder() {

compiler/frontend/serialization/src/org/jetbrains/jet/descriptors/serialization/descriptors/DeserializedClassDescriptor.java

@@ -45,7 +45,6 @@ import java.util.*;
 import static org.jetbrains.jet.descriptors.serialization.TypeDeserializer.TypeParameterResolver.NONE;
 import static org.jetbrains.jet.lang.descriptors.ReceiverParameterDescriptor.NO_RECEIVER_PARAMETER;
 import static org.jetbrains.jet.lang.resolve.DescriptorUtils.getClassObjectName;
-import static org.jetbrains.jet.storage.StorageManager.ReferenceKind.STRONG;
 
 public class DeserializedClassDescriptor extends AbstractClassDescriptor implements ClassDescriptor {
 
@@ -497,7 +496,7 @@ public class DeserializedClassDescriptor extends AbstractClassDescriptor impleme
                            descriptorFinder.findClass(classId.createNestedClassId(name)) :
                            null;
                 }
-            }, STRONG);
+            });
         }
 
         @NotNull

compiler/frontend/serialization/src/org/jetbrains/jet/descriptors/serialization/descriptors/DeserializedMemberScope.java

@@ -33,8 +33,6 @@ import org.jetbrains.jet.storage.StorageManager;
 
 import java.util.*;
 
-import static org.jetbrains.jet.storage.StorageManager.ReferenceKind.STRONG;
-
 public abstract class DeserializedMemberScope implements JetScope {
 
     private static final Filter<ProtoBuf.Callable.CallableKind> FUNCTION = new Filter<ProtoBuf.Callable.CallableKind>() {
@@ -78,13 +76,13 @@ public abstract class DeserializedMemberScope implements JetScope {
             public Collection<FunctionDescriptor> invoke(Name name) {
                 return computeFunctions(name);
             }
-        }, STRONG);
+        });
         this.properties = storageManager.createMemoizedFunction(new Function1<Name, Collection<VariableDescriptor>>() {
             @Override
             public Collection<VariableDescriptor> invoke(Name name) {
                 return computeProperties(name);
             }
-        }, STRONG);
+        });
         this.allDescriptors = storageManager.createLazyValue(new Function0<Collection<DeclarationDescriptor>>() {
             @Override
             public Collection<DeclarationDescriptor> invoke() {

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/LazyImportScope.java

@@ -39,7 +39,6 @@ import java.util.List;
 import java.util.Set;
 
 import static org.jetbrains.jet.lang.resolve.QualifiedExpressionResolver.LookupMode;
-import static org.jetbrains.jet.storage.StorageManager.ReferenceKind.STRONG;
 
 public class LazyImportScope implements JetScope {
     private final ResolveSession resolveSession;
@@ -135,7 +134,7 @@ public class LazyImportScope implements JetScope {
             public ImportDirectiveResolveCache invoke(JetImportDirective directive) {
                 return new ImportDirectiveResolveCache(directive);
             }
-        }, STRONG);
+        });
 
         NamespaceDescriptor rootPackageDescriptor = resolveSession.getPackageDescriptorByFqName(FqName.ROOT);
         if (rootPackageDescriptor == null) {

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/ScopeProvider.java

@@ -36,8 +36,6 @@ import org.jetbrains.jet.storage.NotNullLazyValue;
 import java.util.Collection;
 import java.util.List;
 
-import static org.jetbrains.jet.storage.StorageManager.ReferenceKind.WEAK;
-
 public class ScopeProvider {
     private final ResolveSession resolveSession;
 
@@ -48,12 +46,12 @@ public class ScopeProvider {
     public ScopeProvider(@NotNull ResolveSession resolveSession) {
         this.resolveSession = resolveSession;
 
-        this.fileScopes = resolveSession.getStorageManager().createMemoizedFunction(new Function1<JetFile, JetScope>() {
+        this.fileScopes = resolveSession.getStorageManager().createWeaklyRetainedMemoizedFunction(new Function1<JetFile, JetScope>() {
             @Override
             public JetScope invoke(@NotNull JetFile file) {
                 return createFileScope(file);
             }
-        }, WEAK);
+        });
 
         this.defaultImportsScope = resolveSession.getStorageManager().createLazyValue(new Function0<JetScope>() {
             @Override

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/declarations/FileBasedDeclarationProviderFactory.java

@@ -76,7 +76,7 @@ public class FileBasedDeclarationProviderFactory implements DeclarationProviderF
             public PackageMemberDeclarationProvider invoke(FqName fqName) {
                 return createPackageMemberDeclarationProvider(fqName);
             }
-        }, StorageManager.ReferenceKind.STRONG);
+        });
     }
 
     @NotNull

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/descriptors/AbstractLazyMemberScope.java

@@ -42,7 +42,6 @@ import java.util.List;
 import java.util.Set;
 
 import static org.jetbrains.jet.lang.resolve.lazy.ResolveSessionUtils.safeNameForLazyResolve;
-import static org.jetbrains.jet.storage.StorageManager.ReferenceKind.STRONG;
 
 public abstract class AbstractLazyMemberScope<D extends DeclarationDescriptor, DP extends DeclarationProvider> implements JetScope {
     protected final ResolveSession resolveSession;
@@ -77,26 +76,26 @@ public abstract class AbstractLazyMemberScope<D extends DeclarationDescriptor, D
             public List<ClassDescriptor> invoke(Name name) {
                 return resolveClassOrObjectDescriptor(name, false);
             }
-        }, STRONG);
+        });
         this.objectDescriptors = storageManager.createMemoizedFunction(new Function1<Name, List<ClassDescriptor>>() {
             @Override
             public List<ClassDescriptor> invoke(Name name) {
                 return resolveClassOrObjectDescriptor(name, true);
             }
-        }, STRONG);
+        });
 
         this.functionDescriptors = storageManager.createMemoizedFunction(new Function1<Name, Set<FunctionDescriptor>>() {
             @Override
             public Set<FunctionDescriptor> invoke(Name name) {
                 return doGetFunctions(name);
             }
-        }, STRONG);
+        });
         this.propertyDescriptors = storageManager.createMemoizedFunction(new Function1<Name, Set<VariableDescriptor>>() {
             @Override
             public Set<VariableDescriptor> invoke(Name name) {
                 return doGetProperties(name);
             }
-        }, STRONG);
+        });
 
         this.allDescriptors = storageManager.createLazyValue(new Function0<AllDescriptors>() {
             @Override

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/descriptors/LazyPackageMemberScope.java

@@ -29,7 +29,6 @@ import org.jetbrains.jet.lang.resolve.name.FqName;
 import org.jetbrains.jet.lang.resolve.name.Name;
 import org.jetbrains.jet.lang.resolve.scopes.JetScope;
 import org.jetbrains.jet.storage.MemoizedFunctionToNullable;
-import org.jetbrains.jet.storage.StorageManager;
 
 import java.util.Collection;
 import java.util.Set;
@@ -49,7 +48,7 @@ public class LazyPackageMemberScope extends AbstractLazyMemberScope<NamespaceDes
                     public NamespaceDescriptor invoke(Name name) {
                         return createPackageDescriptor(name);
                     }
-                }, StorageManager.ReferenceKind.STRONG);
+                });
     }
 
     @Override

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/storage/LazyResolveStorageManager.java

@@ -16,11 +16,20 @@
 
 package org.jetbrains.jet.lang.resolve.lazy.storage;
 
+import jet.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.resolve.BindingTrace;
+import org.jetbrains.jet.storage.MemoizedFunctionToNotNull;
+import org.jetbrains.jet.storage.MemoizedFunctionToNullable;
 import org.jetbrains.jet.storage.StorageManager;
 
 public interface LazyResolveStorageManager extends StorageManager {
+    @NotNull
+    <K, V> MemoizedFunctionToNotNull<K, V> createWeaklyRetainedMemoizedFunction(@NotNull Function1<K, V> compute);
+
+    @NotNull
+    <K, V> MemoizedFunctionToNullable<K, V> createWeaklyRetainedMemoizedFunctionWithNullableValues(@NotNull Function1<K, V> compute);
+
     @NotNull
     BindingTrace createSafeTrace(@NotNull BindingTrace originalTrace);
 }

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/storage/LockBasedLazyResolveStorageManager.java

@@ -17,6 +17,8 @@
 package org.jetbrains.jet.lang.resolve.lazy.storage;
 
 import com.google.common.collect.ImmutableMap;
+import com.intellij.util.containers.ConcurrentWeakValueHashMap;
+import jet.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.TestOnly;
@@ -25,6 +27,8 @@ import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.lang.resolve.BindingTrace;
 import org.jetbrains.jet.lang.resolve.Diagnostics;
 import org.jetbrains.jet.storage.LockBasedStorageManager;
+import org.jetbrains.jet.storage.MemoizedFunctionToNotNull;
+import org.jetbrains.jet.storage.MemoizedFunctionToNullable;
 import org.jetbrains.jet.util.slicedmap.ReadOnlySlice;
 import org.jetbrains.jet.util.slicedmap.WritableSlice;
 
@@ -32,6 +36,23 @@ import java.util.Collection;
 import java.util.concurrent.locks.Lock;
 
 public class LockBasedLazyResolveStorageManager extends LockBasedStorageManager implements LazyResolveStorageManager {
+
+    @Override
+    @NotNull
+    public <K, V> MemoizedFunctionToNotNull<K, V> createWeaklyRetainedMemoizedFunction(
+            @NotNull Function1<K, V> compute
+    ) {
+        return super.createMemoizedFunction(compute, new ConcurrentWeakValueHashMap<K, Object>());
+    }
+
+    @NotNull
+    @Override
+    public <K, V> MemoizedFunctionToNullable<K, V> createWeaklyRetainedMemoizedFunctionWithNullableValues(
+            @NotNull Function1<K, V> compute
+    ) {
+        return super.createMemoizedFunctionWithNullableValues(compute, new ConcurrentWeakValueHashMap<K, Object>());
+    }
+
     @NotNull
     @Override
     public BindingTrace createSafeTrace(@NotNull BindingTrace originalTrace) {

compiler/tests/org/jetbrains/jet/storage/StorageManagerTest.java

@@ -9,8 +9,6 @@ import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 
-import static org.jetbrains.jet.storage.StorageManager.ReferenceKind.STRONG;
-
 public class StorageManagerTest extends TestCase {
 
     private StorageManager m;
@@ -102,31 +100,31 @@ public class StorageManagerTest extends TestCase {
 
     public void testFunctionComputesOnce() throws Exception {
         CounterFunction counter = new CounterFunction();
-        MemoizedFunctionToNotNull<String, String> f = m.createMemoizedFunction(counter, STRONG);
+        MemoizedFunctionToNotNull<String, String> f = m.createMemoizedFunction(counter);
         doTestComputesOnce(apply(f, "ok"), "ok1", counter);
     }
 
     public void testNullableFunctionComputesOnce() throws Exception {
         CounterFunction counter = new CounterFunction();
-        MemoizedFunctionToNullable<String, String> f = m.createMemoizedFunctionWithNullableValues(counter, STRONG);
+        MemoizedFunctionToNullable<String, String> f = m.createMemoizedFunctionWithNullableValues(counter);
         doTestComputesOnce(apply(f, "ok"), "ok1", counter);
     }
 
     public void testNullIsNotConfusedForNotComputedInFunction() throws Exception {
         CounterFunctionToNull counter = new CounterFunctionToNull();
-        MemoizedFunctionToNullable<String, String> f = m.createMemoizedFunctionWithNullableValues(counter, STRONG);
+        MemoizedFunctionToNullable<String, String> f = m.createMemoizedFunctionWithNullableValues(counter);
         doTestComputesOnce(apply(f, ""), null, counter);
     }
 
     public void testFunctionPreservesExceptions() throws Exception {
         ExceptionCounterFunction counter = new ExceptionCounterFunction();
-        MemoizedFunctionToNotNull<String, String> f = m.createMemoizedFunction(counter, STRONG);
+        MemoizedFunctionToNotNull<String, String> f = m.createMemoizedFunction(counter);
         doTestExceptionPreserved(apply(f, ""), UnsupportedOperationException.class, counter);
     }
 
     public void testNullableFunctionPreservesExceptions() throws Exception {
         ExceptionCounterFunction counter = new ExceptionCounterFunction();
-        MemoizedFunctionToNullable<String, String> f = m.createMemoizedFunctionWithNullableValues(counter, STRONG);
+        MemoizedFunctionToNullable<String, String> f = m.createMemoizedFunctionWithNullableValues(counter);
         doTestExceptionPreserved(apply(f, ""), UnsupportedOperationException.class, counter);
     }
 

core/descriptor.loader.java/src/org/jetbrains/jet/lang/resolve/kotlin/AnnotationDescriptorDeserializer.java

@@ -42,7 +42,6 @@ import org.jetbrains.jet.lang.resolve.name.Name;
 import org.jetbrains.jet.lang.types.ErrorUtils;
 import org.jetbrains.jet.storage.LockBasedStorageManager;
 import org.jetbrains.jet.storage.MemoizedFunctionToNotNull;
-import org.jetbrains.jet.storage.StorageManager;
 
 import javax.inject.Inject;
 import java.io.IOException;
@@ -75,7 +74,7 @@ public class AnnotationDescriptorDeserializer implements AnnotationDeserializer
                                 return Collections.emptyMap();
                             }
                         }
-                    }, StorageManager.ReferenceKind.STRONG);
+                    });
 
     @Inject
     public void setJavaClassResolver(JavaClassResolver javaClassResolver) {

core/util.runtime/src/org/jetbrains/jet/storage/LockBasedStorageManager.java

@@ -16,7 +16,6 @@
 
 package org.jetbrains.jet.storage;
 
-import com.intellij.util.containers.ConcurrentWeakValueHashMap;
 import jet.Function0;
 import jet.Function1;
 import org.jetbrains.annotations.NotNull;
@@ -50,24 +49,30 @@ public class LockBasedStorageManager implements StorageManager {
 
     @NotNull
     @Override
-    public <K, V> MemoizedFunctionToNotNull<K, V> createMemoizedFunction(
+    public <K, V> MemoizedFunctionToNotNull<K, V> createMemoizedFunction(@NotNull Function1<K, V> compute) {
-            @NotNull Function1<K, V> compute, @NotNull ReferenceKind valuesReferenceKind
+        return createMemoizedFunction(compute, new ConcurrentHashMap<K, Object>());
+    }
+
+    @NotNull
+    protected  <K, V> MemoizedFunctionToNotNull<K, V> createMemoizedFunction(
+            @NotNull Function1<K, V> compute,
+            @NotNull ConcurrentMap<K, Object> map
     ) {
-        ConcurrentMap<K, Object> map = createConcurrentMap(valuesReferenceKind);
         return new MapBasedMemoizedFunctionToNotNull<K, V>(lock, map, compute);
     }
 
     @NotNull
     @Override
-    public <K, V> MemoizedFunctionToNullable<K, V> createMemoizedFunctionWithNullableValues(
+    public <K, V> MemoizedFunctionToNullable<K, V> createMemoizedFunctionWithNullableValues(@NotNull Function1<K, V> compute) {
-            @NotNull Function1<K, V> compute, @NotNull ReferenceKind valuesReferenceKind
+        return createMemoizedFunctionWithNullableValues(compute, new ConcurrentHashMap<K, Object>());
-    ) {
-        ConcurrentMap<K, Object> map = createConcurrentMap(valuesReferenceKind);
-        return new MapBasedMemoizedFunction<K, V>(lock, map, compute);
     }
 
-    private static <K, V> ConcurrentMap<K, V> createConcurrentMap(ReferenceKind referenceKind) {
+    @NotNull
-        return (referenceKind == ReferenceKind.WEAK) ? new ConcurrentWeakValueHashMap<K, V>() : new ConcurrentHashMap<K, V>();
+    protected <K, V> MemoizedFunctionToNullable<K, V> createMemoizedFunctionWithNullableValues(
+            @NotNull Function1<K, V> compute,
+            @NotNull ConcurrentMap<K, Object> map
+    ) {
+        return new MapBasedMemoizedFunction<K, V>(lock, map, compute);
     }
 
     @NotNull

core/util.runtime/src/org/jetbrains/jet/storage/StorageManager.java

@@ -30,9 +30,9 @@ public interface StorageManager {
      * NOTE: if compute() has side-effects the WEAK reference kind is dangerous: the side-effects will be repeated if
      */
     @NotNull
-    <K, V> MemoizedFunctionToNotNull<K, V> createMemoizedFunction(@NotNull Function1<K, V> compute, @NotNull ReferenceKind valuesReferenceKind);
+    <K, V> MemoizedFunctionToNotNull<K, V> createMemoizedFunction(@NotNull Function1<K, V> compute);
     @NotNull
-    <K, V> MemoizedFunctionToNullable<K, V> createMemoizedFunctionWithNullableValues(@NotNull Function1<K, V> compute, @NotNull ReferenceKind valuesReferenceKind);
+    <K, V> MemoizedFunctionToNullable<K, V> createMemoizedFunctionWithNullableValues(@NotNull Function1<K, V> compute);
 
     @NotNull
     <T> NotNullLazyValue<T> createLazyValue(@NotNull Function0<T> computable);
@@ -69,9 +69,4 @@ public interface StorageManager {
     <T> NullableLazyValue<T> createNullableLazyValueWithPostCompute(@NotNull Function0<T> computable, @NotNull Function1<T, Void> postCompute);
 
     <T> T compute(@NotNull Function0<T> computable);
-
-    enum ReferenceKind {
-        STRONG,
-        WEAK
-    }
 }

idea/src/org/jetbrains/jet/plugin/project/ResolveElementCache.java

@@ -39,13 +39,13 @@ import org.jetbrains.jet.lang.resolve.lazy.ResolveSession;
 import org.jetbrains.jet.lang.resolve.lazy.ScopeProvider;
 import org.jetbrains.jet.lang.resolve.lazy.descriptors.LazyClassDescriptor;
 import org.jetbrains.jet.lang.resolve.lazy.descriptors.LazyPackageDescriptor;
-import org.jetbrains.jet.storage.MemoizedFunctionToNotNull;
+import org.jetbrains.jet.lang.resolve.lazy.storage.LazyResolveStorageManager;
-import org.jetbrains.jet.storage.StorageManager;
 import org.jetbrains.jet.lang.resolve.name.FqName;
 import org.jetbrains.jet.lang.resolve.name.Name;
 import org.jetbrains.jet.lang.resolve.scopes.JetScope;
 import org.jetbrains.jet.lang.resolve.scopes.WritableScope;
 import org.jetbrains.jet.lang.types.TypeConstructor;
+import org.jetbrains.jet.storage.MemoizedFunctionToNotNull;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -66,14 +66,14 @@ public class ResolveElementCache {
                             @Nullable
                             @Override
                             public Result<MemoizedFunctionToNotNull<JetElement, BindingContext>> compute() {
-                                StorageManager manager = ResolveElementCache.this.resolveSession.getStorageManager();
+                                LazyResolveStorageManager manager = ResolveElementCache.this.resolveSession.getStorageManager();
                                 MemoizedFunctionToNotNull<JetElement, BindingContext> elementsCacheFunction =
-                                        manager.createMemoizedFunction(new Function1<JetElement, BindingContext>() {
+                                        manager.createWeaklyRetainedMemoizedFunction(new Function1<JetElement, BindingContext>() {
                                             @Override
                                             public BindingContext invoke(JetElement jetElement) {
                                                 return elementAdditionalResolve(jetElement);
                                             }
-                                        }, StorageManager.ReferenceKind.WEAK);
+                                        });
 
                                 return Result.create(elementsCacheFunction, PsiModificationTracker.MODIFICATION_COUNT);
                             }