From e869eec377b0000d59ae9692f28c503631358bf4 Mon Sep 17 00:00:00 2001
From: Taras Tsugrii <ttsugrii@fb.com>
Date: Wed, 23 May 2018 11:24:22 -0700
Subject: [PATCH] [Skylark] Implement specialized RangeListView iterators.

---
 .../build/lib/syntax/SkylarkList.java         | 77 +++++++++++++++++++
 .../build/lib/syntax/MethodLibraryTest.java   |  9 ++-
 2 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/google/devtools/build/lib/syntax/SkylarkList.java b/src/main/java/com/google/devtools/build/lib/syntax/SkylarkList.java
index af4096a866a6c9..000c137b5aaf18 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/SkylarkList.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/SkylarkList.java
@@ -18,6 +18,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
+import com.google.common.collect.UnmodifiableIterator;
 import com.google.devtools.build.lib.events.Location;
 import com.google.devtools.build.lib.skylarkinterface.Param;
 import com.google.devtools.build.lib.skylarkinterface.SkylarkCallable;
@@ -28,7 +29,9 @@
 import java.util.AbstractList;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
+import java.util.NoSuchElementException;
 import java.util.RandomAccess;
 import javax.annotation.Nullable;
 
@@ -226,6 +229,65 @@ public static final class RangeList extends SkylarkList<Integer> implements
 
     /** Provides access to range elements based on their index. */
     private static class RangeListView extends AbstractList<Integer> {
+
+      /** Iterator for increasing sequences. */
+      private static class IncreasingIterator extends UnmodifiableIterator<Integer> {
+        private final int stop;
+        private final int step;
+
+        private int cursor;
+
+        private IncreasingIterator(int start, int stop, int step) {
+          this.cursor = start;
+          this.stop = stop;
+          this.step = step;
+        }
+
+        @Override
+        public boolean hasNext() {
+          return cursor < stop;
+        }
+
+        @Override
+        public Integer next() {
+          if (!hasNext()) {
+            throw new NoSuchElementException();
+          }
+          int current = cursor;
+          cursor += step;
+          return current;
+        }
+      }
+
+      /** Iterator for decreasing sequences. */
+      private static class DecreasingIterator extends UnmodifiableIterator<Integer> {
+        private final int stop;
+        private final int step;
+
+        private int cursor;
+
+        private DecreasingIterator(int start, int stop, int step) {
+          this.cursor = start;
+          this.stop = stop;
+          this.step = step;
+        }
+
+        @Override
+        public boolean hasNext() {
+          return cursor > stop;
+        }
+
+        @Override
+        public Integer next() {
+          if (!hasNext()) {
+            throw new NoSuchElementException();
+          }
+          int current = cursor;
+          cursor += step;
+          return current;
+        }
+      }
+
       private static int computeSize(int start, int stop, int step) {
         final int length = Math.abs(stop - start);
         final int absolute_step = Math.abs(step);
@@ -258,6 +320,19 @@ public Integer get(int index) {
       public int size() {
         return size;
       }
+
+      /**
+       * Returns an iterator optimized for traversing range elements, since it's the most frequent
+       * operation for which ranges are used.
+       */
+      @Override
+      public Iterator<Integer> iterator() {
+        if (step > 0) {
+          return new IncreasingIterator(start, stop, step);
+        } else {
+          return new DecreasingIterator(start, stop, step);
+        }
+      }
     }
 
     private final AbstractList<Integer> contents;
@@ -279,6 +354,7 @@ public ImmutableList<Integer> getImmutableList() {
     @Override
     public SkylarkList<Integer> getSlice(Object start, Object end, Object step, Location loc,
         Mutability mutability) throws EvalException {
+      // TODO: use lazy slice implementation
       List<Integer> sliceIndices = EvalUtils.getSliceIndices(start, end, step, this.size(), loc);
       ImmutableList.Builder<Integer> builder = ImmutableList.builderWithExpectedSize(sliceIndices.size());
       for (int pos : sliceIndices) {
@@ -317,6 +393,7 @@ public Tuple<Integer> toTuple() {
     }
 
     public static RangeList of(int start, int stop, int step) {
+      Preconditions.checkArgument(step != 0);
       return new RangeList(start, stop, step);
     }
   }
diff --git a/src/test/java/com/google/devtools/build/lib/syntax/MethodLibraryTest.java b/src/test/java/com/google/devtools/build/lib/syntax/MethodLibraryTest.java
index 1146e9d6cea2c7..cf2f63e997f252 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/MethodLibraryTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/MethodLibraryTest.java
@@ -486,7 +486,14 @@ public void testRange() throws Exception {
         .testStatement("str(range(5, 0, -1))", "[5, 4, 3, 2, 1]")
         .testStatement("str(range(5, 0, -10))", "[5]")
         .testStatement("str(range(0, -3, -2))", "[0, -2]")
-        .testStatement("str(range(3)[1:2])", "[1]")
+        .testStatement("str(range(5)[1:])", "[1, 2, 3, 4]")
+        .testStatement("len(range(5)[1:])", 4)
+        .testStatement("str(range(5)[:2])", "[0, 1]")
+        .testStatement("str(range(10)[1:9:2])", "[1, 3, 5, 7]")
+        .testStatement("str(range(10)[1:10:2])", "[1, 3, 5, 7, 9]")
+        .testStatement("str(range(10)[1:11:2])", "[1, 3, 5, 7, 9]")
+        .testStatement("str(range(0, 10, 2)[::2])", "[0, 4, 8]")
+        .testStatement("str(range(0, 10, 2)[::-2])", "[8, 4, 0]")
         .testIfErrorContains("step cannot be 0", "range(2, 3, 0)");
   }