Optimize Stream.zip and Stream.zip_with with only lists

Author: josevalim

lib/elixir/lib/stream.ex

@@ -1119,7 +1119,27 @@ defmodule Stream do
   """
   @spec zip(Enumerable.t(), Enumerable.t()) :: Enumerable.t()
   def zip(enumerable1, enumerable2) do
-    zip([enumerable1, enumerable2])
+    zip_with(enumerable1, enumerable2, fn left, right -> {left, right} end)
+  end
+
+  @doc """
+  Zips corresponding elements from a finite collection of enumerables
+  into one stream of tuples.
+
+  The zipping finishes as soon as any enumerable in the given collection completes.
+
+  ## Examples
+
+      iex> concat = Stream.concat(1..3, 4..6)
+      iex> cycle = Stream.cycle(["foo", "bar", "baz"])
+      iex> Stream.zip([concat, [:a, :b, :c], cycle]) |> Enum.to_list()
+      [{1, :a, "foo"}, {2, :b, "bar"}, {3, :c, "baz"}]
+
+  """
+  @doc since: "1.4.0"
+  @spec zip(enumerables) :: Enumerable.t() when enumerables: [Enumerable.t()] | Enumerable.t()
+  def zip(enumerables) do
+    zip_with(enumerables, &List.to_tuple(&1))
   end
 
   @doc """
@@ -1139,28 +1159,28 @@ defmodule Stream do
   """
   @doc since: "1.12.0"
   @spec zip_with(Enumerable.t(), Enumerable.t(), (term, term -> term)) :: Enumerable.t()
-  def zip_with(enumerable1, enumerable2, zip_fun) when is_function(zip_fun, 2) do
-    zip_with([enumerable1, enumerable2], &apply(zip_fun, &1))
+  def zip_with(enumerable1, enumerable2, zip_fun)
+      when is_list(enumerable1) and is_list(enumerable2) and is_function(zip_fun, 2) do
+    &zip_pair(enumerable1, enumerable2, &1, &2, zip_fun)
   end
 
-  @doc """
-  Zips corresponding elements from a finite collection of enumerables
-  into one stream of tuples.
+  def zip_with(enumerable1, enumerable2, zip_fun) when is_function(zip_fun, 2) do
+    zip_with([enumerable1, enumerable2], fn [left, right] -> zip_fun.(left, right) end)
+  end
 
-  The zipping finishes as soon as any enumerable in the given collection completes.
+  defp zip_pair(_list1, _list2, {:halt, acc}, _fun, _zip_fun) do
+    {:halted, acc}
+  end
 
-  ## Examples
+  defp zip_pair(list1, list2, {:suspend, acc}, fun, zip_fun) do
+    {:suspended, acc, &zip_pair(list1, list2, &1, fun, zip_fun)}
+  end
 
-      iex> concat = Stream.concat(1..3, 4..6)
-      iex> cycle = Stream.cycle(["foo", "bar", "baz"])
-      iex> Stream.zip([concat, [:a, :b, :c], cycle]) |> Enum.to_list()
-      [{1, :a, "foo"}, {2, :b, "bar"}, {3, :c, "baz"}]
+  defp zip_pair([], _list2, {:cont, acc}, _fun, _zip_fun), do: {:done, acc}
+  defp zip_pair(_list1, [], {:cont, acc}, _fun, _zip_fun), do: {:done, acc}
 
-  """
-  @doc since: "1.4.0"
-  @spec zip(enumerables) :: Enumerable.t() when enumerables: [Enumerable.t()] | Enumerable.t()
-  def zip(enumerables) do
-    zip_with(enumerables, &List.to_tuple(&1))
+  defp zip_pair([head1 | tail1], [head2 | tail2], {:cont, acc}, fun, zip_fun) do
+    zip_pair(tail1, tail2, fun.(zip_fun.(head1, head2), acc), fun, zip_fun)
   end
 
   @doc """
@@ -1188,10 +1208,41 @@ defmodule Stream do
   @spec zip_with(enumerables, (Enumerable.t() -> term)) :: Enumerable.t()
         when enumerables: [Enumerable.t()] | Enumerable.t()
   def zip_with(enumerables, zip_fun) when is_function(zip_fun, 1) do
-    &prepare_zip(enumerables, &1, &2, zip_fun)
+    if is_list(enumerables) and :lists.all(&is_list/1, enumerables) do
+      &zip_list(enumerables, &1, &2, zip_fun)
+    else
+      &zip_enum(enumerables, &1, &2, zip_fun)
+    end
+  end
+
+  defp zip_list(_enumerables, {:halt, acc}, _fun, _zip_fun) do
+    {:halted, acc}
+  end
+
+  defp zip_list(enumerables, {:suspend, acc}, fun, zip_fun) do
+    {:suspended, acc, &zip_list(enumerables, &1, fun, zip_fun)}
+  end
+
+  defp zip_list(enumerables, {:cont, acc}, fun, zip_fun) do
+    case zip_list_heads_tails(enumerables, [], []) do
+      {heads, tails} -> zip_list(tails, fun.(zip_fun.(heads), acc), fun, zip_fun)
+      :error -> {:done, acc}
+    end
+  end
+
+  defp zip_list_heads_tails([[head | tail] | rest], heads, tails) do
+    zip_list_heads_tails(rest, [head | heads], [tail | tails])
+  end
+
+  defp zip_list_heads_tails([[] | _rest], _heads, _tails) do
+    :error
+  end
+
+  defp zip_list_heads_tails([], heads, tails) do
+    {:lists.reverse(heads), :lists.reverse(tails)}
   end
 
