%% -*- 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.