skylark: zip: support iterables of unknown length (#24)

2 files changed, 30 insertions, 16 deletions

diff --git a/library.go b/library.go
index b8f4b15..5f88c32 100644
--- a/library.go
+++ b/library.go
@@ -1117,23 +1117,35 @@ func zip(thread *Thread, _ *Builtin, args Tuple, kwargs []Tuple) (Value, error)
 		}
 		iters[i] = it
 		n := Len(seq)
-		if n < 0 {
-			// TODO(adonovan): support iterables of unknown length.
-			return nil, fmt.Errorf("zip: argument #%d has unknown length", i+1)
-		}
 		if i == 0 || n < rows {
-			rows = n
+			rows = n // possibly -1
+		}
+	}
+	var result []Value
+	if rows >= 0 {
+		// length known
+		result = make([]Value, rows)
+		array := make(Tuple, cols*rows) // allocate a single backing array
+		for i := 0; i < rows; i++ {
+			tuple := array[:cols:cols]
+			array = array[cols:]
+			for j, iter := range iters {
+				iter.Next(&tuple[j])
+			}
+			result[i] = tuple
 		}
-	}
-	result := make([]Value, rows)
-	array := make(Tuple, cols*rows) // allocate a single backing array
-	for i := 0; i < rows; i++ {
-		tuple := array[:cols:cols]
-		array = array[cols:]
-		for j, iter := range iters {
-			iter.Next(&tuple[j])
+	} else {
+		// length not known
+	outer:
+		for {
+			tuple := make(Tuple, cols)
+			for i, iter := range iters {
+				if !iter.Next(&tuple[i]) {
+					break outer
+				}
+			}
+			result = append(result, tuple)
 		}
-		result[i] = tuple
 	}
 	return NewList(result), nil
 }
diff --git a/testdata/builtins.sky b/testdata/builtins.sky
index d6774a8..d5f8f04 100644
--- a/testdata/builtins.sky
+++ b/testdata/builtins.sky
@@ -129,10 +129,12 @@ assert.eq(enumerate([False, True, None], 42), [(42, False), (43, True), (44, Non
 
 # zip
 assert.eq(zip(), [])
+assert.eq(zip([]), [])
 assert.eq(zip([1, 2, 3]), [(1,), (2,), (3,)])
-assert.eq(zip(list("abc".split_bytes()),
+assert.eq(zip("".split_bytes()), [])
+assert.eq(zip("abc".split_bytes(),
               list("def".split_bytes()),
-              list("hijk".split_bytes())),
+              "hijk".split_bytes()),
           [("a", "d", "h"), ("b", "e", "i"), ("c", "f", "j")])
 z1 = [1]
 assert.eq(zip(z1), [(1,)])