diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/LockBasedStorageManager.java b/compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/LockBasedStorageManager.java index 4b83605d7da..c9f61bd8188 100644 --- a/compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/LockBasedStorageManager.java +++ b/compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/LockBasedStorageManager.java @@ -44,38 +44,49 @@ public class LockBasedStorageManager implements StorageManager { @NotNull @Override - public Function createMemoizedFunction(@NotNull final Function compute, @NotNull final MemoizationMode mode) { - return createMemoizedFunction(compute, mode, false); + public Function createMemoizedFunction(@NotNull final Function compute, @NotNull final MemoizationMode modeForValues) { + return new Function() { + private final ConcurrentMap cache = createConcurrentMap(modeForValues); + + @Override + public V fun(@NotNull final K input) { + V value = cache.get(input); + if (value != null) return value; + + synchronized (lock) { + value = cache.get(input); + if (value != null) return value; + + value = compute.fun(input); + + V oldValue = cache.put(input, value); + assert oldValue == null : "Race condition detected"; + } + + return value; + } + }; } @NotNull @Override public Function createMemoizedFunctionWithNullableValues( - @NotNull Function compute, @NotNull MemoizationMode modeForValues + @NotNull final Function compute, @NotNull final MemoizationMode modeForValues ) { - return createMemoizedFunction(compute, modeForValues, true); - } - - @NotNull - private Function createMemoizedFunction(final Function compute, final MemoizationMode mode, final boolean nullable) { return new Function() { - private final ConcurrentMap> cache; - { - cache = (mode == WEAK) ? new ConcurrentWeakValueHashMap>() : new ConcurrentHashMap>(); - } + private final ConcurrentMap> cache = createConcurrentMap(modeForValues); @Override public V fun(@NotNull final K input) { LazyValue lazyValue = cache.get(input); if (lazyValue != null) return lazyValue.get(); - Computable computable = new Computable() { + lazyValue = createNullableLazyValue(new Computable() { @Override public V compute() { return compute.fun(input); } - }; - lazyValue = nullable ? createNullableLazyValue(computable) : createLazyValue(computable); + }); LazyValue oldValue = cache.putIfAbsent(input, lazyValue); if (oldValue != null) return oldValue.get(); @@ -85,6 +96,10 @@ public class LockBasedStorageManager implements StorageManager { }; } + private static ConcurrentMap createConcurrentMap(MemoizationMode mode) { + return (mode == WEAK) ? new ConcurrentWeakValueHashMap() : new ConcurrentHashMap(); + } + @NotNull @Override public LazyValue createLazyValue(@NotNull Computable computable) {