Simple (not thread-safe) memoized functions supported

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

@@ -27,6 +27,7 @@ import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.lang.resolve.BindingTrace;
 import org.jetbrains.jet.util.slicedmap.ReadOnlySlice;
 import org.jetbrains.jet.util.slicedmap.WritableSlice;
+import org.jetbrains.jet.utils.Nulls;
 
 import java.util.Collection;
 import java.util.concurrent.ConcurrentHashMap;
@@ -107,23 +108,6 @@ public class LockBasedStorageManager implements StorageManager {
         return new LockProtectedTrace(lock, originalTrace);
     }
 
-    private static class Nulls {
-        private static final Object NULL_VALUE = new Object();
-    
-        @Nullable
-        @SuppressWarnings("unchecked")
-        private static <V> V unescape(@NotNull Object value) {
-            if (value == NULL_VALUE) return null;
-            return (V) value;
-        }
-    
-        @NotNull
-        private static <V> Object escape(@Nullable V value) {
-            if (value == null) return NULL_VALUE;
-            return value;
-        }
-    }
-
     private static class LockBasedLazyValue<T> implements NullableLazyValue<T> {
         private final Object lock;
         private final Computable<T> computable;

compiler/util/src/org/jetbrains/jet/utils/NotNullMemoizedFunction.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2010-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.utils;
+
+import com.intellij.util.NotNullFunction;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Map;
+
+public abstract class NotNullMemoizedFunction<K, V> extends NullableMemoizedFunction<K, V>  implements NotNullFunction<K, V> {
+
+    public static <K, V> NotNullFunction<K, V> create(@NotNull final NotNullFunction<K, V> compute) {
+        return new NotNullMemoizedFunction<K, V>() {
+            @NotNull
+            @Override
+            protected V compute(@NotNull K input) {
+                return compute.fun(input);
+            }
+        };
+    }
+
+    public NotNullMemoizedFunction(@NotNull Map<K, Object> map) {
+        super(map);
+    }
+
+    public NotNullMemoizedFunction() {
+        super();
+    }
+
+    @NotNull
+    @Override
+    public V fun(@NotNull K input) {
+        //noinspection ConstantConditions
+        return super.fun(input);
+    }
+
+    @NotNull
+    @Override
+    protected abstract V compute(@NotNull K input);
+}

compiler/util/src/org/jetbrains/jet/utils/NullableMemoizedFunction.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2010-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.utils;
+
+import com.intellij.util.Function;
+import com.intellij.util.NullableFunction;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public abstract class NullableMemoizedFunction<K, V> implements NullableFunction<K, V> {
+
+    public static <K, V> NullableFunction<K, V> create(@NotNull final Function<K, V> compute) {
+        return new NullableMemoizedFunction<K, V>() {
+            @Nullable
+            @Override
+            protected V compute(@NotNull K input) {
+                return compute.fun(input);
+            }
+        };
+    }
+
+    private final Map<K, Object> cache;
+
+    public NullableMemoizedFunction(@NotNull Map<K, Object> map) {
+        this.cache = map;
+    }
+
+    public NullableMemoizedFunction() {
+        this(new HashMap<K, Object>());
+    }
+
+    @Override
+    @Nullable
+    public V fun(@NotNull K input) {
+        Object value = cache.get(input);
+        if (value != null) return Nulls.unescape(value);
+
+        V typedValue = compute(input);
+
+        Object oldValue = cache.put(input, Nulls.escape(typedValue));
+        assert oldValue == null : "Race condition detected";
+
+        return typedValue;
+    }
+
+    @Nullable
+    protected abstract V compute(@NotNull K input);
+}

compiler/util/src/org/jetbrains/jet/utils/Nulls.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.utils;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+public class Nulls {
+    private static final Object NULL_VALUE = new Object();
+
+    @Nullable
+    @SuppressWarnings("unchecked")
+    public static <V> V unescape(@NotNull Object value) {
+        if (value == NULL_VALUE) return null;
+        return (V) value;
+    }
+
+    @NotNull
+    public static <V> Object escape(@Nullable V value) {
+        if (value == null) return NULL_VALUE;
+        return value;
+    }
+}