fingertree_simple.erl

%%%-------------------------------------------------------------------
%%% @author Alex S
%%% @copyright (C) 2016, Alex S
%%% @doc
%%% Homegrown Hinze/Paterson FingerTree implementation optimized for speed and
%%% occasionally size.
%%% Note that despite Erlang being an eager language, we do not use
%%% Kaplan-Tarjan implementation because I didn't want to descend into
%%% madness of stacks of stacks upon stacks; thus, this implementation
%%% gives no real-time guarantees but good amortized guarantees.
%%% @end
%%% Created : 2016-09-19 14:22
%%%-------------------------------------------------------------------
-module(fingertree_simple).
-author("alex0player@gmail.com").


%% API
-export([new/0, is_empty/1,
         peek_left/1, peek_right/1,
         emplace_left/2, emplace_right/2,
         push_left/2, push_right/2,
         pop_left/1, pop_right/1,
         find_index/2, find_count/2, get_count/1]).

%% Types. We use lots of untagged shit, so help us dialyzer.
-export_type([element/0, tree/0, split/0]).

-type element() :: {number(), number()}.  % index, count
-type measure() :: {number() | undefined, number()}. % max(index), sum(count)
-type node(X) ::
  {measure(), X, X} |
  {measure(), X, X, X}.
-type digit(X) ::
  {measure(), X} |
  {measure(), X, X} |
  {measure(), X, X, X} |
  {measure(), X, X, X, X}.
-type tree(X) ::
  {} |
  {X} |
  {measure(), digit(X), tree(node(X)), digit(X)}.
-type some(X) :: {} | {X} | {X, X} | {X, X, X} | {X, X, X, X}.

-type desc() :: element() | node(desc()).

-opaque tree() :: tree(element()).
-type split() :: {tree(), tree()} | {tree(), element(), tree()}.

%% Constructors.
-spec tree(Prefix :: digit(X), Middle :: tree(node(X)), Suffix :: digit(X)) -> tree(X) when X :: desc().
tree(Prefix, Middle, Suffix) ->
  {measure_combine(digit_measure(Prefix), tree_measure(Middle), digit_measure(Suffix)),
    Prefix, Middle, Suffix}.


-compile({inline, [node/2, node/3]}).

-spec node(X, X) -> node(X) when X :: desc().
node(Left, Right) ->
  {measure_combine(
    desc_measure(Left), desc_measure(Right)
  ), Left, Right}.
-spec node(X, X, X) -> node(X) when X :: desc().
node(Left, Middle, Right) ->
  {measure_combine(
    desc_measure(Left), desc_measure(Middle), desc_measure(Right)
  ), Left, Middle, Right}.


-compile({inline, [digit/1, digit/2, digit/3, digit/4]}).
-spec digit(X) -> digit(X) when X :: desc().
digit(X) ->
  {desc_measure(X), X}.
-spec digit(X, X) -> digit(X) when X :: desc().
digit(X, Y) ->
  {measure_combine(
    desc_measure(X),
    desc_measure(Y)
  ), X, Y}.
-spec digit(X, X, X) -> digit(X) when X :: desc().
digit(X, Y, Z) ->
  {measure_combine(
    desc_measure(X),
    desc_measure(Y),
    desc_measure(Z)
  ), X, Y, Z}.
-spec digit(X, X, X, X) -> digit(X) when X :: desc().
digit(X, Y, Z, W) ->
  {measure_combine(
    desc_measure(X),
    desc_measure(Y),
    desc_measure(Z),
    desc_measure(W)
  ), X, Y, Z, W}.


%% Measures.
-compile({inline, [desc_measure/1, digit_measure/1, tree_measure/1]}).
desc_measure(_Node3 = {Measure, _, _, _}) ->
  Measure;
desc_measure(_Node2 = {Measure, _, _}) ->
  Measure;
desc_measure({Index, Count}) ->
  {Index, Count}.

digit_measure(Digit) ->
  element(1, Digit).

-define(NULL, {undefined, 0}).

tree_measure({}) ->
  ?NULL;
tree_measure({X}) ->
  desc_measure(X);
