Implement a bias free rand:shuffle algorithm

Author: RaimoNiskanen

lib/stdlib/src/rand.erl

@@ -405,8 +405,7 @@ the generator's range:
          bytes/1, bytes_s/2,
          jump/0, jump/1,
          normal/0, normal/2, normal_s/1, normal_s/3,
-         shuffle1/1, shuffle1_s/2, shuffle2/1, shuffle2_s/2,
-         shuffle3/1, shuffle3_s/2, shuffle4/1, shuffle4_s/2
+         shuffle/1, shuffle_s/2
 	]).
 
 %% Utilities
@@ -1317,16 +1316,15 @@ normal_s(Mean, Variance, State0) when 0 =< Variance ->
     {X, State} = normal_s(State0),
     {Mean + (math:sqrt(Variance) * X), State}.
 
-%% -------
 
--spec shuffle1(list()) -> list().
-shuffle1(List) ->
-    {ShuffledList, State} = shuffle1_s(List, seed_get()),
+-spec shuffle(list()) -> list().
+shuffle(List) ->
+    {ShuffledList, State} = shuffle_s(List, seed_get()),
     _ = seed_put(State),
     ShuffledList.
 
--spec shuffle1_s(list(), state()) -> {list(), state()}.
-shuffle1_s(List, {#{bits:=_, next:=Next} = AlgHandler, R0} = State)
+-spec shuffle_s(list(), state()) -> {list(), state()}.
+shuffle_s(List, {#{bits:=_, next:=Next} = AlgHandler, R0} = State)
   when is_list(List) ->
     case List of
         [] ->
@@ -1335,10 +1333,10 @@ shuffle1_s(List, {#{bits:=_, next:=Next} = AlgHandler, R0} = State)
             {List, State};
         _ ->
             WeakLowBits = maps:get(weak_low_bits, AlgHandler, 0),
-            {ShuffledList, R1} = shuffle1_r(List, Next, R0, WeakLowBits, []),
+            {ShuffledList, R1} = shuffle_r(List, Next, R0, WeakLowBits, []),
             {ShuffledList, {AlgHandler, R1}}
     end;
-shuffle1_s(List, {#{max:=Mask, next:=Next} = AlgHandler, R0} = State)
+shuffle_s(List, {#{max:=Mask, next:=Next} = AlgHandler, R0} = State)
   when is_list(List), ?MASK(58) =< Mask ->
     case List of
         [] ->
@@ -1348,7 +1346,7 @@ shuffle1_s(List, {#{max:=Mask, next:=Next} = AlgHandler, R0} = State)
         _ ->
             %% Old spec - assume 2 weak low bits
             WeakLowBits = 2,
-            {ShuffledList, R1} = shuffle1_r(List, Next, R0, WeakLowBits, []),
+            {ShuffledList, R1} = shuffle_r(List, Next, R0, WeakLowBits, []),
             {ShuffledList, {AlgHandler, R1}}
     end.
 
@@ -1360,204 +1358,42 @@ shuffle1_s(List, {#{max:=Mask, next:=Next} = AlgHandler, R0} = State)
 %% produces a bias free shuffle.
 
 %% Recursion entry point
-shuffle1_r([X, Y], Next, R0, _WeakLowBits, Acc) ->
+shuffle_r([X, Y], Next, R0, _WeakLowBits, Acc) ->
     %% Optimization for 2 elements; the most common case for duplicates
     {V, R1} = Next(R0),
     if
         %% Bit 7 should not be weak in any of the generators
         V band 128 =:= 0 -> {[Y, X | Acc], R1};
         true             -> {[X, Y | Acc], R1}
     end;
-shuffle1_r(L, Next, R0, WeakLowBits, Acc) ->
-    shuffle1_tag(L, Next, R0, WeakLowBits, Acc, []).
+shuffle_r(L, Next, R0, WeakLowBits, Acc) ->
+    shuffle_tag(L, Next, R0, WeakLowBits, Acc, []).
 
 %% Tag elements with random integers
-shuffle1_tag([], Next, R, WeakLowBits, Acc, TL) ->
-    %% Shuffle1; sort by random tag
-    shuffle1_untag(lists:keysort(1, TL), Next, R, WeakLowBits, Acc);
-shuffle1_tag([X | L], Next, R0, WeakLowBits, Acc, TL) ->
+shuffle_tag([], Next, R, WeakLowBits, Acc, TL) ->
+    %% Shuffle; sort by random tag
+    shuffle_untag(lists:keysort(1, TL), Next, R, WeakLowBits, Acc);
+shuffle_tag([X | L], Next, R0, WeakLowBits, Acc, TL) ->
     {V, R1} = Next(R0),
     T = V bsr WeakLowBits,
-    shuffle1_tag(L, Next, R1, WeakLowBits, Acc, [{T,X} | TL]).
+    shuffle_tag(L, Next, R1, WeakLowBits, Acc, [{T,X} | TL]).
 
 %% Strip the tag integers
-shuffle1_untag([{T,X}, {T,Y} | TL], Next, R, WeakLowBits, Acc) ->
+shuffle_untag([{T,X}, {T,Y} | TL], Next, R, WeakLowBits, Acc) ->
     %% Random number duplicate
-    shuffle1_untag(TL, Next, R, WeakLowBits, Acc, [Y, X], T);
-shuffle1_untag([{_,X} | TL], Next, R, WeakLowBits, Acc) ->
-    shuffle1_untag(TL, Next, R, WeakLowBits, [X | Acc]);
-shuffle1_untag([], _Next, R, _WeakLowBits, Acc) ->
+    shuffle_untag(TL, Next, R, WeakLowBits, Acc, [Y, X], T);
+shuffle_untag([{_,X} | TL], Next, R, WeakLowBits, Acc) ->
+    shuffle_untag(TL, Next, R, WeakLowBits, [X | Acc]);
+shuffle_untag([], _Next, R, _WeakLowBits, Acc) ->
     {Acc, R}.
 %%
 %% Collect duplicates
-shuffle1_untag([{T,X} | TL], Next, R, WeakLowBits, Acc, Dups, T) ->
-    shuffle1_untag(TL, Next, R, WeakLowBits, Acc, [X | Dups], T);
-shuffle1_untag(TL, Next, R0, WeakLowBits, Acc0, Dups, _T) ->
-    %% Shuffle1 the duplicates onto the result
-    {Acc1, R1} = shuffle1_r(Dups, Next, R0, WeakLowBits, Acc0),
-    shuffle1_untag(TL, Next, R1, WeakLowBits, Acc1).
-
-%% -------
-
--spec shuffle2(list()) -> list().
-shuffle2(List) ->
-    {ShuffledList, State} = shuffle2_s(List, seed_get()),
-    _ = seed_put(State),
-    ShuffledList.
-
--spec shuffle2_s(list(), state()) -> {list(), state()}.
-shuffle2_s(List, State)
-  when is_list(List) ->
-    case List of
-        [] ->
-            {List, State};
-        [_] ->
-            {List, State};
-        _ ->
-            M = maps:from_list(lists:enumerate(List)),
-            N = maps:size(M),
-            shuffle2_s(M, State, N, [])
-    end.
-
-%% Classical Fisher-Yates shuffle, a.k.a Knuth shuffle.
-%% See the Wikipedia article "Fisher-Yates shuffle".
-%%
-%% This variant uses a map with integer keys as array
-%% and is optimized in that it minimizes map updates
-%% since the high index is never used again, so an overwrite
-%% can be used instead of an exchange.
-
-shuffle2_s(M0, State0, N, Acc)
-  when is_map(M0), is_integer(N) ->
-    if
-        N =:= 0 -> {Acc, State0};
-        true ->
-            X = maps:get(N, M0),
-            case uniform_s(N, State0) of
-                {N, State1} ->
-                    shuffle2_s(M0, State1, N - 1, [X | Acc]);
-                {K, State1} when is_integer(K) ->
-                    Y = maps:get(K, M0),
-                    M1 = maps:update(K, X, M0),
-                    shuffle2_s(M1, State1, N - 1, [Y | Acc])
-            end
-    end.
-
-%% -------
-
--spec shuffle3(list()) -> list().
-shuffle3(List) ->
-    {ShuffledList, State} = shuffle3_s(List, seed_get()),
-    _ = seed_put(State),
-    ShuffledList.
-
--spec shuffle3_s(list(), state()) -> {list(), state()}.
-shuffle3_s(List, {#{bits:=_, next:=Next} = AlgHandler, R0} = State)
-  when is_list(List) ->
-    case List of
-        [] ->
-            {List, State};
-        [_] ->
-            {List, State};
-        _ ->
-            WeakLowBits = maps:get(weak_low_bits, AlgHandler, 0),
-            T = gb_trees:empty(),
-            {ShuffledList, R1} = shuffle3_r(List, Next, R0, WeakLowBits, T),
-            {ShuffledList, {AlgHandler, R1}}
-    end;
-shuffle3_s(List, {#{max:=Mask, next:=Next} = AlgHandler, R0} = State)
-  when is_list(List), ?MASK(58) =< Mask ->
-    case List of
-        [] ->
-            {List, State};
-        [_] ->
-            {List, State};
-        _ ->
-            %% Old spec - assume 2 weak low bits
-            WeakLowBits = 2,
-            T = gb_trees:empty(),
-            {ShuffledList, R1} = shuffle3_r(List, Next, R0, WeakLowBits, T),
-            {ShuffledList, {AlgHandler, R1}}
-    end.
-
-%% See the Wikipedia article "Fisher-Yates shuffle", section "Sorting".
-%%
-%% To avoid bias due to duplicate random numbers, a gb_tree
-%% is used to check if a random number has already been used,
-%% and if so generate a new random number.
-%%
-%% Because a gb_tree is sorted no sorting needs to be done,
-%% it is enough to extract the values of the gb_tree that are
-%% ordered in key sort order.
-
-shuffle3_r([], _Next, R, _WeakLowBits, T) ->
-    {gb_trees:values(T), R};
-shuffle3_r([X | L] , Next, R0, WeakLowBits, T) ->
-    {V, R1} = Next(R0),
-    K = V bsr WeakLowBits,
-    case gb_trees:is_defined(K, T) of
-        false ->
-            shuffle3_r(L, Next, R1, WeakLowBits, gb_trees:insert(K, X, T));
-        true ->
-            shuffle3_r([X | L], Next, R1, WeakLowBits, T)
-    end.
-
-%% -------
-
--spec shuffle4(list()) -> list().
-shuffle4(List) ->
-    {ShuffledList, State} = shuffle4_s(List, seed_get()),
-    _ = seed_put(State),
-    ShuffledList.
-
--spec shuffle4_s(list(), state()) -> {list(), state()}.
-shuffle4_s(List, {#{bits:=_, next:=Next} = AlgHandler, R0} = State)
-  when is_list(List) ->
-    case List of
-        [] ->
-            {List, State};
-        [_] ->
-            {List, State};
-        _ ->
-            WeakLowBits = maps:get(weak_low_bits, AlgHandler, 0),
-            {ShuffledList, R1} = shuffle4_r(List, Next, R0, WeakLowBits, #{}),
-            {ShuffledList, {AlgHandler, R1}}
-    end;
-shuffle4_s(List, {#{max:=Mask, next:=Next} = AlgHandler, R0} = State)
-  when is_list(List), ?MASK(58) =< Mask ->
-    case List of
-        [] ->
-            {List, State};
-        [_] ->
-            {List, State};
-        _ ->
-            %% Old spec - assume 2 weak low bits
-            WeakLowBits = 2,
-            {ShuffledList, R1} = shuffle4_r(List, Next, R0, WeakLowBits, #{}),
-            {ShuffledList, {AlgHandler, R1}}
-    end.
-
-%% See the Wikipedia article "Fisher-Yates shuffle", section "Sorting".
-%%
-%% To avoid bias due to duplicate random numbers, a map
-%% is used to check if a random number has already been used,
-%% and if so generate a new random number.
-%%
-%% Actual sorting doesn't is not needed.  A map is ordered by key
-%% and therefore it is enough to extract the values of the map.
-%% The internal map key order will do just fine.
-
-shuffle4_r([], _Next, R, _WeakLowBits, M) ->
-    {maps:values(M), R};
-shuffle4_r([X | L] , Next, R0, WeakLowBits, M) ->
-    {V, R1} = Next(R0),
-    K = V bsr WeakLowBits,
-    case maps:is_key(K, M) of
-        true ->
-            shuffle4_r([X | L], Next, R1, WeakLowBits, M);
-        false ->
-            shuffle4_r(L, Next, R1, WeakLowBits, maps:put(K, X, M))
-    end.
+shuffle_untag([{T,X} | TL], Next, R, WeakLowBits, Acc, Dups, T) ->
+    shuffle_untag(TL, Next, R, WeakLowBits, Acc, [X | Dups], T);
+shuffle_untag(TL, Next, R0, WeakLowBits, Acc0, Dups, _T) ->
+    %% Shuffle the duplicates onto the result
+    {Acc1, R1} = shuffle_r(Dups, Next, R0, WeakLowBits, Acc0),
+    shuffle_untag(TL, Next, R1, WeakLowBits, Acc1).
 
 %% =====================================================================
 %% Internal functions