Ensure certain streams and enumerables are zippable

Author: José Valim

lib/elixir/lib/hash_dict.ex

@@ -402,7 +402,7 @@ defmodule HashDict do
   end
 
   defp node_reduce(list, -1, { :suspend, acc }, fun, count, next) do
-    { :suspend, acc, &node_reduce(list, -1, &1, fun, count, next) }
+    { :suspended, acc, &node_reduce(list, -1, &1, fun, count, next) }
   end
 
   defp node_reduce([[k|v]|t], -1, { :cont, acc }, fun, _count, next) do

lib/elixir/lib/stream.ex

@@ -235,35 +235,49 @@ defmodule Stream do
 
   """
   @spec flat_map(Enumerable.t, (element -> any)) :: Enumerable.t
-  def flat_map(enum, f) do
-    lazy enum, fn(f1) ->
-      fn(entry, acc) ->
-        enum = f.(entry)
-        fun  = &do_flat_map_each(f1, &1, &2)
-        do_flat_map(&Enumerable.reduce(enum, &1, fun), { :cont, acc })
-      end
+  def flat_map(enum, mapper) do
+    &do_flat_map(enum, mapper, &1, &2)
+  end
+
+  defp do_flat_map(enumerables, mapper, acc, fun) do
+    fun  = &do_flat_map_each(fun, &1, &2)
+    step = &do_flat_map_step/2
+    next = &Enumerable.reduce(enumerables, &1, step)
+    do_flat_map([], next, mapper, acc, fun)
+  end
+
+  defp do_flat_map(next_acc, next, mapper, acc, fun) do
+    case next.({ :cont, next_acc }) do
+      { :suspended, [val|next_acc], next } ->
+        enum = mapper.(val)
+        do_flat_map(next_acc, next, mapper, acc, fun, &Enumerable.reduce(enum, &1, fun))
+      { reason, _ } ->
+        { reason, elem(acc, 1) }
     end
   end
 
-  defp do_flat_map(reduce, acc) do
+  defp do_flat_map(next_acc, next, mapper, acc, fun, reduce) do
     try do
       reduce.(acc)
     catch
-      { :stream_flat_map, acc } -> acc
+      { :stream_flat_map, h } -> { :halted, h }
     else
-      { :done, acc }         -> { :cont, acc }
-      { :halted, acc }       -> { :cont, acc }
-      { :suspended, acc, c } -> { :suspend, acc, &do_flat_map(c, &1) }
+      { _, acc }              -> do_flat_map(next_acc, next, mapper, { :cont, acc }, fun)
+      { :suspended, acc, c }  -> { :suspended, acc, &do_flat_map(next_acc, next, mapper, &1, fun, c) }
     end
   end
 
-  defp do_flat_map_each(f1, x, acc) do
-    case f1.(x, acc) do
-      { :halt, _ } = h -> throw({ :stream_flat_map, h })
-      { _, _ }     = o -> o
+  defp do_flat_map_each(f, x, acc) do
+    case f.(x, acc) do
+      { :halt, h } -> throw({ :stream_flat_map, h })
+      { _, _ } = o -> o
     end
   end
 
+  defp do_flat_map_step(x, acc) do
+    { :suspend, [x|acc] }
+  end
+
   @doc """
   Creates a stream that will reject elements according to
   the given function on enumeration.
@@ -374,7 +388,7 @@ defmodule Stream do
   """
   @spec concat(Enumerable.t) :: Enumerable.t
   def concat(enumerables) do
-    &do_concat(enumerables, &1, &2)
+    flat_map(enumerables, &(&1))
   end
 
   @doc """
@@ -395,34 +409,7 @@ defmodule Stream do
   """
   @spec concat(Enumerable.t, Enumerable.t) :: Enumerable.t
   def concat(first, second) do
-    &do_concat([first, second], &1, &2)
-  end
-
-  defp do_concat(enumerables, acc, fun) do
-    step = &do_concat_step/2
-    next = &Enumerable.reduce(enumerables, &1, step)
-    do_concat([], next, acc, fun)
-  end
-
-  defp do_concat(next_acc, next, acc, fun) do
-    case next.({ :cont, next_acc }) do
-      { :suspended, [enum|next_acc], next } ->
-        do_concat(next_acc, next, acc, fun, &Enumerable.reduce(enum, &1, fun))
-      { reason, _ } ->
-        { reason, elem(acc, 1) }
-    end
-  end
-
-  defp do_concat(next_acc, next, acc, fun, reduce) do
-    case reduce.(acc) do
-      { :done, acc }         -> do_concat(next_acc, next, { :cont, acc }, fun)
-      { :halted, acc } = h   -> h
-      { :suspended, acc, c } -> { :suspended, acc, &do_concat(next_acc, next, &1, fun, c) }
-    end
-  end
-
-  defp do_concat_step(x, acc) do
-    { :suspend, [x|acc] }
+    flat_map([first, second], &(&1))
   end
 
   @doc """