tree_measure({Measure, _, _, _}) ->
  Measure.


%% Assuming sorted input.
-compile({inline, [measure_combine/2, measure_combine/3, measure_combine/4]}).
measure_combine({_I1, C1}, {I2, C2}) ->
  {I2, C1 + C2}.
measure_combine({_I1, C1}, {_I2, C2}, {I3, C3}) ->
  {I3, C1 + C2 + C3}.
measure_combine({_I1, C1}, {_I2, C2}, {_I3, C3}, {I4, C4}) ->
  {I4, C1 + C2 + C3 + C4}.

%% conversions
-compile({inline, [node_to_digit/1, digit_to_node/1]}).
-spec digit_to_tree(digit(X)) -> tree(X) when X :: desc().
digit_to_tree({_Measure, X}) ->
  {X};
digit_to_tree({Measure, X, Y}) ->
  {Measure, digit(X), {}, digit(Y)};
digit_to_tree({Measure, X, Y, Z}) ->
  {Measure, digit(X, Y), {}, digit(Z)};
digit_to_tree({Measure, X, Y, Z, W}) ->
  {Measure, digit(X, Y), {}, digit(Z, W)}.

-spec node_to_digit(node(X)) -> digit(X) when X :: desc().
node_to_digit(X) ->
  X.

-spec some_to_digit(some(X)) -> digit(X) when X :: desc().
some_to_digit({X}) ->
  digit(X);
some_to_digit({X, Y}) ->
  digit(X, Y);
some_to_digit({X, Y, Z}) ->
  digit(X, Y, Z);
some_to_digit({X, Y, Z, W}) ->
  digit(X, Y, Z, W).

digit_to_some(X) when tuple_size(X) >= 2; tuple_size(X) =< 5 ->
  delete_element(1, X).

some_to_tree(X) ->
  digit_to_tree(some_to_digit(X)).

%% deque operations

-spec tree_left(some(X), tree(node(X)), digit(X)) -> tree(X) when X :: desc().
tree_left({}, {}, Suffix) ->
  digit_to_tree(Suffix);
tree_left({}, Middle, Suffix) ->
  {Node, Middle1} = pop_left(Middle),
  tree(node_to_digit(Node), Middle1, Suffix);
tree_left(Some, Middle, Suffix) ->
  tree(some_to_digit(Some), Middle, Suffix).

-spec pop_left(tree(X)) -> {X, tree(X)} when X :: desc().
pop_left({}) ->
  error(badarg);
pop_left({X}) ->
  {X, {}};
pop_left({_, Prefix, Middle, Suffix}) ->
  case Prefix of
    {_, X} ->
      {X, tree_left({}, Middle, Suffix)};
    {_, X, Y} ->
      {X, tree(digit(Y), Middle, Suffix)};
    {_, X, Y, Z} ->
      {X, tree(digit(Y, Z), Middle, Suffix)};
    {_, X, Y, Z, W} ->
      {X, tree(digit(Y, Z, W), Middle, Suffix)}
  end.

-spec tree_right(digit(X), tree(node(X)), some(X)) -> tree(X) when X :: desc().
tree_right(Prefix, {}, {}) ->
  digit_to_tree(Prefix);
tree_right(Prefix, Middle, {}) ->
  {Middle1, Node} = pop_right(Middle),
  tree(Prefix, Middle1, node_to_digit(Node));
tree_right(Prefix, Middle, Some) ->
  tree(Prefix, Middle, some_to_digit(Some)).

-spec pop_right(tree(X)) -> {tree(X), X} when X :: desc().
pop_right({}) ->
  error(badarg);
pop_right({X}) ->
  {{}, X};
pop_right({_, Prefix, Middle, Suffix}) ->
  case Suffix of
    {_, X} ->
      {tree_right(Prefix, Middle, {}), X};
    {_, Y, X} ->
      {tree(Prefix, Middle, digit(Y)), X};
    {_, Z, Y, X} ->
      {tree(Prefix, Middle, digit(Z, Y)), X};
    {_, W, Z, Y, X} ->
      {tree(Prefix, Middle, digit(W, Z, Y)), X}
  end.

