diff --git a/lib/elixir/lib/enum.ex b/lib/elixir/lib/enum.ex index df8a6cfebc8..c92775a2d95 100644 --- a/lib/elixir/lib/enum.ex +++ b/lib/elixir/lib/enum.ex @@ -1855,6 +1855,27 @@ defmodule Enum do do_split_reverse(reverse(collection), abs(count), []) end + @doc """ + Chunks the `collection` each time `fun` returns `true` or at the beginning of + `collection`. + + ## Examples + + iex> Enum.slice_before([1, 2, 3, 6, 7, 9], fn(x) -> rem(x, 3) == 0 end) + [[1, 2], [3], [6, 7], [9]] + + """ + @spec slice_before(t, (element -> any)) :: [list] + def slice_before(collection, fun) do + {_, {acc, res}} = + Enumerable.reduce(collection, {:cont, {[], []}}, R.slice_before(fun)) + + case res do + [] -> [] + buffer -> :lists.reverse([:lists.reverse(buffer) | acc]) + end + end + @doc """ Splits `collection` in two while `fun` returns `true`. diff --git a/lib/elixir/lib/stream.ex b/lib/elixir/lib/stream.ex index b334eae19d1..a9435f51112 100644 --- a/lib/elixir/lib/stream.ex +++ b/lib/elixir/lib/stream.ex @@ -490,6 +490,33 @@ defmodule Stream do lazy enum, acc, fn(f1) -> R.scan_3(fun, f1) end end + @doc """ + Chunks the `enum` each time `fun` returns `true` or at the beginning of + `enum`. + + ## Examples + + iex> stream = Stream.slice_before(1..8, &(rem(&1, 3) == 0)) + iex> Enum.to_list(stream) + [[1,2], [3, 4, 5], [6, 7, 8]] + + """ + @spec slice_before(Enumerable.t, (element -> any)) :: Enumerable.t + def slice_before(enum, fun) do + lazy enum, + [], + fn(f1) -> R.slice_before(fun, f1) end, + &do_slice_before(&1, &2) + end + + defp do_slice_before(acc(_, nil, _) = acc, _) do + {:cont, acc} + end + + defp do_slice_before(acc(h, buffer, t), f1) do + cont_with_acc(f1, :lists.reverse(buffer), h, nil, t) + end + @doc """ Lazily takes the next `n` items from the enumerable and stops enumeration. diff --git a/lib/elixir/lib/stream/reducers.ex b/lib/elixir/lib/stream/reducers.ex index 18659600fd5..f79417b4e8d 100644 --- a/lib/elixir/lib/stream/reducers.ex +++ b/lib/elixir/lib/stream/reducers.ex @@ -129,6 +129,21 @@ defmodule Stream.Reducers do end end + defmacro slice_before(callback, f \\ nil) do + quote do + fn + entry, acc(h, [], t) -> + {:cont, acc(h, [entry], t)} + entry, acc(h, buffer, t) -> + if unquote(callback).(entry) do + cont_with_acc(unquote(f), :lists.reverse(buffer), h, [entry], t) + else + {:cont, acc(h, [entry|buffer], t)} + end + end + end + end + defmacro take(f \\ nil) do quote do fn(entry, acc(h, n, t) = orig) -> diff --git a/lib/elixir/test/elixir/enum_test.exs b/lib/elixir/test/elixir/enum_test.exs index 7bf39705dd4..a8223eabed0 100644 --- a/lib/elixir/test/elixir/enum_test.exs +++ b/lib/elixir/test/elixir/enum_test.exs @@ -374,6 +374,14 @@ defmodule EnumTest.List do assert Enum.split([1, 2, 3], -10) == {[], [1, 2, 3]} end + test :slice_before do + assert Enum.slice_before([1, 2, 2, 3, 4, 4, 6, 7, 7], &(rem(&1, 2) == 1)) == [[1, 2, 2], [3, 4, 4, 6], [7], [7]] + assert Enum.slice_before([1, 2, 3, 4], fn _ -> false end) == [[1, 2, 3, 4]] + assert Enum.slice_before([1, 2, 3, 4], fn _ -> true end) == [[1], [2], [3], [4]] + assert Enum.slice_before([], fn _ -> true end) == [] + assert Enum.slice_before([1], fn _ -> true end) == [[1]] + end + test :split_while do assert Enum.split_while([1, 2, 3], fn(_) -> false end) == {[], [1, 2, 3]} assert Enum.split_while([1, 2, 3], fn(_) -> true end) == {[1, 2, 3], []} @@ -1006,6 +1014,12 @@ defmodule EnumTest.Range do assert Enum.split(range, 3) == {[1, 0], []} end + test :slice_before do + assert Enum.slice_before(1..4, fn _ -> false end) == [[1, 2, 3, 4]] + assert Enum.slice_before(1..4, fn _ -> true end) == [[1], [2], [3], [4]] + assert Enum.slice_before(1..4, &(rem(&1, 2) == 1)) == [[1, 2], [3, 4]] + end + test :split_while do range = 1..3 assert Enum.split_while(range, fn(_) -> false end) == {[], [1, 2, 3]} diff --git a/lib/elixir/test/elixir/stream_test.exs b/lib/elixir/test/elixir/stream_test.exs index 9793c41468d..408e058c482 100644 --- a/lib/elixir/test/elixir/stream_test.exs +++ b/lib/elixir/test/elixir/stream_test.exs @@ -537,6 +537,22 @@ defmodule StreamTest do assert Stream.scan([], 0, &(&1 + &2)) |> Enum.to_list == [] end + test "slice_before/2" do + stream = Stream.slice_before([1, 2, 1, 1, 2, 3, 4, 1, 2, 3], &(&1 == 1)) + + assert is_lazy(stream) + assert Enum.to_list(stream) == + [[1, 2], [1], [1, 2, 3, 4], [1, 2, 3]] + assert stream |> Stream.take(3) |> Enum.to_list == + [[1, 2], [1], [1, 2, 3, 4]] + end + + test "slice_before/2 is zippable" do + stream = Stream.slice_before([1, 2, 2, 3], &(&1 == 2)) + list = Enum.to_list(stream) + assert Enum.zip(list, list) == Enum.zip(stream, stream) + end + test "take/2" do stream = Stream.take(1..1000, 5) assert is_lazy(stream)