-  defp prepare_zip(enumerables, acc, fun, zip_fun) do
+  defp zip_enum(enumerables, acc, fun, zip_fun) do
     step = fn x, acc ->
       {:suspend, :lists.reverse([x | acc])}
     end
@@ -1201,42 +1252,42 @@ defmodule Stream do
         {&Enumerable.reduce(enum, &1, step), [], :cont}
       end)
 
-    do_zip(enum_funs, acc, fun, zip_fun)
+    do_zip_enum(enum_funs, acc, fun, zip_fun)
   end
 
-  # This implementation of do_zip/4 works for any number of streams to zip
-  defp do_zip(zips, {:halt, acc}, _fun, _zip_fun) do
+  # This implementation of do_zip_enum/4 works for any number of streams to zip
+  defp do_zip_enum(zips, {:halt, acc}, _fun, _zip_fun) do
     do_zip_close(zips)
     {:halted, acc}
   end
 
-  defp do_zip(zips, {:suspend, acc}, fun, zip_fun) do
-    {:suspended, acc, &do_zip(zips, &1, fun, zip_fun)}
+  defp do_zip_enum(zips, {:suspend, acc}, fun, zip_fun) do
+    {:suspended, acc, &do_zip_enum(zips, &1, fun, zip_fun)}
   end
 
-  defp do_zip([], {:cont, acc}, _callback, _zip_fun) do
+  defp do_zip_enum([], {:cont, acc}, _callback, _zip_fun) do
     {:done, acc}
   end
 
-  defp do_zip(zips, {:cont, acc}, callback, zip_fun) do
+  defp do_zip_enum(zips, {:cont, acc}, callback, zip_fun) do
     try do
-      do_zip_next_tuple(zips, acc, callback, [], [], zip_fun)
+      do_zip_next(zips, acc, callback, [], [], zip_fun)
     catch
       kind, reason ->
         do_zip_close(zips)
         :erlang.raise(kind, reason, __STACKTRACE__)
     else
       {:next, buffer, acc} ->
-        do_zip(buffer, acc, callback, zip_fun)
+        do_zip_enum(buffer, acc, callback, zip_fun)
 
       {:done, _acc} = other ->
         other
     end
   end
 
-  # do_zip_next_tuple/6 computes the next tuple formed by
+  # do_zip_next/6 computes the next tuple formed by
   # the next element of each zipped stream.
-  defp do_zip_next_tuple(
+  defp do_zip_next(
          [{_, [], :halt} | zips],
          acc,
          _callback,
@@ -1248,15 +1299,15 @@ defmodule Stream do
     {:done, acc}
   end
 
-  defp do_zip_next_tuple([{fun, [], :cont} | zips], acc, callback, yielded_elems, buffer, zip_fun) do
+  defp do_zip_next([{fun, [], :cont} | zips], acc, callback, yielded_elems, buffer, zip_fun) do
     case fun.({:cont, []}) do
       {:suspended, [elem | next_acc], fun} ->
         next_buffer = [{fun, next_acc, :cont} | buffer]
-        do_zip_next_tuple(zips, acc, callback, [elem | yielded_elems], next_buffer, zip_fun)
+        do_zip_next(zips, acc, callback, [elem | yielded_elems], next_buffer, zip_fun)
 
       {_, [elem | next_acc]} ->
         next_buffer = [{fun, next_acc, :halt} | buffer]
-        do_zip_next_tuple(zips, acc, callback, [elem | yielded_elems], next_buffer, zip_fun)
+        do_zip_next(zips, acc, callback, [elem | yielded_elems], next_buffer, zip_fun)
 
       {_, []} ->
         # The current zipped stream terminated, so we close all the streams
@@ -1266,7 +1317,7 @@ defmodule Stream do
     end
   end
 
-  defp do_zip_next_tuple(
+  defp do_zip_next(
          [{fun, zip_acc, zip_op} | zips],
          acc,
          callback,
@@ -1276,10 +1327,10 @@ defmodule Stream do
        ) do
     [elem | rest] = zip_acc
     next_buffer = [{fun, rest, zip_op} | buffer]
-    do_zip_next_tuple(zips, acc, callback, [elem | yielded_elems], next_buffer, zip_fun)
+    do_zip_next(zips, acc, callback, [elem | yielded_elems], next_buffer, zip_fun)
   end
 
-  defp do_zip_next_tuple([] = _zips, acc, callback, yielded_elems, buffer, zip_fun) do
+  defp do_zip_next([] = _zips, acc, callback, yielded_elems, buffer, zip_fun) do
     # "yielded_elems" is a reversed list of results for the current iteration of
     # zipping. That is to say, the nth element from each of the enums being zipped.
     # It needs to be reversed and passed to the zipping function so it can do it's thing.