Migrating to memoized functions and nullable lazy values

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

@@ -17,6 +17,7 @@
 package org.jetbrains.jet.lang.resolve.lazy;
 
 import com.google.common.collect.Sets;
+import com.intellij.util.Function;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.*;
@@ -26,21 +27,24 @@ import org.jetbrains.jet.lang.resolve.name.LabelName;
 import org.jetbrains.jet.lang.resolve.name.Name;
 import org.jetbrains.jet.lang.resolve.scopes.JetScope;
 
-import java.util.*;
+import java.util.Collection;
-import java.util.concurrent.ConcurrentMap;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
 
 import static org.jetbrains.jet.lang.resolve.lazy.ResolveSessionUtils.safeNameForLazyResolve;
+import static org.jetbrains.jet.lang.resolve.lazy.StorageManager.MemoizationMode.STRONG;
 
 public abstract class AbstractLazyMemberScope<D extends DeclarationDescriptor, DP extends DeclarationProvider> implements JetScope {
     protected final ResolveSession resolveSession;
     protected final DP declarationProvider;
     protected final D thisDescriptor;
 
-    private final ConcurrentMap<Name, ClassDescriptor> classDescriptors;
+    private final Function<Name, ClassDescriptor> classDescriptors;
-    private final ConcurrentMap<Name, ClassDescriptor> objectDescriptors;
+    private final Function<Name, ClassDescriptor> objectDescriptors;
 
-    private final ConcurrentMap<Name, Set<FunctionDescriptor>> functionDescriptors;
+    private final Function<Name, Set<FunctionDescriptor>> functionDescriptors;
-    private final ConcurrentMap<Name, Set<VariableDescriptor>> propertyDescriptors;
+    private final Function<Name, Set<VariableDescriptor>> propertyDescriptors;
 
     private final Collection<DeclarationDescriptor> allDescriptors;
     protected volatile boolean allDescriptorsComputed = false;
@@ -55,25 +59,40 @@ public abstract class AbstractLazyMemberScope<D extends DeclarationDescriptor, D
         this.thisDescriptor = thisDescriptor;
 
         StorageManager storageManager = resolveSession.getStorageManager();
-        this.classDescriptors = storageManager.createConcurrentMap();
+        this.classDescriptors = storageManager.createMemoizedFunctionWithNullableValues(new Function<Name, ClassDescriptor>() {
-        this.objectDescriptors = storageManager.createConcurrentMap();
+            @Override
-        this.functionDescriptors = storageManager.createConcurrentMap();
+            public ClassDescriptor fun(Name name) {
-        this.propertyDescriptors = storageManager.createConcurrentMap();
+                return resolveClassOrObjectDescriptor(name, false);
+            }
+        }, STRONG);
+        this.objectDescriptors = storageManager.createMemoizedFunctionWithNullableValues(new Function<Name, ClassDescriptor>() {
+            @Override
+            public ClassDescriptor fun(Name name) {
+                return resolveClassOrObjectDescriptor(name, true);
+            }
+        }, STRONG);
+
+        this.functionDescriptors = storageManager.createMemoizedFunction(new Function<Name, Set<FunctionDescriptor>>() {
+            @Override
+            public Set<FunctionDescriptor> fun(Name name) {
+                return doGetFunctions(name);
+            }
+        }, STRONG);
+        this.propertyDescriptors = storageManager.createMemoizedFunction(new Function<Name, Set<VariableDescriptor>>() {
+            @Override
+            public Set<VariableDescriptor> fun(Name name) {
+                return doGetProperties(name);
+            }
+        }, STRONG);
+
         this.allDescriptors = storageManager.createConcurrentCollection();
     }
 
     @Nullable
