StorageManager can create thread-safe memoized functions

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

@@ -17,6 +17,8 @@
 package org.jetbrains.jet.lang.resolve.lazy;
 
 import com.intellij.openapi.util.Computable;
+import com.intellij.util.Function;
+import com.intellij.util.containers.ConcurrentWeakValueHashMap;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.diagnostics.Diagnostic;
@@ -30,6 +32,8 @@ import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.ConcurrentMap;
 
+import static org.jetbrains.jet.lang.resolve.lazy.StorageManager.MemoizationMode.WEAK;
+
 public class LockBasedStorageManager implements StorageManager {
 
     private final Object lock = new Object() {
@@ -45,6 +49,35 @@ public class LockBasedStorageManager implements StorageManager {
         return new ConcurrentHashMap<K, V>();
     }
 
+    @NotNull
+    @Override
+    public <K, V> Function<K, V> createMemoizedFunction(@NotNull final Function<K, V> compute, @NotNull final MemoizationMode mode) {
+        return new Function<K, V>() {
+            private final ConcurrentMap<K, LazyValue<V>> cache;
+            {
+                cache = (mode == WEAK) ? new ConcurrentWeakValueHashMap<K, LazyValue<V>>() : new ConcurrentHashMap<K, LazyValue<V>>();
+            }
+
+            @Override
+            public V fun(@NotNull final K input) {
+                LazyValue<V> lazyValue = cache.get(input);
+                if (lazyValue != null) return lazyValue.get();
+
+                lazyValue = createLazyValue(new Computable<V>() {
+                    @Override
+                    public V compute() {
+                        return compute.fun(input);
+                    }
+                });
+
+                LazyValue<V> oldValue = cache.putIfAbsent(input, lazyValue);
+                if (oldValue != null) return oldValue.get();
+
+                return lazyValue.get();
+            }
+        };
+    }
+
     @Override
     public <E> Collection<E> createConcurrentCollection() {
         return new ConcurrentLinkedQueue<E>();

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

@@ -17,6 +17,7 @@
 package org.jetbrains.jet.lang.resolve.lazy;
 
 import com.intellij.openapi.util.Computable;
+import com.intellij.util.Function;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.resolve.BindingTrace;
 
@@ -27,6 +28,14 @@ public interface StorageManager {
     @NotNull
     <K, V> ConcurrentMap<K, V> createConcurrentMap();
 
+    /**
+     * Given a function compute: K -> V create a memoized version of it that computes a value only once for each key
+     * @param compute the function to be memoized
+     * @param modeForValues how to store teh memoized values
+     */
+    @NotNull
+    <K, V> Function<K, V> createMemoizedFunction(@NotNull Function<K, V> compute, @NotNull MemoizationMode modeForValues);
+
     <E> Collection<E> createConcurrentCollection();
 
     @NotNull
@@ -34,4 +43,9 @@ public interface StorageManager {
 
     @NotNull
     BindingTrace createSafeTrace(@NotNull BindingTrace originalTrace);
+
+    enum MemoizationMode {
+        STRONG,
+        WEAK
+    }
 }