-spec push_left(X, tree(X)) -> tree(X) when X :: desc().
push_left(X, {}) ->
  {X};
push_left(X, {Y}) ->
  tree(digit(X), {}, digit(Y));
push_left(X, {_, Prefix, Middle, Suffix}) ->
  case Prefix of
    {_, Y} ->
      tree(digit(X, Y), Middle, Suffix);
    {_, Y, Z} ->
      tree(digit(X, Y, Z), Middle, Suffix);
    {_, Y, Z, W} ->
      tree(digit(X, Y, Z, W), Middle, Suffix);
    {_, Y, Z, W, TooMuch} ->
      Middle1 = push_left(node(Z, W, TooMuch), Middle),
      tree(digit(X, Y), Middle1, Suffix)
  end.

-spec push_right(tree(X), X) -> tree(X) when X :: desc().
push_right({}, X) ->
  {X};
push_right({Y}, X) ->
  tree(digit(Y), {}, digit(X));
push_right({_, Prefix, Middle, Suffix}, X) ->
  case Suffix of
    {_, Y} ->
      tree(Prefix, Middle, digit(Y, X));
    {_, Z, Y} ->
      tree(Prefix, Middle, digit(Z, Y, X));
    {_, W, Z, Y} ->
      tree(Prefix, Middle, digit(W, Z, Y, X));
    {_, TooMuch, W, Z, Y} ->
      Middle1 = push_right(Middle, node(TooMuch, W, Z)),
      tree(digit(X, Y), Middle1, Suffix)
  end.

%% Merging.
push_left_some({}, Tree) ->
  Tree;
push_left_some({X}, Tree) ->
  push_left(X, Tree);
push_left_some({X, Y}, Tree) ->
  push_left(X, push_left(Y, Tree));
push_left_some({X, Y, Z}, Tree) ->
  push_left(X, push_left(Y, push_left(Z, Tree)));
push_left_some({X, Y, Z, W}, Tree) ->
  push_left(X, push_left(Y, push_left(Z, push_left(W, Tree)))).

push_right_some(Tree, {}) ->
  Tree;
push_right_some(Tree, {X}) ->
  push_right(Tree, X);
push_right_some(Tree, {Y, X}) ->
  push_right(push_right(Tree, Y), X);
push_right_some(Tree, {Z, Y, X}) ->
  push_right(push_right(push_right(Tree, Z), Y), X);
push_right_some(Tree, {W, Z, Y, X}) ->
  push_right(push_right(push_right(push_right(Tree, W), Z), Y), X).

-spec join_some(digit(X), some(X), digit(X)) -> some(node(X)) when X :: desc().
-include("ftree_generate.hrl").

-spec join_trees(tree(X), some(X), tree(X)) -> tree(X) when X :: desc().
join_trees({}, Some, TreeR) ->
  push_left_some(Some, TreeR);
join_trees(TreeL, Some, {}) ->
  push_right_some(TreeL, Some);
join_trees({X}, Some, TreeR) ->
  push_left(X, push_left_some(Some, TreeR));
join_trees(TreeL, Some, {X}) ->
  push_right(push_right_some(TreeL, Some), X);
join_trees({_, Left, MiddleLeft, InnerLeft}, Some, {_, InnerRight, MiddleRight, Right}) ->
  SomeDeep = join_some(InnerLeft, Some, InnerRight),
  tree(Left, join_trees(MiddleLeft, SomeDeep, MiddleRight), Right).

%% Splitting on a predicate, it may return true | false.

-type predicate() :: fun((measure()) -> boolean()).
-type split(X, C) :: {C, X, C}.

%% For ease, we always split with the
-spec split_digit(predicate(), measure(), digit(X)) -> split(X, some(X)) when X :: desc().
split_digit(_Pred, _P0, {_, X}) ->
  {{}, X, {}};
split_digit(Pred, P0, {_, X, Y}) ->
  Xm = measure_combine(P0, desc_measure(X)),
  case Pred(desc_measure(X)) of
    false ->
      {{X}, Y, {}};
    _ ->
      {{}, X, {Y}}
  end;
