memstore

%% Copyright (c) 2017-present. Benoit Chesneau
%%
%% This source code is licensed under the MIT license found in the
%% LICENSE file in the root directory of this source tree.

-module(memstore_iterator).
-author("benoitc").

%% API
-export([
  iterator_move/2,
  iterator_close/1
]).

-export([
  lookup/3,
  lookup_last/2,
  start_link/2
]).

-export([
  init/1,
  handle_call/3,
  handle_cast/2,
  handle_info/2,
  terminate/2,
  code_change/3
]).

-include_lib("stdlib/include/ms_transform.hrl").


lookup(Name, Key, Seq) ->
  MS = ets:fun2ms(
    fun({{K, S}, Res}) when K =:= Key, S =< Seq->
      Res
    end
  ),
  case ets:select_reverse(Name, MS, 1) of
    '$end_of_table' -> not_found;
    {[], _} -> not_found;
    {[deleted], _} -> not_found;
    {[{ok, Val}], _} -> {ok, Val}
  end.

lookup_last(Name, Key) ->
  MS = ets:fun2ms(
    fun({{K, _S}, Res}) when K =:= Key -> Res end
  ),
  case ets:select_reverse(Name, MS, 1) of
    '$end_of_table' -> not_found;
    {[], _} -> not_found;
    {[deleted], _} -> not_found;
    {[{ok, Val}], _} -> {ok, Val}
  end.

iterator_move(Itr, Action) ->
  gen_server:call(Itr, {move, Action}).

iterator_close(Itr) ->
  gen_server:call(Itr, close).

start_link(Name, Options) ->
  gen_server:start_link(?MODULE, [Name, Options], []).

init([Name, Options]) ->
  State = init_iterator(Name, Options),
  {ok, State}.

init_iterator(Name, Options) ->
  {Seq, Sp} = case proplists:get_value(snapshot, Options) of
                #{ seq := S } -> {S, undefined};
                undefined ->
                  #{ seq := S} = Snapshot = memstore:new_snapshot(Name),
                  {S, Snapshot}
              end,
  Tab = ets:lookup_element(Name, tab, 2),
  #{ name => Tab, seq => Seq, snapshot => Sp }.

handle_call({move, Action}, _From, State) ->
  {Result, NewState} = handle_action(Action, State),
  {reply, Result, NewState};

handle_call(close, _From, State) ->
  {stop, normal, ok, State};

handle_call(_Msg, _From, State) ->
  {reply, badarg, State}.

handle_cast(_Msg, State) -> {noreply, State}.

handle_info(_Info, State) -> {noreply, State}.

terminate(_Reason, _Config) -> ok.

code_change(_Vsn, State, _Extra) -> {ok, State}.

handle_action(first, State = #{ name := Name, seq := Seq}) ->
  case seek_to_first(Seq, Name) of
    false -> {false, State};
    {ok, Key, _Val}=OK -> {{OK, State#{ next => Key}}}
  end;
handle_action(last, State = #{ name := Name, seq := Seq}) ->
  MS = ets:fun2ms(
    fun({{K, S}, Res}) when S =< Seq, Res /= deleted -> {K, Res} end
  ),
  case ets:select_reverse(Name, MS, 1) of
    '$end_of_table' ->
      false;
    {[{Key, {ok, Val}}], _} ->
      {{ok, Key, Val}, State#{ next => Key }}
  end;
handle_action(next, State= #{ next := Next, seq := Seq, name := Name } ) ->
  case next(Next, Seq, Name) of
    {ok, Key, _}=OK ->
      {OK, State#{next => Key}};
    false ->
      {false, State}
  end;
handle_action(next, State) ->
  handle_action(first, State);
handle_action(prev, State= #{ next := Next, seq := Seq, name := Name }) ->
  case prev(Next, Seq, Name) of
    {ok, Key, _}=OK ->
      {OK, State#{next => Key}};
    false ->
      {false, State}
  end;
handle_action(prev, State) ->
  {false, State};
handle_action({seek, Key}, State = #{ name := Name, seq := Seq }) ->
  case lookup(Name, Key, Seq) of
    {ok, Val} ->
      {{ok, Key, Val}, State#{ next => Key }};
    not_found ->
      handle_action(next, State#{ next => Key })
  end;
handle_action({seek_for_prev, Key}, State = #{ name := Name, seq := Seq }) ->
  case lookup(Name, Key, Seq) of
    {ok, Val} ->
      {{ok, Key, Val}, State#{ next => Key }};
    not_found ->
      handle_action(prev, State#{ next => Key })
  end;
handle_action(_, _) ->
  {error, badarg}.


seek_to_first(Seq, Tab) ->
  ets:safe_fixtable(Tab, true),
  MS = ets:fun2ms(
    fun({{K, S}, R}) when S =< Seq -> {K, R} end
  ),
  try forward(ets:select(Tab, MS, 1), Tab)
  after ets:safe_fixtable(Tab, false)
  end.


next(Key, Seq, Tab) ->
  ets:safe_fixtable(Tab, true),
  MS = ets:fun2ms(
    fun({{K, S}, R}) when K > Key, S =< Seq -> {K, R} end
  ),
  try forward(ets:select(Tab, MS, 1), Tab)
  after ets:safe_fixtable(Tab, false)
  end.

prev(Key, Seq, Tab ) ->
  MS = ets:fun2ms(
    fun({{K, S}, Res}) when K < Key, S =< Seq -> {K, Res} end
  ),
  ets:safe_fixtable(Tab, true),
  try reverse(ets:select_reverse(Tab, MS, 1), Tab)
  after ets:safe_fixtable(Tab, false)
  end.

forward('$end_of_table', _) -> false;
forward({[Prev], _}=Res, Tab) ->
  forward(Res, Prev, Tab).

forward('$end_of_table', {_, deleted}, _) ->
  false;
forward({[{Key, _}=Next], Cont}, {Key, _}, Tab) ->
  forward(ets:select(Cont), Next, Tab);
forward({[{_, _}=Next], Cont}, {_, deleted}, Tab) ->
  forward(ets:select(Cont), Next, Tab);
forward(_, {Key, {ok, Val}}, _) ->
  {ok, Key, Val}.

reverse('$end_of_table', _) -> false;
reverse({[{Key, {ok, Val}}], _}, _Tab) ->
  {ok, Key, Val};
reverse({[{_, deleted}=Prev], _}=Res, Tab) ->
  reverse(Res, Prev, Tab).

reverse('$end_of_table', _, _) ->
  false;
reverse({[{Key, _}=Prev], Cont}, {Key, _}, Tab) ->
  reverse(ets:select(Cont), Prev, Tab);
reverse({[{_, deleted}=Prev], Cont}, _, Tab) ->
  reverse(ets:select(Cont), Prev, Tab);
reverse({[{Key, {ok, Val}}], _}, _, _) ->
  {ok, Key, Val}.