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> }