-    private ClassDescriptor getClassOrObjectDescriptor(@NotNull ConcurrentMap<Name, ClassDescriptor> cache, @NotNull Name name, boolean object) {
+    private ClassDescriptor resolveClassOrObjectDescriptor(@NotNull Name name, boolean object) {
-        ClassDescriptor known = cache.get(name);
-        if (known != null) return known;
-
-        if (allDescriptorsComputed) return null;
-
         Collection<JetClassOrObject> classOrObjectDeclarations = declarationProvider.getClassOrObjectDeclarations(name);
         if (classOrObjectDeclarations.isEmpty()) return null;
 
-        ClassDescriptor result = null;
-
         for (JetClassOrObject classOrObjectDeclaration : classOrObjectDeclarations) {
 
             // TODO: when enum entries with constructors are dropped, replace with declaresObjectOrEnumConstant()
@@ -82,16 +101,13 @@ public abstract class AbstractLazyMemberScope<D extends DeclarationDescriptor, D
             ClassDescriptor classDescriptor = new LazyClassDescriptor(resolveSession, thisDescriptor, name,
                                                                       JetClassInfoUtil.createClassLikeInfo(classOrObjectDeclaration));
 
-            ClassDescriptor oldValue = cache.putIfAbsent(name, classDescriptor);
-            if (oldValue != null) return oldValue;
-
             if (!object) {
                 registerDescriptor(classDescriptor);
             }
 
-            result = classDescriptor;
+            return classDescriptor;
         }
-        return result;
+        return null;
     }
 
     private static boolean declaresObjectOrEnumConstant(JetClassOrObject declaration) {
@@ -100,23 +116,22 @@ public abstract class AbstractLazyMemberScope<D extends DeclarationDescriptor, D
 
     @Override
     public ClassifierDescriptor getClassifier(@NotNull Name name) {
-        return getClassOrObjectDescriptor(classDescriptors, name, false);
+        return classDescriptors.fun(name);
     }
 
     @Override
     public ClassDescriptor getObjectDescriptor(@NotNull Name name) {
-        return getClassOrObjectDescriptor(objectDescriptors, name, true);
+        return objectDescriptors.fun(name);
     }
 
     @NotNull
     @Override
     public Set<FunctionDescriptor> getFunctions(@NotNull Name name) {
-        Set<FunctionDescriptor> known = functionDescriptors.get(name);
+        return functionDescriptors.fun(name);
-        if (known != null) return known;
+    }
-
-        // If all descriptors are already computed, we are
-        if (allDescriptorsComputed) return Collections.emptySet();
 
+    @NotNull
+    private Set<FunctionDescriptor> doGetFunctions(@NotNull Name name) {
         Set<FunctionDescriptor> result = Sets.newLinkedHashSet();
 
         Collection<JetNamedFunction> declarations = declarationProvider.getFunctionDeclarations(name);
@@ -130,9 +145,6 @@ public abstract class AbstractLazyMemberScope<D extends DeclarationDescriptor, D
         getNonDeclaredFunctions(name, result);
 
         if (!result.isEmpty()) {
-            Set<FunctionDescriptor> oldValue = functionDescriptors.putIfAbsent(name, result);
-            if (oldValue != null) return oldValue;
-
             registerDescriptors(result);
         }
         return result;
@@ -146,12 +158,11 @@ public abstract class AbstractLazyMemberScope<D extends DeclarationDescriptor, D
     @NotNull
     @Override
     public Set<VariableDescriptor> getProperties(@NotNull Name name) {
-        Set<VariableDescriptor> known = propertyDescriptors.get(name);
+        return propertyDescriptors.fun(name);
-        if (known != null) return known;
+    }
-
-        // If all descriptors are already computed, we are
-        if (allDescriptorsComputed) return Collections.emptySet();
 
+    @NotNull
+    public Set<VariableDescriptor> doGetProperties(@NotNull Name name) {
         Set<VariableDescriptor> result = Sets.newLinkedHashSet();
 
         Collection<JetProperty> declarations = declarationProvider.getPropertyDeclarations(name);
@@ -178,12 +189,8 @@ public abstract class AbstractLazyMemberScope<D extends DeclarationDescriptor, D
 
         getNonDeclaredProperties(name, result);
 
-        if (!result.isEmpty()) {
+        registerDescriptors(result);
-            Set<VariableDescriptor> oldValue = propertyDescriptors.putIfAbsent(name, result);
-            if (oldValue != null) return oldValue;
 
-            registerDescriptors(result);
-        }
         return result;
     }
 

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

@@ -23,7 +23,9 @@ import com.google.common.collect.HashMultimap;
 import com.google.common.collect.Multimap;
 import com.google.common.collect.Sets;
 import com.intellij.openapi.util.Computable;
+import com.intellij.util.Function;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.JetFile;
 import org.jetbrains.jet.lang.psi.JetNamespaceHeader;
 import org.jetbrains.jet.lang.resolve.lazy.data.JetClassLikeInfo;
@@ -31,7 +33,6 @@ import org.jetbrains.jet.lang.resolve.name.FqName;
 
 import java.util.Collection;
 import java.util.Set;
-import java.util.concurrent.ConcurrentMap;
 
 public class FileBasedDeclarationProviderFactory implements DeclarationProviderFactory {
 
@@ -45,7 +46,7 @@ public class FileBasedDeclarationProviderFactory implements DeclarationProviderF
     private final StorageManager storageManager;
     private final LazyValue<Index> index;
 
-    private final ConcurrentMap<FqName, PackageMemberDeclarationProvider> packageDeclarationProviders;
+    private final Function<FqName, PackageMemberDeclarationProvider> packageDeclarationProviders;
 
     public FileBasedDeclarationProviderFactory(@NotNull StorageManager storageManager, @NotNull Collection<JetFile> files) {
         this(storageManager, files, Predicates.<FqName>alwaysFalse());
@@ -64,7 +65,12 @@ public class FileBasedDeclarationProviderFactory implements DeclarationProviderF
                 return computeFilesByPackage(files);
             }
         });
-        this.packageDeclarationProviders = storageManager.createConcurrentMap();
+        this.packageDeclarationProviders = storageManager.createMemoizedFunctionWithNullableValues(new Function<FqName, PackageMemberDeclarationProvider>() {
+            @Override
+            public PackageMemberDeclarationProvider fun(FqName fqName) {
+                return createPackageMemberDeclarationProvider(fqName);
+            }
+        }, StorageManager.MemoizationMode.STRONG);
     }
 
     @NotNull
@@ -109,11 +115,11 @@ public class FileBasedDeclarationProviderFactory implements DeclarationProviderF
 
     @Override
     public PackageMemberDeclarationProvider getPackageMemberDeclarationProvider(@NotNull FqName packageFqName) {
-        PackageMemberDeclarationProvider cached = packageDeclarationProviders.get(packageFqName);
+        return packageDeclarationProviders.fun(packageFqName);
-        if (cached != null) {
+    }
-            return cached;
-        }
 
+    @Nullable
+    public PackageMemberDeclarationProvider createPackageMemberDeclarationProvider(@NotNull FqName packageFqName) {
         if (!isPackageDeclaredExplicitly(packageFqName)) {
             if (isPackageDeclaredExternally.apply(packageFqName)) {
                 return EmptyPackageMemberDeclarationProvider.INSTANCE;
@@ -121,15 +127,7 @@ public class FileBasedDeclarationProviderFactory implements DeclarationProviderF
             return null;
         }
 
-        FileBasedPackageMemberDeclarationProvider provider =
+        return new FileBasedPackageMemberDeclarationProvider(storageManager, packageFqName, this, index.get().filesByPackage.get(packageFqName));
-                new FileBasedPackageMemberDeclarationProvider(storageManager, packageFqName, this, index.get().filesByPackage.get(packageFqName));
-
-        PackageMemberDeclarationProvider oldValue = packageDeclarationProviders.putIfAbsent(packageFqName, provider);
-        if (oldValue != null) {
-            return oldValue;
-        }
-
-        return provider;
     }
 
     @NotNull

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

@@ -16,7 +16,6 @@
 
 package org.jetbrains.jet.lang.resolve.lazy;
 
-import com.google.common.base.Optional;
 import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
 import com.google.common.collect.Collections2;
@@ -71,7 +70,7 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements LazyDesc
 
     private final LazyValue<ReceiverParameterDescriptor> thisAsReceiverParameter;
     private final LazyValue<List<AnnotationDescriptor>> annotations;
-    private final LazyValue<Optional<ClassDescriptor>> classObjectDescriptor;
+    private final LazyValue<ClassDescriptor> classObjectDescriptor;
 
     private final LazyClassMemberScope unsubstitutedMemberScope;
     private final JetScope unsubstitutedInnerClassesScope;
@@ -129,10 +128,10 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements LazyDesc
                 return resolveAnnotations();
             }
         });
