Skip to content

Commit

Permalink
Optimize Enum.random/1 for ranges
Browse files Browse the repository at this point in the history
  • Loading branch information
josevalim committed Sep 20, 2023
1 parent 4e96f92 commit d101450
Showing 1 changed file with 20 additions and 18 deletions.
38 changes: 20 additions & 18 deletions lib/elixir/lib/enum.ex
Original file line number Diff line number Diff line change
Expand Up @@ -2390,7 +2390,14 @@ defmodule Enum do
def random(enumerable) when is_list(enumerable) do
case length(enumerable) do
0 -> raise Enum.EmptyError
length -> enumerable |> drop_list(random_integer(0, length - 1)) |> hd()
length -> enumerable |> drop_list(random_count(length)) |> hd()
end
end

def random(first.._//step = range) do
case Range.size(range) do
0 -> raise Enum.EmptyError
size -> first + random_count(size) * step
end
end

Expand All @@ -2401,14 +2408,14 @@ defmodule Enum do
[]

{:ok, count, fun} when is_function(fun, 1) ->
slice_list(fun.(enumerable), random_integer(0, count - 1), 1, 1)
slice_list(fun.(enumerable), random_count(count), 1, 1)

# TODO: Deprecate me in Elixir v1.18.
{:ok, count, fun} when is_function(fun, 2) ->
fun.(random_integer(0, count - 1), 1)
fun.(random_count(count), 1)

{:ok, count, fun} when is_function(fun, 3) ->
fun.(random_integer(0, count - 1), 1, 1)
fun.(random_count(count), 1, 1)

{:error, _} ->
take_random(enumerable, 1)
Expand All @@ -2420,6 +2427,10 @@ defmodule Enum do
end
end

defp random_count(count) do
:rand.uniform(count) - 1
end

@doc """
Invokes `fun` for each element in the `enumerable` with the
accumulator.
Expand Down Expand Up @@ -3609,7 +3620,7 @@ defmodule Enum do
sample = Tuple.duplicate(nil, count)

reducer = fn elem, {idx, sample} ->
jdx = random_integer(0, idx)
jdx = random_index(idx)

cond do
idx < count ->
Expand All @@ -3630,7 +3641,7 @@ defmodule Enum do

def take_random(enumerable, count) when is_integer(count) and count >= 0 do
reducer = fn elem, {idx, sample} ->
jdx = random_integer(0, idx)
jdx = random_index(idx)

cond do
idx < count ->
Expand Down Expand Up @@ -3666,6 +3677,9 @@ defmodule Enum do

defp take_random_list_one([], current, _), do: [current]

defp random_index(0), do: 0
defp random_index(idx), do: :rand.uniform(idx + 1) - 1

@doc """
Takes the elements from the beginning of the `enumerable` while `fun` returns
a truthy value.
Expand Down Expand Up @@ -4149,18 +4163,6 @@ defmodule Enum do
end)
end

defp random_integer(limit, limit) when is_integer(limit) do
limit
end

defp random_integer(lower_limit, upper_limit) when upper_limit < lower_limit do
random_integer(upper_limit, lower_limit)
end

defp random_integer(lower_limit, upper_limit) do
lower_limit + :rand.uniform(upper_limit - lower_limit + 1) - 1
end

## Implementations

## all?/1
Expand Down

0 comments on commit d101450

Please sign in to comment.