From f0eaaa4a5da2bb6d9c702b034fe7b702c1d3e3ec Mon Sep 17 00:00:00 2001
From: DeedleFake <DeedleFake@users.noreply.github.com>
Date: Thu, 27 Jun 2019 14:46:36 -0400
Subject: [PATCH] std/stream: Add `zip`. (#188)

---
 std/stream/middle.go | 56 +++++++++++++++++++++++++++++++++++++++++++-
 std/stream/stream.go |  1 +
 wdte_test.go         |  9 +++++++
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/std/stream/middle.go b/std/stream/middle.go
index 5fd360d..2858250 100644
--- a/std/stream/middle.go
+++ b/std/stream/middle.go
@@ -1,6 +1,9 @@
 package stream
 
-import "github.com/DeedleFake/wdte"
+import (
+	"github.com/DeedleFake/wdte"
+	"github.com/DeedleFake/wdte/wdteutil"
+)
 
 // Map is a WDTE function with the following signature:
 //
@@ -246,3 +249,54 @@ func Limit(frame wdte.Frame, args ...wdte.Func) wdte.Func {
 		})
 	})
 }
+
+// Zip is a WDTE function with the following signatures:
+//
+//    (zip s1) ...
+//    zip ...
+//
+// Zip returns a Stream which yields the streams that it is given
+// simultaneuously as arrays. In other words,
+//
+//    zip (a.stream [1; 2; 3]) (a.stream ['a'; 'b'; 'c'])
+//
+// will yield
+//
+//    [1; 'a']
+//    [2; 'b']
+//    [3; 'c']
+//
+// The order of the yielded arrays matches the order that the streams
+// are given in. If one of the streams ends before the other ones, End
+// will be yielded for that stream after that point.
+func Zip(frame wdte.Frame, args ...wdte.Func) wdte.Func {
+	frame = frame.Sub("zip")
+
+	if len(args) < 2 {
+		return wdteutil.SaveArgs(wdte.GoFunc(Zip), args...)
+	}
+
+	streams := make([]Stream, 0, len(args))
+	for _, arg := range args {
+		streams = append(streams, arg.Call(frame).(Stream))
+	}
+
+	return NextFunc(func(frame wdte.Frame) (wdte.Func, bool) {
+		frame = frame.Sub("zip")
+
+		r := make(wdte.Array, 0, len(streams))
+		var more bool
+		for _, stream := range streams {
+			n, ok := stream.Next(frame)
+			if !ok {
+				r = append(r, end{})
+				continue
+			}
+
+			r = append(r, n)
+			more = true
+		}
+
+		return r, more
+	})
+}
diff --git a/std/stream/stream.go b/std/stream/stream.go
index d4504cb..6683730 100644
--- a/std/stream/stream.go
+++ b/std/stream/stream.go
@@ -49,6 +49,7 @@ var Scope = wdte.S().Map(map[wdte.ID]wdte.Func{
 	"enumerate": wdte.GoFunc(Enumerate),
 	"repeat":    wdte.GoFunc(Repeat),
 	"limit":     wdte.GoFunc(Limit),
+	"zip":       wdte.GoFunc(Zip),
 
 	"end":     End(),
 	"collect": wdte.GoFunc(Collect),
diff --git a/wdte_test.go b/wdte_test.go
index cadb244..6a53b2e 100644
--- a/wdte_test.go
+++ b/wdte_test.go
@@ -16,6 +16,7 @@ import (
 	wdteio "github.com/DeedleFake/wdte/std/io"
 	_ "github.com/DeedleFake/wdte/std/math"
 	_ "github.com/DeedleFake/wdte/std/rand"
+	"github.com/DeedleFake/wdte/std/stream"
 	_ "github.com/DeedleFake/wdte/std/stream"
 	_ "github.com/DeedleFake/wdte/std/strings"
 )
@@ -567,6 +568,14 @@ func TestStream(t *testing.T) {
 				wdte.Number(0), wdte.Number(1), wdte.Number(2),
 			},
 		},
+		{
+			name:   "Zip",
+			script: `let s => import 'stream'; s.zip (s.range 2) (s.range 1 2) -> s.collect;`,
+			ret: wdte.Array{
+				wdte.Array{wdte.Number(0), wdte.Number(1)},
+				wdte.Array{wdte.Number(1), stream.End()},
+			},
+		},
 		{
 			name: "Drain",
 			script: `