lib/elixir/test/elixir/hash_dict_test.exs

@@ -5,57 +5,57 @@ defmodule HashDictTest do
 
   @dict HashDict.new(foo: :bar)
 
-  test :is_serializable_as_attribute do
+  test "is serializable as attribute" do
     assert @dict == HashDict.new(foo: :bar)
   end
 
-  test :access_the_serialized_as_attribute do
+  test "is accessible as attribute" do
     assert @dict[:foo] == :bar
   end
 
-  test :smoke_small_range_test do
+  test "small dict smoke test" do
     smoke_test(1..8)
     smoke_test(8..1)
   end
 
-  test :smoke_medium_range_test do
+  test "medium dict smoke test" do
     smoke_test(1..80)
     smoke_test(80..1)
   end
 
-  test :smoke_large_range_test do
+  test "large dict smoke test" do
     smoke_test(1..1200)
     smoke_test(1200..1)
   end
 
-  test :fetch! do
+  test "fetch!/2" do
     dict = filled_dict(8)
     assert HashDict.fetch!(dict, 1) == 1
     assert_raise KeyError, fn ->
       HashDict.fetch!(dict, 11)
     end
   end
 
-  test :empty do
+  test "empty/1" do
     assert HashDict.empty filled_dict(8)   == HashDict.new
     assert HashDict.empty filled_dict(20)  == HashDict.new
     assert HashDict.empty filled_dict(120) == HashDict.new
   end
 
-  test :fetch do
+  test "fetch/2" do
     dict = filled_dict(8)
     assert HashDict.fetch(dict, 4)  == { :ok, 4 }
     assert HashDict.fetch(dict, 16) == :error
   end
 
-  test :equal? do
+  test "equal?/2" do
     assert HashDict.equal?(filled_dict(3), filled_dict(3)) == true
 
     assert HashDict.equal?(HashDict.new([{:a, 1}, {:b, 2}]),
                            HashDict.new([{:a, 2}, {:b, 3}])) == false
   end
 
-  test :has_key? do
+  test "has_key?/2" do
     dict = filled_dict(8)
     assert HashDict.has_key? dict, 4
     refute HashDict.has_key? dict, 16
@@ -69,7 +69,7 @@ defmodule HashDictTest do
     refute HashDict.has_key? dict, 240
   end
 
-  test :put_new do
+  test "put_new/3" do
     dict = filled_dict(8)
 
     dict = HashDict.put_new(dict, 1, 11)
@@ -86,7 +86,7 @@ defmodule HashDictTest do
     assert HashDict.size(dict) == 10
   end
 
-  test :update do
+  test "update/3" do
     dict = filled_dict(8)
 
     dict = HashDict.update!(dict, 1, &(&1 * 2))
@@ -119,7 +119,7 @@ defmodule HashDictTest do
     assert HashDict.size(dict) == 10
   end
 
-  test :to_list do
+  test "to_list/1" do
     dict = filled_dict(8)
     list = dict |> HashDict.to_list
     assert length(list) == 8
@@ -139,7 +139,7 @@ defmodule HashDictTest do
     assert list == Enum.to_list(dict)
   end
 
-  test :keys do
+  test "keys/1" do
     list = filled_dict(8) |> HashDict.keys
     assert length(list) == 8
     assert 1 in list
@@ -153,7 +153,7 @@ defmodule HashDictTest do
     assert 1 in list
   end
 
-  test :values do
+  test "values/1" do
     list = filled_dict(8) |> HashDict.values
     assert length(list) == 8
     assert 1 in list
@@ -167,7 +167,7 @@ defmodule HashDictTest do
     assert 1 in list
   end
 
-  test :enum do
+  test "implements Enumerable" do
     dict = filled_dict(10)
     assert Enum.empty?(HashDict.new)
     refute Enum.empty?(dict)
@@ -177,17 +177,17 @@ defmodule HashDictTest do
     assert Enum.map(filled_dict(3), fn({ k, v }) -> k + v end) == [2, 4, 6]
   end
 