-        this.classObjectDescriptor = storageManager.createLazyValue(new Computable<Optional<ClassDescriptor>>() {
+        this.classObjectDescriptor = storageManager.createNullableLazyValue(new Computable<ClassDescriptor>() {
             @Override
-            public Optional<ClassDescriptor> compute() {
+            public ClassDescriptor compute() {
-                return Optional.fromNullable(computeClassObjectDescriptor());
+                return computeClassObjectDescriptor();
             }
         });
         this.scopeForClassHeaderResolution = storageManager.createLazyValue(new Computable<JetScope>() {
@@ -273,7 +272,7 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements LazyDesc
 
     @Override
     public ClassDescriptor getClassObjectDescriptor() {
-        return classObjectDescriptor.get().orNull();
+        return classObjectDescriptor.get();
     }
 
     @Nullable

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

@@ -16,7 +16,6 @@
 
 package org.jetbrains.jet.lang.resolve.lazy;
 
-import com.google.common.base.Optional;
 import com.google.common.collect.Lists;
 import com.intellij.openapi.util.Computable;
 import com.intellij.psi.PsiElement;
@@ -66,7 +65,7 @@ public class LazyClassMemberScope extends AbstractLazyMemberScope<LazyClassDescr
         Collection<T> extract(@NotNull JetType extractFrom, @NotNull Name name);
     }
 
-    private final LazyValue<Optional<ConstructorDescriptor>> primaryConstructor;
+    private final LazyValue<ConstructorDescriptor> primaryConstructor;
 
     public LazyClassMemberScope(
             @NotNull ResolveSession resolveSession,
@@ -74,10 +73,10 @@ public class LazyClassMemberScope extends AbstractLazyMemberScope<LazyClassDescr
             @NotNull LazyClassDescriptor thisClass
     ) {
         super(resolveSession, declarationProvider, thisClass);
-        this.primaryConstructor = resolveSession.getStorageManager().createLazyValue(new Computable<Optional<ConstructorDescriptor>>() {
+        this.primaryConstructor = resolveSession.getStorageManager().createNullableLazyValue(new Computable<ConstructorDescriptor>() {
             @Override
-            public Optional<ConstructorDescriptor> compute() {
+            public ConstructorDescriptor compute() {
-                return Optional.fromNullable(resolvePrimaryConstructor());
+                return resolvePrimaryConstructor();
             }
         });
     }
@@ -338,7 +337,7 @@ public class LazyClassMemberScope extends AbstractLazyMemberScope<LazyClassDescr
 
     @Nullable
     public ConstructorDescriptor getPrimaryConstructor() {
-        return primaryConstructor.get().orNull();
+        return primaryConstructor.get();
     }
 
     @Nullable

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

@@ -16,7 +16,9 @@
 
 package org.jetbrains.jet.lang.resolve.lazy;
 
+import com.intellij.util.Function;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.psi.JetDeclaration;
 import org.jetbrains.jet.lang.psi.JetFile;
@@ -26,27 +28,31 @@ import org.jetbrains.jet.lang.resolve.name.Name;
 import org.jetbrains.jet.lang.resolve.scopes.JetScope;
 
 import java.util.Set;
-import java.util.concurrent.ConcurrentMap;
 
 public class LazyPackageMemberScope extends AbstractLazyMemberScope<NamespaceDescriptor, PackageMemberDeclarationProvider> {
 
-    private final ConcurrentMap<Name, NamespaceDescriptor> packageDescriptors;
+    private final Function<Name, NamespaceDescriptor> packageDescriptors;
 
     public LazyPackageMemberScope(@NotNull ResolveSession resolveSession,
             @NotNull PackageMemberDeclarationProvider declarationProvider,
             @NotNull NamespaceDescriptor thisPackage) {
         super(resolveSession, declarationProvider, thisPackage);
 
-        this.packageDescriptors = resolveSession.getStorageManager().createConcurrentMap();
+        this.packageDescriptors = resolveSession.getStorageManager().createMemoizedFunctionWithNullableValues(new Function<Name, NamespaceDescriptor>() {
+            @Override
+            public NamespaceDescriptor fun(Name name) {
+                return createPackageDescriptor(name);
+            }
+        }, StorageManager.MemoizationMode.STRONG);
     }
 
     @Override
     public NamespaceDescriptor getNamespace(@NotNull Name name) {
-        NamespaceDescriptor known = packageDescriptors.get(name);
+        return packageDescriptors.fun(name);
-        if (known != null) return known;
+    }
-
-        if (allDescriptorsComputed) return null;
 
+    @Nullable
+    public NamespaceDescriptor createPackageDescriptor(@NotNull Name name) {
         if (!declarationProvider.isPackageDeclared(name)) return null;
 
         PackageMemberDeclarationProvider packageMemberDeclarationProvider = resolveSession.getDeclarationProviderFactory().getPackageMemberDeclarationProvider(
@@ -54,9 +60,6 @@ public class LazyPackageMemberScope extends AbstractLazyMemberScope<NamespaceDes
         assert packageMemberDeclarationProvider != null : "Package is declared, but declaration provider is not found: " + name;
         NamespaceDescriptor namespaceDescriptor = new LazyPackageDescriptor(thisDescriptor, name, resolveSession, packageMemberDeclarationProvider);
 
-        NamespaceDescriptor oldValue = packageDescriptors.putIfAbsent(name, namespaceDescriptor);
-        if (oldValue != null) return oldValue;
-
         registerDescriptor(namespaceDescriptor);
 
         return namespaceDescriptor;