Merge pull request #1281 from bettio/add-more-enum-functions-1

Add more Enum functions

Add some missing Enum functions and make find functions compatible with
Enumerable protocol.

These changes are made under both the "Apache 2.0" and the "GNU Lesser General
Public License 2.1 or later" license terms (dual license).

SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later

CHANGELOG.md

@@ -15,14 +15,15 @@ Elixir standard library modules
 - Support for Elixir `String.Chars` protocol, now functions such as `Enum.join` are able to take
 also non string parameters (e.g. `Enum.join([1, 2], ",")`
 - Support for Elixir `Enum.at/3`
-- Add support to Elixir `Enumerable` protocol also for `Enum.all?`, `Enum.any?`, `Enum.each` and
-`Enum.filter`
 - Add support for `is_bitstring/1` construct which is used in Elixir protocols runtime.
+- Add support to Elixir `Enumerable` protocol also for `Enum.all?`, `Enum.any?`, `Enum.each`,
+`Enum.filter`, `Enum.flat_map`, `Enum.reject`, `Enum.chunk_by` and `Enum.chunk_while`
 
 ### Changed
 
 - ESP32: Elixir library is not shipped anymore with `esp32boot.avm`. Use `elixir_esp32boot.avm`
 instead
+- `Enum.find_index` and `Enum.find_value` support Enumerable and not just lists
 
 ### Fixed
 

libs/exavmlib/lib/Enum.ex

@@ -24,6 +24,7 @@ defmodule Enum do
   @compile {:autoload, false}
 
   @type t :: Enumerable.t()
+  @type acc :: any
   @type index :: integer
   @type element :: any
 
@@ -209,6 +210,98 @@ defmodule Enum do
     end
   end
 
+  @doc """
+  Chunks the `enumerable` with fine grained control when every chunk is emitted.
+
+  `chunk_fun` receives the current element and the accumulator and
+  must return `{:cont, chunk, acc}` to emit the given chunk and
+  continue with accumulator or `{:cont, acc}` to not emit any chunk
+  and continue with the return accumulator.
+
+  `after_fun` is invoked when iteration is done and must also return
+  `{:cont, chunk, acc}` or `{:cont, acc}`.
+
+  Returns a list of lists.
+
+  ## Examples
+
+      iex> chunk_fun = fn element, acc ->
+      ...>   if rem(element, 2) == 0 do
+      ...>     {:cont, Enum.reverse([element | acc]), []}
+      ...>   else
+      ...>     {:cont, [element | acc]}
+      ...>   end
+      ...> end
+      iex> after_fun = fn
+      ...>   [] -> {:cont, []}
+      ...>   acc -> {:cont, Enum.reverse(acc), []}
+      ...> end
+      iex> Enum.chunk_while(1..10, [], chunk_fun, after_fun)
+      [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]
+
+  """
+  @doc since: "1.5.0"
+  @spec chunk_while(
+          t,
+          acc,
+          (element, acc -> {:cont, chunk, acc} | {:cont, acc} | {:halt, acc}),
+          (acc -> {:cont, chunk, acc} | {:cont, acc})
+        ) :: Enumerable.t()
+        when chunk: any
+  def chunk_while(enumerable, acc, chunk_fun, after_fun) do
+    {_, {res, acc}} =
+      Enumerable.reduce(enumerable, {:cont, {[], acc}}, fn entry, {buffer, acc} ->
+        case chunk_fun.(entry, acc) do
+          {:cont, emit, acc} -> {:cont, {[emit | buffer], acc}}
+          {:cont, acc} -> {:cont, {buffer, acc}}
+          {:halt, acc} -> {:halt, {buffer, acc}}
+        end
+      end)
+
+    case after_fun.(acc) do
+      {:cont, _acc} -> :lists.reverse(res)
+      {:cont, elem, _acc} -> :lists.reverse([elem | res])
+    end
+  end
+
+  @doc """
+  Splits enumerable on every element for which `fun` returns a new
+  value.
+
+  Returns a list of lists.
+
+  ## Examples
+
+      iex> Enum.chunk_by([1, 2, 2, 3, 4, 4, 6, 7, 7], &(rem(&1, 2) == 1))
+      [[1], [2, 2], [3], [4, 4, 6], [7, 7]]
+
+  """
+  @spec chunk_by(t, (element -> any)) :: [list]
+  def chunk_by(enumerable, fun) do
+    reducers_chunk_by(&chunk_while/4, enumerable, fun)
+  end
+
+  # Taken from Stream.Reducers
+  defp reducers_chunk_by(chunk_by, enumerable, fun) do
+    chunk_fun = fn
+      entry, nil ->
+        {:cont, {[entry], fun.(entry)}}
+
+      entry, {acc, value} ->
+        case fun.(entry) do
+          ^value -> {:cont, {[entry | acc], value}}
+          new_value -> {:cont, :lists.reverse(acc), {[entry], new_value}}
+        end
+    end
+
+    after_fun = fn
+      nil -> {:cont, :done}
+      {acc, _value} -> {:cont, :lists.reverse(acc), :done}
+    end
+
+    chunk_by.(enumerable, nil, chunk_fun, after_fun)
+  end
+
   @doc """
   Invokes the given `fun` for each element in the `enumerable`.
 
@@ -305,14 +398,101 @@ defmodule Enum do
     |> elem(1)
   end
 
+  @doc """
+  Similar to `find/3`, but returns the index (zero-based)
+  of the element instead of the element itself.
+
+  ## Examples
+
+      iex> Enum.find_index([2, 4, 6], fn x -> rem(x, 2) == 1 end)
+      nil
+
+      iex> Enum.find_index([2, 3, 4], fn x -> rem(x, 2) == 1 end)
+      1
+
+  """
+  @spec find_index(t, (element -> any)) :: non_neg_integer | nil
   def find_index(enumerable, fun) when is_list(enumerable) do
     find_index_list(enumerable, 0, fun)
   end
 
+  def find_index(enumerable, fun) do
+    result =
+      Enumerable.reduce(enumerable, {:cont, {:not_found, 0}}, fn entry, {_, index} ->
+        if fun.(entry), do: {:halt, {:found, index}}, else: {:cont, {:not_found, index + 1}}
+      end)
+
+    case elem(result, 1) do
+      {:found, index} -> index
+      {:not_found, _} -> nil
+    end
+  end
+
+  @doc """
+  Similar to `find/3`, but returns the value of the function
+  invocation instead of the element itself.
+
+  ## Examples
+
+      iex> Enum.find_value([2, 4, 6], fn x -> rem(x, 2) == 1 end)
+      nil
+
+      iex> Enum.find_value([2, 3, 4], fn x -> rem(x, 2) == 1 end)
+      true
+
+      iex> Enum.find_value([1, 2, 3], "no bools!", &is_boolean/1)
+      "no bools!"
+
+  """
+  @spec find_value(t, any, (element -> any)) :: any | nil
+  def find_value(enumerable, default \\ nil, fun)
+
   def find_value(enumerable, default, fun) when is_list(enumerable) do
     find_value_list(enumerable, default, fun)
   end
 
+  def find_value(enumerable, default, fun) do
+    Enumerable.reduce(enumerable, {:cont, default}, fn entry, default ->
+      fun_entry = fun.(entry)
+      if fun_entry, do: {:halt, fun_entry}, else: {:cont, default}
+    end)
+    |> elem(1)
+  end
+
+  @doc """
+  Maps the given `fun` over `enumerable` and flattens the result.
+
+  This function returns a new enumerable built by appending the result of invoking `fun`
+  on each element of `enumerable` together; conceptually, this is similar to a
+  combination of `map/2` and `concat/1`.
+
+  ## Examples
+
+      iex> Enum.flat_map([:a, :b, :c], fn x -> [x, x] end)
+      [:a, :a, :b, :b, :c, :c]
+
+      iex> Enum.flat_map([{1, 3}, {4, 6}], fn {x, y} -> x..y end)
+      [1, 2, 3, 4, 5, 6]
+
+      iex> Enum.flat_map([:a, :b, :c], fn x -> [[x]] end)
+      [[:a], [:b], [:c]]
+
+  """
+  @spec flat_map(t, (element -> t)) :: list
+  def flat_map(enumerable, fun) when is_list(enumerable) do
+    flat_map_list(enumerable, fun)
+  end
+
+  def flat_map(enumerable, fun) do
+    reduce(enumerable, [], fn entry, acc ->
+      case fun.(entry) do
+        list when is_list(list) -> :lists.reverse(list, acc)
+        other -> reduce(other, acc, &[&1 | &2])
+      end
+    end)
+    |> :lists.reverse()
+  end
+
   @doc """
   Returns a list where each element is the result of invoking
   `fun` on each corresponding element of `enumerable`.
@@ -415,10 +595,6 @@ defmodule Enum do
     end
   end
 
-  def reject(enumerable, fun) when is_list(enumerable) do
-    reject_list(enumerable, fun)
-  end
-
   ## all?
 
   defp all_list([h | t], fun) do
@@ -499,6 +675,19 @@ defmodule Enum do
     default
   end
 
+  ## flat_map
+
+  defp flat_map_list([head | tail], fun) do
+    case fun.(head) do
+      list when is_list(list) -> list ++ flat_map_list(tail, fun)
+      other -> to_list(other) ++ flat_map_list(tail, fun)
+    end
+  end
+
+  defp flat_map_list([], _fun) do
+    []
+  end
+
   @doc """
   Inserts the given `enumerable` into a `collectable`.
 
@@ -650,6 +839,27 @@ defmodule Enum do
     []
   end
 
+  @doc """
+  Returns a list of elements in `enumerable` excluding those for which the function `fun` returns
+  a truthy value.
+
+  See also `filter/2`.
+
+  ## Examples
+
+      iex> Enum.reject([1, 2, 3], fn x -> rem(x, 2) == 0 end)
+      [1, 3]
+
+  """
+  @spec reject(t, (element -> as_boolean(term))) :: list
+  def reject(enumerable, fun) when is_list(enumerable) do
+    reject_list(enumerable, fun)
+  end
+
+  def reject(enumerable, fun) do
+    reduce(enumerable, [], R.reject(fun)) |> :lists.reverse()
+  end
+
   @doc """
   Returns a list of elements in `enumerable` in reverse order.
 

tests/libs/exavmlib/Tests.ex

@@ -39,9 +39,12 @@ defmodule Tests do
     [2, 3] = Enum.slice([1, 2, 3], 1, 2)
     :test = Enum.at([0, 1, :test, 3], 2)
     :atom = Enum.find([1, 2, :atom, 3, 4], -1, fn item -> not is_integer(item) end)
+    1 = Enum.find_index([:a, :b, :c], fn item -> item == :b end)
+    true = Enum.find_value([-2, -3, -1, 0, 1], fn item -> item >= 0 end)
     true = Enum.all?([1, 2, 3], fn n -> n >= 0 end)
     true = Enum.any?([1, -2, 3], fn n -> n < 0 end)
     [2] = Enum.filter([1, 2, 3], fn n -> rem(n, 2) == 0 end)
+    [1, 3] = Enum.reject([1, 2, 3], fn n -> rem(n, 2) == 0 end)
     :ok = Enum.each([1, 2, 3], fn n -> true = is_integer(n) end)
 
     # map
@@ -63,9 +66,11 @@ defmodule Tests do
     true = at_0 != at_1
     {:c, :atom} = Enum.find(%{a: 1, b: 2, c: :atom, d: 3}, fn {_k, v} -> not is_integer(v) end)
     {:d, 3} = Enum.find(%{a: 1, b: 2, c: :atom, d: 3}, fn {k, _v} -> k == :d end)
+    true = Enum.find_value(%{"a" => 1, b: 2}, fn {k, _v} -> is_atom(k) end)
     true = Enum.all?(%{a: 1, b: 2}, fn {_k, v} -> v >= 0 end)
     true = Enum.any?(%{a: 1, b: -2}, fn {_k, v} -> v < 0 end)
     [b: 2] = Enum.filter(%{a: 1, b: 2, c: 3}, fn {_k, v} -> rem(v, 2) == 0 end)
+    [] = Enum.reject(%{a: 1, b: 2, c: 3}, fn {_k, v} -> v > 0 end)
     :ok = Enum.each(%{a: 1, b: 2}, fn {_k, v} -> true = is_integer(v) end)
 
     # map set
@@ -80,9 +85,11 @@ defmodule Tests do
     true = ms_at_0 == 1 or ms_at_0 == 2
     true = ms_at_1 == 1 or ms_at_1 == 2
     :atom = Enum.find(MapSet.new([1, 2, :atom, 3, 4]), fn item -> not is_integer(item) end)
+    nil = Enum.find_value([-2, -3, -1, 0, 1], fn item -> item > 100 end)
     true = Enum.all?(MapSet.new([1, 2, 3]), fn n -> n >= 0 end)
     true = Enum.any?(MapSet.new([1, -2, 3]), fn n -> n < 0 end)
     [2] = Enum.filter(MapSet.new([1, 2, 3]), fn n -> rem(n, 2) == 0 end)
+    [1] = Enum.reject(MapSet.new([1, 2, 3]), fn n -> n > 1 end)
     :ok = Enum.each(MapSet.new([1, 2, 3]), fn n -> true = is_integer(n) end)
 
     # range
@@ -94,9 +101,11 @@ defmodule Tests do
     [6, 7, 8, 9, 10] = Enum.slice(1..10, 5, 100)
     7 = Enum.at(1..10, 6)
     8 = Enum.find(-10..10, fn item -> item >= 8 end)
+    true = Enum.find_value(-10..10, fn item -> item >= 0 end)
     true = Enum.all?(0..10, fn n -> n >= 0 end)
     true = Enum.any?(-1..10, fn n -> n < 0 end)
     [0, 1, 2] = Enum.filter(-10..2, fn n -> n >= 0 end)
+    [-1] = Enum.reject(-1..10, fn n -> n >= 0 end)
     :ok = Enum.each(-5..5, fn n -> true = is_integer(n) end)
 
     # into
@@ -105,6 +114,11 @@ defmodule Tests do
     expected_mapset = MapSet.new([1, 2, 3])
     ^expected_mapset = Enum.into([1, 2, 3], MapSet.new())
 
+    # Enum.flat_map
+    [:a, :a, :b, :b, :c, :c] = Enum.flat_map([:a, :b, :c], fn x -> [x, x] end)
+    [1, 2, 3, 4, 5, 6] = Enum.flat_map([{1, 3}, {4, 6}], fn {x, y} -> x..y end)
+    [[:a], [:b], [:c]] = Enum.flat_map([:a, :b, :c], fn x -> [[x]] end)
+
     # Enum.join
     "1, 2, 3" = Enum.join(["1", "2", "3"], ", ")
     "1, 2, 3" = Enum.join([1, 2, 3], ", ")
@@ -113,6 +127,9 @@ defmodule Tests do
     # Enum.reverse
     [4, 3, 2] = Enum.reverse([2, 3, 4])
 
+    # other enum functions
+    test_enum_chunk_while()
+
     undef =
       try do
         Enum.map({1, 2}, fn x -> x end)
@@ -132,6 +149,28 @@ defmodule Tests do
     :ok
   end
 
+  defp test_enum_chunk_while() do
+    initial_col = 4
+    lines_list = '-1234567890\nciao\n12345\nabcdefghijkl\n12'
+    columns = 5
+
+    chunk_fun = fn char, {count, rchars} ->
+      cond do
+        char == ?\n -> {:cont, Enum.reverse(rchars), {0, []}}
+        count == columns -> {:cont, Enum.reverse(rchars), {1, [char]}}
+        true -> {:cont, {count + 1, [char | rchars]}}
+      end
+    end
+
+    after_fun = fn
+      {_count, []} -> {:cont, [], []}
+      {_count, rchars} -> {:cont, Enum.reverse(rchars), []}
+    end
+
+    ['-', '12345', '67890', 'ciao', '12345', 'abcde', 'fghij', 'kl', '12'] =
+      Enum.chunk_while(lines_list, {initial_col, []}, chunk_fun, after_fun)
+  end
+
   defp test_exception() do
     ex1 =
       try do