Support memoizing lazy iterable

compiler/frontend/src/org/jetbrains/jet/storage/LockBasedLazyResolveStorageManager.java

@@ -32,6 +32,7 @@ import org.jetbrains.jet.util.slicedmap.ReadOnlySlice;
 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 @@ public class LockBasedLazyResolveStorageManager implements LazyResolveStorageMan
         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;

compiler/tests/org/jetbrains/jet/storage/StorageManagerTest.java

@@ -6,10 +6,7 @@ import jet.Unit;
 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 @@ public class StorageManagerTest extends TestCase {
         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) {

core/util.runtime/src/org/jetbrains/jet/storage/LockBasedStorageManager.java

@@ -24,6 +24,7 @@ import org.jetbrains.annotations.Nullable;
 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;
@@ -168,6 +169,12 @@ public class LockBasedStorageManager implements StorageManager {
         }
     }
 
+    @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");

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

core/util.runtime/src/org/jetbrains/jet/storage/StorageManager.kt

@@ -54,4 +54,7 @@ public trait StorageManager {
     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>
 }