-  test :access do
+  test "access" do
     assert filled_dict(8)[1] == 1
     assert filled_dict(8)[5] == 5
     assert filled_dict(8)[9] == nil
   end
 
-  test :inspect do
+  test "inspect" do
     assert inspect(filled_dict(8)) =~ "#HashDict<"
   end
 
-  test :small_range_merge do
+  test "small dict merge" do
     dict1 = filled_dict(8)
     dict2 = Enum.reduce 6..10, HashDict.new, fn(i, d) -> HashDict.put(d, i, i * 2) end
 
@@ -217,7 +217,7 @@ defmodule HashDictTest do
     assert HashDict.size(dict) == 10
   end
 
-  test :medium_range_merge do
+  test "medium dict merge" do
     dict1 = filled_dict(20)
     dict2 = Enum.reduce 18..22, HashDict.new, fn(i, d) -> HashDict.put(d, i, i * 2) end
 
@@ -247,7 +247,7 @@ defmodule HashDictTest do
     assert HashDict.size(dict) == 22
   end
 
-  test :large_range_merge do
+  test "large dict merge" do
     dict1 = filled_dict(120)
     dict2 = Enum.reduce 118..122, HashDict.new, fn(i, d) -> HashDict.put(d, i, i * 2) end
 
@@ -277,12 +277,22 @@ defmodule HashDictTest do
     assert HashDict.size(dict) == 122
   end
 
-  test :trie_contract do
+  test "trie contraction" do
     dict = filled_dict(120)
     dict = Enum.reduce 16..120, dict, fn(x, acc) -> HashDict.delete(acc, x) end
     assert (Enum.filter 1..120, fn(x) -> HashDict.get(dict, x) == x end) == (Enum.sort 1..15)
   end
 
+  test "is zippable" do
+    dict = filled_dict(8)
+    list = Dict.to_list(dict)
+    assert Enum.zip(list, list) == Enum.zip(dict, dict)
+
+    dict = filled_dict(120)
+    list = Dict.to_list(dict)
+    assert Enum.zip(list, list) == Enum.zip(dict, dict)
+  end
+
   defp smoke_test(range) do
     { dict, _ } = Enum.reduce range, { HashDict.new, 1 }, fn(x, { acc, i }) ->
       acc = HashDict.put(acc, x, x)

lib/elixir/test/elixir/hash_set_test.exs

@@ -57,27 +57,27 @@ defmodule HashSetTest do
     refute HashSet.member?(filled_set(120), 200)
   end
 
-  test :subset? do
+  test "subset?" do
     assert HashSet.subset?(HashSet.new, HashSet.new)
     assert HashSet.subset?(filled_set(6), filled_set(10))
     assert HashSet.subset?(filled_set(6), filled_set(120))
 
     refute HashSet.subset?(filled_set(120), filled_set(6))
   end
 
-  test :equal? do
+  test "equal?" do
     assert HashSet.equal?(HashSet.new, HashSet.new)
     assert HashSet.equal?(filled_set(20), HashSet.delete(filled_set(21), 21))
     assert HashSet.equal?(filled_set(120), filled_set(120))
   end
 
-  test :empty do
+  test "empty" do
     assert HashSet.empty filled_set(8)   == HashSet.new
     assert HashSet.empty filled_set(20)  == HashSet.new
     assert HashSet.empty filled_set(120) == HashSet.new
   end
 
-  test :to_list do
+  test "to_list" do
     set = filled_set(8)
     list = set |> HashSet.to_list
     assert length(list) == 8
@@ -97,7 +97,7 @@ defmodule HashSetTest do
     assert list == Enum.to_list(set)
   end
 
-  test :delete do
+  test "delete" do
     assert HashSet.delete(filled_set(8), 8) == filled_set(7)
     assert HashSet.delete(filled_set(8), 9) == filled_set(8)
     assert HashSet.delete(HashSet.new, 10) == HashSet.new
@@ -107,6 +107,16 @@ defmodule HashSetTest do
     assert HashSet.delete(filled_set(121), 121) == filled_set(120)
   end
 
+  test "is zippable" do
+    set  = filled_set(8)
+    list = Dict.to_list(set)
+    assert Enum.zip(list, list) == Enum.zip(set, set)
+
+    set  = filled_set(120)
+    list = Dict.to_list(set)
+    assert Enum.zip(list, list) == Enum.zip(set, set)
+  end
+
   defp filled_set(range) do
     Enum.reduce 1..range, HashSet.new, &HashSet.put(&2, &1)
   end