Merge pull request #4415 from benwilson512/mapset-optimizations

MapSet Optimizations

lib/elixir/lib/map.ex

@@ -76,8 +76,14 @@ defmodule Map do
 
   """
   @spec new(Enum.t) :: map
-  def new(enumerable) do
-    Enum.reduce(enumerable, %{}, fn {k, v}, acc -> put(acc, k, v) end)
+  def new(%{__struct__: _} = struct), do: new_from_enum(struct)
+  def new(%{} = map), do: map
+  def new(enum), do: new_from_enum(enum)
+
+  defp new_from_enum(enumerable) do
+    enumerable
+    |> Enum.to_list
+    |> :maps.from_list
   end
 
   @doc """
@@ -93,11 +99,18 @@ defmodule Map do
   """
   @spec new(Enum.t, (term -> {key, value})) :: map
   def new(enumerable, transform) do
-    fun = fn el, acc ->
-      {k, v} = transform.(el)
-      put(acc, k, v)
-    end
-    Enum.reduce(enumerable, %{}, fun)
+    enumerable
+    |> Enum.to_list
+    |> do_new_transform(transform, [])
+  end
+
+  defp do_new_transform([], _fun, acc) do
+    acc
+    |> :lists.reverse
+    |> :maps.from_list
+  end
+  defp do_new_transform([item | rest], fun, acc) do
+    do_new_transform(rest, fun, [fun.(item) | acc])
   end
 
   @doc """
@@ -209,15 +222,20 @@ defmodule Map do
       %{a: 1, c: 3}
 
   """
-  @spec take(map, [key]) :: map
+  @spec take(map, Enumerable.t) :: map
   def take(map, keys) do
-    Enum.reduce(keys, [], fn key, acc ->
-      case fetch(map, key) do
-        {:ok, value} -> [{key, value} | acc]
-        :error -> acc
-      end
-    end)
-    |> :maps.from_list
+    keys
+    |> Enum.to_list
+    |> do_take(map, [])
+  end
+
+  defp do_take([], _map, acc), do: :maps.from_list(acc)
+  defp do_take([key | rest], map, acc) do
+    acc = case fetch(map, key) do
+      {:ok, value} -> [{key, value} | acc]
+      :error -> acc
+    end
+    do_take(rest, map, acc)
   end
 
   @doc """
@@ -423,9 +441,16 @@ defmodule Map do
       %{a: 1, c: 3}
 
   """
-  @spec drop(map, [key]) :: map
+  @spec drop(map, Enumerable.t) :: map
   def drop(map, keys) do
-    Enum.reduce(keys, map, &delete(&2, &1))
+    keys
+    |> Enum.to_list
+    |> drop_list(map)
+  end
+
+  defp drop_list([], acc), do: acc
+  defp drop_list([key | rest], acc) do
+    drop_list(rest, Map.delete(acc, key))
   end
 
   @doc """
@@ -442,16 +467,23 @@ defmodule Map do
       {%{a: 1, c: 3}, %{b: 2}}
 
   """
-  @spec split(map, [key]) :: {map, map}
+  @spec split(map, Enumerable.t) :: {map, map}
   def split(map, keys) do
-    Enum.reduce(keys, {new, map}, fn key, {inc, exc} = acc ->
-      case fetch(exc, key) do
-        {:ok, value} ->
-          {put(inc, key, value), delete(exc, key)}
-        :error ->
-          acc
-      end
-    end)
+    keys
+    |> Enum.to_list
+    |> do_split([], map)
+  end
+
+  defp do_split([], inc, exc) do
+    {:maps.from_list(inc), exc}
+  end
+  defp do_split([key | rest], inc, exc) do
+    case fetch(exc, key) do
+      {:ok, value} ->
+        do_split(rest, [{key, value} | inc], delete(exc, key))
+      :error ->
+        do_split(rest, inc, exc)
+    end
   end
 
   @doc """

lib/elixir/lib/map_set.ex

@@ -37,8 +37,14 @@ defmodule MapSet do
 
   """
   @spec new(Enum.t) :: t
+  def new(%__MODULE__{} = mapset), do: mapset
   def new(enumerable) do
-    Enum.reduce(enumerable, %MapSet{}, &put(&2, &1))
+    map =
+      enumerable
+      |> Enum.to_list
+      |> do_new([])
+
+    %MapSet{map: map}
   end
 
   @doc """
@@ -52,7 +58,30 @@ defmodule MapSet do
   """
   @spec new(Enum.t, (term -> term)) :: t
   def new(enumerable, transform) do
-    Enum.reduce(enumerable, %MapSet{}, &put(&2, transform.(&1)))
+    map =
+      enumerable
+      |> Enum.to_list
+      |> do_new_transform(transform, [])
+
+    %MapSet{map: map}
+  end
+
+  defp do_new([], acc) do
+    acc
+    |> :lists.reverse
+    |> :maps.from_list
+  end
+  defp do_new([item | rest], acc) do
+    do_new(rest, [{item, true} | acc])
+  end
+
+  defp do_new_transform([], _fun, acc) do
+    acc
+    |> :lists.reverse
+    |> :maps.from_list
+  end
+  defp do_new_transform([item | rest], fun, acc) do
+    do_new_transform(rest, fun, [{fun.(item), true} | acc])
   end
 
   @doc """