split_digit(Pred, P0, {_, X, Y, Z}) ->
  Xm = measure_combine(P0, desc_measure(X)),
  XYm = measure_combine(Xm, desc_measure(Y)),
  case Pred(Xm) of
    true ->
      {{}, X, {Y, Z}};
    _ ->
      case Pred(XYm) of
        true ->
          {{X}, Y, {Z}};
        _ ->
          {{X, Y}, Z}
      end
  end;
split_digit(Pred, P0, {_, X, Y, Z, W}) ->
  Xm = measure_combine(P0, desc_measure(X)),
  XYm = measure_combine(Xm, desc_measure(Y)),
  XYZm = measure_combine(XYm, desc_measure(Z)),
  case Pred(XYm) of
    true ->
      case Pred(XYZm) of
        true ->
          {{X, Y}, Z, {W}};
        _ ->
          {{X, Y, Z}, W, {}}
      end;
    _ ->
      case Pred(Xm) of
        true ->
          {{}, X, {Y, Z, W}};
        _ ->
          {{X}, Y, {Z, W}}
      end
  end.

-spec split_tree(predicate(), measure(), tree(X)) -> split(X, tree(X)).
split_tree(_, _, {}) ->
  error(badarg);
split_tree(_, _P0, {X}) ->
  {{}, X, {}};
split_tree(Pred, P0, {_, Left, Middle, Right}) ->
  Pm = measure_combine(P0, digit_measure(Left)),
  PMm = measure_combine(Pm, tree_measure(Middle)),
  case Pred(Pm) of
    true ->
      {SomeLeft, X, SomeRight} = split_digit(Pred, P0, Left),
      {some_to_tree(SomeLeft), X, tree_left(SomeRight, Middle, Right)};
    _ ->
      case Pred(PMm) of
        true ->
          {LeftSub, X, RightSub} = split_tree(Pred, Pm, Middle),
          {push_left_some(digit_to_some(Left), LeftSub), X, push_right_some(RightSub, digit_to_some(Right))};
        _ ->
          {SomeLeft, X, SomeRight} = split_digit(Pred, PMm, Right),
          {tree_right(Left, Middle, SomeLeft), X, some_to_tree(SomeRight)}
      end
  end.

%% API
-spec new() -> tree().
new() ->
  {}.

-spec is_empty(tree()) -> boolean().
is_empty({}) ->
  true;
is_empty(_) ->
  false.

-spec peek_left(tree()) -> element().
peek_left({X}) ->
  X;
peek_left({_, {_, X}, _, _}) ->
  X;
peek_left({_, {_, X, _}, _, _}) ->
  X;
peek_left({_, {_, X, _, _}, _, _}) ->
  X;
peek_left({_, {_, X, _, _, _}, _, _}) ->
  X.

-spec emplace_left(element(), tree()) -> tree().
emplace_left(X, {_}) ->
  {X};
emplace_left(X, {_, {_, _}, Middle, Right}) ->
  tree(digit(X), Middle, Right);
emplace_left(X, {_, {_, _, Y}, Middle, Right}) ->
  tree(digit(X, Y), Middle, Right);
emplace_left(X, {_, {_, _, Y, Z}, Middle, Right}) ->
  tree(digit(X, Y, Z), Middle, Right);
emplace_left(X, {_, {_, _, Y, Z, W}, Middle, Right}) ->
  tree(digit(X, Y, Z, W), Middle, Right).

-spec peek_right(tree()) -> element().
peek_right({X}) ->
  X;
peek_right({_, _, _, {_, X}}) ->
  X;
peek_right({_, _, _, {_, _, X}}) ->
  X;
peek_right({_, _, _, {_, _, _, X}}) ->
  X;
peek_right({_, _, _, {_, _, _, _, X}}) ->
  X.

-spec emplace_right(tree(), element()) -> tree().
emplace_right({_}, X) ->
  {X};
emplace_right({_, Left, Middle, {_, _}}, X) ->
  tree(Left, Middle, digit(X));
emplace_right({_, Left, Middle, {_, Y, _}}, X) ->
  tree(Left, Middle, digit(Y, X));
