hipe: Extract disjoint sets to its own module

1 files changed, 84 insertions, 0 deletions

diff --git a/lib/hipe/util/hipe_dsets.erl b/lib/hipe/util/hipe_dsets.erl
new file mode 100644
index 0000000000..9492cab0ff
--- /dev/null
+++ b/lib/hipe/util/hipe_dsets.erl
@@ -0,0 +1,84 @@
+%% -*- erlang-indent-level: 2 -*-
+%%
+%% Licensed under the Apache License, Version 2.0 (the "License");
+%% you may not use this file except in compliance with the License.
+%% You may obtain a copy of the License at
+%%
+%%     http://www.apache.org/licenses/LICENSE-2.0
+%%
+%% Unless required by applicable law or agreed to in writing, software
+%% distributed under the License is distributed on an "AS IS" BASIS,
+%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+%% See the License for the specific language governing permissions and
+%% limitations under the License.
+%%
+%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+%%@doc
+%%	       IMMUTABLE DISJOINT SETS OF ARBITRARY TERMS
+%%
+%% The disjoint set forests data structure, for elements of arbitrary types.
+%% Note that the find operation mutates the set.
+%%
+%% We could do this more efficiently if we restricted the elements to integers,
+%% and used the (mutable) hipe arrays. For arbitrary terms ETS could be used,
+%% for a persistent interface (which isn't that nice when even accessors return
+%% modified copies), the array module could be used.
+-module(hipe_dsets).
+
+-export([new/1, find/2, union/3, to_map/1, to_rllist/1]).
+-export_type([dsets/1]).
+
+-opaque dsets(X) :: #{X => {node, X} | {root, non_neg_integer()}}.
+
+-spec new([E]) -> dsets(E).
+new(Elems) -> maps:from_list([{E,{root,0}} || E <- Elems]).
+
+-spec find(E, dsets(E)) -> {E, dsets(E)}.
+find(E, DS0) ->
+  case DS0 of
+    #{E := {root,_}} -> {E, DS0};
+    #{E := {node,N}} ->
+      case find(N, DS0) of
+	{N, _}=T -> T;
+	{R, DS1} -> {R, DS1#{E := {node,R}}}
+      end;
+    _ -> error(badarg, [E, DS0])
+  end.
+
+-spec union(E, E, dsets(E)) -> dsets(E).
+union(X, Y, DS0) ->
+  {XRoot, DS1} = find(X, DS0),
+  case find(Y, DS1) of
+    {XRoot, DS2} -> DS2;
+    {YRoot, DS2} ->
+      #{XRoot := {root,XRR}, YRoot := {root,YRR}} = DS2,
+      if XRR < YRR -> DS2#{XRoot := {node,YRoot}};
+	 XRR > YRR -> DS2#{YRoot := {node,XRoot}};
+	 true -> DS2#{YRoot := {node,XRoot}, XRoot := {root,XRR+1}}
+      end
+  end.
+
+-spec to_map(dsets(E)) -> {#{Elem::E => Root::E}, dsets(E)}.
+to_map(DS) ->
+  to_map(maps:keys(DS), DS, #{}).
+
+to_map([], DS, Acc) -> {Acc, DS};
+to_map([K|Ks], DS0, Acc) ->
+  {KR, DS} = find(K, DS0),
+  to_map(Ks, DS, Acc#{K => KR}).
+
+-spec to_rllist(dsets(E)) -> {[{Root::E, Elems::[E]}], dsets(E)}.
+to_rllist(DS0) ->
+  {Lists, DS} = to_rllist(maps:keys(DS0), #{}, DS0),
+  {maps:to_list(Lists), DS}.
+
+to_rllist([], Acc, DS) -> {Acc, DS};
+to_rllist([E|Es], Acc, DS0) ->
+  {ERoot, DS} = find(E, DS0),
+  to_rllist(Es, map_append(ERoot, E, Acc), DS).
+
+map_append(Key, Elem, Map) ->
+  case Map of
+    #{Key := List} -> Map#{Key := [Elem|List]};
+    #{} -> Map#{Key => [Elem]}
+  end.