@@ -84,13 +113,36 @@ defmodule MapSet do
 
   """
   @spec difference(t, t) :: t
-  def difference(%MapSet{map: map1}, %MapSet{map: map2}) do
-    map = :maps.fold(fn value, _, acc ->
-      Map.delete(acc, value)
-    end, map1, map2)
+  # If the first set is less than twice the size of the second map,
+  # it is fastest to re-accumulate items in the first set that are not
+  # present in the second set.
+  def difference(%MapSet{map: map1}, %MapSet{map: map2})
+  when map_size(map1) < map_size(map2) * 2 do
+    map = map1
+    |> Map.keys
+    |> filter_not_in(map2)
+
     %MapSet{map: map}
   end
 
+  # If the second set is less than half the size of the first set, it's fastest
+  # to simply iterate through each item in the second set, deleting them from
+  # the first set.
+  def difference(%MapSet{map: map1}, %MapSet{map: map2}) do
+    %MapSet{map: Map.drop(map1, Map.keys(map2))}
+  end
+
+  defp filter_not_in(keys, map2, acc \\ [])
+  defp filter_not_in([], _map2, acc), do: :maps.from_list(acc)
+  defp filter_not_in([key | rest], map2, acc) do
+    acc = if Map.has_key?(map2, key) do
+      acc
+    else
+      [{key, true} | acc]
+    end
+    filter_not_in(rest, map2, acc)
+  end
+
   @doc """
   Checks if `set1` and `set2` have no members in common.
 
@@ -105,15 +157,20 @@ defmodule MapSet do
   @spec disjoint?(t, t) :: boolean
   def disjoint?(%MapSet{map: map1}, %MapSet{map: map2}) do
     {map1, map2} = order_by_size(map1, map2)
-    :maps.fold(fn value, _, _ ->
-      if Map.has_key?(map2, value) do
-        throw({:halt, false})
-      else
-        true
-      end
-    end, true, map1)
-  catch
-    {:halt, false} -> false
+
+    map1
+    |> Map.keys
+    |> none_in?(map2)
+  end
+
+  defp none_in?([], _) do
+    true
+  end
+  defp none_in?([key | rest], map2) do
+    case Map.has_key?(map2, key) do
+      true -> false
+      false -> none_in?(rest, map2)
+    end
   end
 
   @doc """
@@ -149,14 +206,8 @@ defmodule MapSet do
   @spec intersection(t, t) :: t
   def intersection(%MapSet{map: map1}, %MapSet{map: map2}) do
     {map1, map2} = order_by_size(map1, map2)
-    map = :maps.fold(fn value, _, acc ->
-      if Map.has_key?(map2, value) do
-        Map.put(acc, value, true)
-      else
-        acc
-      end
-    end, %{}, map1)
-    %MapSet{map: map}
+
+    %MapSet{map: Map.take(map2, Map.keys(map1))}
   end
 
   @doc """
@@ -221,18 +272,21 @@ defmodule MapSet do
   @spec subset?(t, t) :: boolean
   def subset?(%MapSet{map: map1}, %MapSet{map: map2}) do
     if map_size(map1) <= map_size(map2) do
-      :maps.fold(fn value, _, _ ->
-        if Map.has_key?(map2, value) do
-          true
-        else
-          throw({:halt, false})
-        end
-      end, true, map1)
+      map1
+      |> Map.keys
+      |> do_subset?(map2)
+    else
+      false
+    end
+  end
+
+  defp do_subset?([], _), do: true
+  defp do_subset?([key | rest], map2) do
+    if Map.has_key?(map2, key) do
+      do_subset?(rest, map2)
     else
       false
     end
-  catch
-    {:halt, false} -> false
   end
 
   @doc """

lib/elixir/test/elixir/map_test.exs

@@ -38,6 +38,21 @@ defmodule MapTest do
     assert map_size(@sample) == 2
   end
 
+  test "take/2" do
+    assert Map.take(%{a: 1, b: 2, c: 3}, [:b, :c]) == %{b: 2, c: 3}
+    assert Map.take(%{a: 1, b: 2, c: 3}, MapSet.new([:b, :c])) == %{b: 2, c: 3}
+  end
+
+  test "drop/2" do
+    assert Map.drop(%{a: 1, b: 2, c: 3}, [:b, :c]) == %{a: 1}
+    assert Map.drop(%{a: 1, b: 2, c: 3}, MapSet.new([:b, :c])) == %{a: 1}
+  end
+
+  test "split/2" do
+    assert Map.split(%{a: 1, b: 2, c: 3}, [:b, :c]) == {%{b: 2, c: 3}, %{a: 1}}
+    assert Map.split(%{a: 1, b: 2, c: 3}, MapSet.new([:b, :c])) == {%{b: 2, c: 3}, %{a: 1}}
+  end
+
   test "maps with optional comma" do
     assert %{a: :b,} == %{a: :b}
     assert %{1 => 2,} == %{1 => 2}