emplace_right({_, Left, Middle, {_, Z, Y, _}}, X) ->
  tree(Left, Middle, digit(Z, Y, X));
emplace_right({_, Left, Middle, {_, W, Z, Y, _}}, X) ->
  tree(Left, Middle, digit(W, Z, Y, X)).

-spec find_index(tree(), number()) -> split().
find_index(Tree, Idx) ->
  Pred = fun({X, _}) when is_number(X), X >= Idx -> true; (_) -> false end,
  FullM = tree_measure(Tree),
  case Pred(FullM) of
    true ->
      {TreeLeft, Middle, TreeRight} = split_tree(Pred, ?NULL, Tree),
      {TreeLeft, push_left(Middle, TreeRight)};
    _ ->
      {Tree, {}}
  end.

-spec find_count(tree(), number()) -> split().
find_count(Tree, Count) ->
  Pred = fun({_, C}) when C >= Count -> true; (_) -> false end,
  {TreeLeft, Middle, TreeRight} = split_tree(Pred, ?NULL, Tree),
  {TreeLeft, Middle, TreeRight}.

-spec get_count(tree()) -> number().
get_count(Tree) ->
  {_, Count} = tree_measure(Tree),
  Count.

-spec to_list(tree()) -> [element()].
to_list(Tree) ->
  lists:reverse(tree_to_list(Tree)).

-spec tree_to_list(tree(X)) -> [X] when X :: desc().
tree_to_list({}) ->
  [];
tree_to_list({X}) ->
  [X];
tree_to_list({_, Prefix, Middle, Suffix}) ->
  Acc1 = digit_to_list(Prefix, []),
  Acc2 = explode_node(tree_to_list(Middle), Acc1),
  Acc3 = digit_to_list(Suffix, Acc2),
  Acc3.

digit_to_list({_, X}, Acc) ->
  [X | Acc];
digit_to_list({_, X, Y}, Acc) ->
  [Y, X | Acc];
digit_to_list({_, X, Y, Z}, Acc) ->
  [Z, Y, X | Acc];
digit_to_list({_, X, Y, Z, W}, Acc) ->
  [W, Z, Y, X | Acc].

explode_node(List, Acc) ->
  lists:foldl(fun({_, X, Y}, Acc1) -> [Y, X | Acc1];
                 ({_, X, Y, Z}, Acc1) -> [Z, Y, X | Acc1]
  end, Acc, List).


-spec from_list([element()]) -> tree().
from_list(List) ->
  tree_from_list(List, length(List)).

-spec tree_from_list([X], non_neg_integer()) -> tree(X) when X :: desc().
tree_from_list([], 0) ->
  {};
tree_from_list([A], 1) ->
  {A};
tree_from_list(List, N) when N =< 8 ->
  LeftSize = N div 2,
  {Left, Right} = lists:split(LeftSize, List),
  tree(digit_from_list(Left), {}, digit_from_list(Right));
tree_from_list([A, B, C, D, E, F, G, H, I], 9) ->
  tree(digit(A, B, C), {node(D, E, F)}, digit(G, H, I));
tree_from_list(List, N) ->
  {Left, MidRight} = lists:split(4, List),
  {Mid, Right} = conjoin_node(N - 8, MidRight, []),
  tree(digit_from_list(Left), tree_from_list(Mid, length(Mid)), digit_from_list(Right)).

conjoin_node(2, [A, B | Rest], Acc) ->
  {lists:reverse(Acc, [node(A, B)]), Rest};
conjoin_node(3, [A, B, C | Rest], Acc) ->
  {lists:reverse(Acc, [node(A, B, C)]), Rest};
conjoin_node(4, [A, B, C, D | Rest], Acc) ->
  {lists:reverse(Acc, [node(A, B), node(C, D)]), Rest};
conjoin_node(N, [A, B, C | Rest], Acc) ->
  conjoin_node(N - 3, Rest, [node(A, B, C) | Acc]).

digit_from_list([A]) ->
  digit(A);
digit_from_list([A, B]) ->
  digit(A, B);
digit_from_list([A, B, C]) ->
  digit(A, B, C);
digit_from_list([A, B, C, D]) ->
  digit(A, B, C, D).