Support memoizing lazy iterable
diff --git a/compiler/frontend/src/org/jetbrains/jet/storage/LockBasedLazyResolveStorageManager.java b/compiler/frontend/src/org/jetbrains/jet/storage/LockBasedLazyResolveStorageManager.java
index 91084ed..c1efbdd 100644
--- a/compiler/frontend/src/org/jetbrains/jet/storage/LockBasedLazyResolveStorageManager.java
+++ b/compiler/frontend/src/org/jetbrains/jet/storage/LockBasedLazyResolveStorageManager.java
@@ -32,6 +32,7 @@
import org.jetbrains.jet.util.slicedmap.WritableSlice;
import java.util.Collection;
+import java.util.Iterator;
import java.util.concurrent.locks.Lock;
// This class is kept under the same package as LockBasedStorageManager to get access to its protected members
@@ -134,6 +135,12 @@
return storageManager.compute(computable);
}
+ @NotNull
+ @Override
+ public <T, D> Iterable<T> createLazyIterable(Iterator<? extends D> data, Function1<? super D, ? extends T> compute) {
+ return storageManager.createLazyIterable(data, compute);
+ }
+
private static class LockProtectedContext implements BindingContext {
private final Lock lock;
private final BindingContext context;
diff --git a/compiler/tests/org/jetbrains/jet/storage/StorageManagerTest.java b/compiler/tests/org/jetbrains/jet/storage/StorageManagerTest.java
index db4c727..e8de4de 100644
--- a/compiler/tests/org/jetbrains/jet/storage/StorageManagerTest.java
+++ b/compiler/tests/org/jetbrains/jet/storage/StorageManagerTest.java
@@ -6,10 +6,7 @@
import junit.framework.TestCase;
import org.jetbrains.jet.util.ReenteringLazyValueComputationException;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.List;
+import java.util.*;
public class StorageManagerTest extends TestCase {
@@ -426,6 +423,64 @@
assertEquals(2, c.getCount());
}
+ // Lazy iterable
+
+ public void testEmpty() throws Exception {
+ Iterable<Object> iterable = m.createLazyIterable(
+ Collections.emptyList().iterator(),
+ new Function1<Object, Object>() {
+ @Override
+ public Object invoke(Object o) {
+ fail("Should not be called for empty data. Argument: " + o);
+ return null;
+ }
+ }
+ );
+ for (Object o : iterable) {
+ fail("Should be empty. Found: " + o);
+ }
+ }
+
+ public void testOne() throws Exception {
+ doTestMutableIterable("a");
+ }
+
+ public void testTwo() throws Exception {
+ doTestMutableIterable("a", "b");
+ }
+
+ public void testThree() throws Exception {
+ doTestMutableIterable("a", "b", "c");
+ }
+
+ private void doTestMutableIterable(String... items) {
+ List<String> data = Arrays.asList(items);
+ final CounterImpl c = new CounterImpl();
+ Iterable<String> iterable = m.createLazyIterable(
+ data.iterator(),
+ new Function1<String, String>() {
+ @Override
+ public String invoke(String s) {
+ c.inc();
+ return s;
+ }
+ }
+ );
+
+ for (int i = 0; i < data.size(); i++) {
+ // iterate up to i
+ int j = 0;
+ for (String s : iterable) {
+ assertEquals("Iterating up to " + i + " at index " + j, data.get(j), s);
+ if (j == i) break;
+ j++;
+ }
+ assertEquals("Too few iterations", i, j);
+ }
+
+ assertEquals(data.size(), c.getCount());
+ }
+
// Utilities
private static <K, V> Function0<V> apply(final Function1<K, V> f, final K x) {
diff --git a/core/util.runtime/src/org/jetbrains/jet/storage/LockBasedStorageManager.java b/core/util.runtime/src/org/jetbrains/jet/storage/LockBasedStorageManager.java
index e1b43ca..ed97614 100644
--- a/core/util.runtime/src/org/jetbrains/jet/storage/LockBasedStorageManager.java
+++ b/core/util.runtime/src/org/jetbrains/jet/storage/LockBasedStorageManager.java
@@ -24,6 +24,7 @@
import org.jetbrains.jet.utils.ExceptionUtils;
import org.jetbrains.jet.utils.WrappedValues;
+import java.util.Iterator;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.locks.Lock;
@@ -169,6 +170,12 @@
}
@NotNull
+ @Override
+ public <T, D> Iterable<T> createLazyIterable(Iterator<? extends D> data, Function1<? super D, ? extends T> compute) {
+ return new MemoizingLazyIterable<T, D>(lock, data, compute);
+ }
+
+ @NotNull
protected <T> RecursionDetectedResult<T> recursionDetectedDefault() {
throw new IllegalStateException("Recursive call in a lazy value");
}
diff --git a/core/util.runtime/src/org/jetbrains/jet/storage/MemoizingLazyIterable.kt b/core/util.runtime/src/org/jetbrains/jet/storage/MemoizingLazyIterable.kt
new file mode 100644
index 0000000..50dd80a
--- /dev/null
+++ b/core/util.runtime/src/org/jetbrains/jet/storage/MemoizingLazyIterable.kt
@@ -0,0 +1,80 @@
+/*
+ * Copyright 2010-2014 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.storage
+
+import java.util.concurrent.locks.Lock
+import java.util.ArrayList
+import java.util.NoSuchElementException
+import org.jetbrains.kotlin.util.printAndReturn
+
+class MemoizingLazyIterable<T, D>(
+ private val lock: Lock,
+ // all access to this iterator is guarded by the lock
+ private val data: Iterator<D>,
+ // all calls to this function are guarded by the lock
+ private val compute : (D) -> T
+) : Iterable<T> {
+
+ // this list contains results of applying compute() to items in data
+ // data is traversed once, the size of this list reflects how many data items were seen
+ private val storedValues = ArrayList<T>()
+
+ private fun isElementAvailableAt(index: Int): Boolean {
+ lock.lock()
+ try {
+ val size = storedValues.size()
+ assert(index <= size) {"Trying to look too far ahead: $index, only ${size} elements are computed"}
+ return index < size || data.hasNext()
+ }
+ finally {
+ lock.unlock()
+ }
+ }
+
+ private fun getByIndex(index: Int): T {
+ lock.lock()
+ try {
+ if (storedValues.size() > index) return storedValues[index]
+
+ for (i in storedValues.size()..index) {
+ if (!data.hasNext()) throw NoSuchElementException("Index: $index")
+
+ storedValues.add(compute(data.next()))
+ }
+
+ if (!data.hasNext()) storedValues.trimToSize()
+
+ return storedValues[index]
+ }
+ finally {
+ lock.unlock()
+ }
+ }
+
+ // Accessing the same iterator from different threads is NOT safe
+ // Having many iterators, each of which is accessed from only one thread, is safe
+ override fun iterator(): Iterator<T> {
+ return object : Iterator<T> {
+ private var index = 0
+
+ override fun next() = getByIndex(index++)
+
+ override fun hasNext() = isElementAvailableAt(index)
+ }
+ }
+
+}
\ No newline at end of file
diff --git a/core/util.runtime/src/org/jetbrains/jet/storage/StorageManager.kt b/core/util.runtime/src/org/jetbrains/jet/storage/StorageManager.kt
index 9a93217..db9bd23 100644
--- a/core/util.runtime/src/org/jetbrains/jet/storage/StorageManager.kt
+++ b/core/util.runtime/src/org/jetbrains/jet/storage/StorageManager.kt
@@ -54,4 +54,7 @@
fun createNullableLazyValueWithPostCompute<T: Any>(computable: () -> T?, postCompute: (T?) -> Unit): NullableLazyValue<T>
fun compute<T>(computable: () -> T): T
+
+ // data and computable are accessed only synchronously
+ fun <T, D> createLazyIterable(data: Iterator<D>, compute: (D) -> T): Iterable<T>
}
