From 29d5d6f170964769f094ee68f2fd0c358e369808 Mon Sep 17 00:00:00 2001
From: Ulf Wiger <ulf@wiger.net>
Date: Wed, 2 Jul 2025 14:36:25 +0200
Subject: [PATCH] Add mrdb_index:select() et al

---
 src/mrdb_index.erl                    | 36 +++++++++++++++++++++++++++
 src/mrdb_select.erl                   | 18 +++++++++++---
 test/mnesia_rocksdb_indexes_SUITE.erl | 10 ++++++++
 3 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/src/mrdb_index.erl b/src/mrdb_index.erl
index 2f901a7..d81413e 100644
--- a/src/mrdb_index.erl
+++ b/src/mrdb_index.erl
@@ -9,6 +9,10 @@
         , fold/4
         , rev_fold/4
         , index_ref/2
+        , select/3
+        , select/4
+        , select_reverse/3
+        , select_reverse/4
         ]).
 
 -record(mrdb_ix_iter, { i    :: mrdb:mrdb_iterator()
@@ -88,6 +92,22 @@ fold_(Tab, IxPos, Start, Dir, FoldFun, Acc) ->
               iter_fold(I, Start, Dir, FoldFun, Acc)
       end).
 
+select(Tab, Ix, MS) ->
+    select(Tab, Ix, MS, infinity).
+
+select(Tab, Ix, MS, Limit) ->
+    #{} = R = mrdb:ensure_ref(Tab),
+    #{} = IxR = ensure_index_ref(Ix, R),
+    mrdb_select:select(IxR#{derive_obj_f => mk_derive_obj_f(R)}, MS, Limit).
+
+select_reverse(Tab, Ix, MS) ->
+    select_reverse(Tab, Ix, MS, infinity).
+
+select_reverse(Tab, Ix, MS, Limit) ->
+    #{} = R = mrdb:ensure_ref(Tab),
+    #{} = IxR = ensure_index_ref(Ix, R),
+    mrdb_select:select_reverse(IxR#{derive_obj_f => mk_derive_obj_f(R)}, MS, Limit).
+
 iter_fold(I, Start, Dir, Fun, Acc) ->
     iter_fold_(iterator_move(I, Start), I, Dir, Fun, Acc).
 
@@ -121,6 +141,22 @@ iterator_move_bag(#mrdb_ix_iter{i = I, sub = Sub}, Dir) ->
             Other
     end.
 
+%% Used for mrdb_select:select()
+%% The select operation folds over index keys, and for matching keys,
+%% calls the `derive_obj_f/1` fun, which normally just does `Obj -> [Obj]`.
+%% In the case of indexes, it fetches the object, if it exists, and then
+%% returns `[{IndexKey, Object}]`, which is what the matchspec should operate on.
+%%
+mk_derive_obj_f(Ref) ->
+    fun({{IxK, Key}}) ->
+            case mrdb:read(Ref, Key, []) of
+                [Obj] ->
+                    [{IxK, Obj}];
+                [] ->
+                    []
+            end
+    end.
+
 -spec opt_read(mrdb:ref_or_tab(), Key :: any()) -> any().
 opt_read(R, Key) ->
     case mrdb:read(R, Key, []) of
diff --git a/src/mrdb_select.erl b/src/mrdb_select.erl
index 2415746..06dd5bf 100644
--- a/src/mrdb_select.erl
+++ b/src/mrdb_select.erl
@@ -30,6 +30,7 @@
              , limit
              , key_only = false                 % TODO: not used
              , direction = forward
+             , derive_obj_f = fun unit_l/1
              }).
 
 select(Ref, MS, Limit) when is_map(Ref), is_list(MS) ->
@@ -58,7 +59,9 @@ mk_sel(#{name := Tab} = Ref, MS, Dir, Limit) ->
          compiled_ms = ets:match_spec_compile(MS),
          key_only = needs_key_only(MS),
          direction = Dir,
-         limit = Limit}.
+         limit = Limit,
+         derive_obj_f = derive_f(Ref)}.
+
 
 select(Cont) ->
     case Cont of
@@ -266,8 +269,9 @@ select_traverse({ok, K, V}, Limit, Pfx, MS, I, #sel{ref = R, direction = Dir} =
     case is_prefix(Pfx, K) of
         true ->
             DecKey = decode_key(K, R),
-            Rec = decode_val(V, DecKey, R),
-            case ets:match_spec_run([Rec], MS) of
+            Rec0 = decode_val(V, DecKey, R),
+            RecL = derive_object(Rec0, Sel),
+            case ets:match_spec_run(RecL, MS) of
                 [] ->
                     select_traverse(
                       rocksdb:iterator_move(I, next_or_prev(Dir)), Limit, Pfx, MS,
@@ -341,6 +345,9 @@ i_move(I, K, reverse) ->
 i_move(I, K, forward) ->
     rocksdb:iterator_move(I, K).
 
+derive_object(R, #sel{derive_obj_f = F}) ->
+    F(R).
+
 keypat([H|T], KeyPos, Ref) ->
     keypat(T, KeyPos, Ref, keypat_pfx(H, KeyPos, Ref)).
 
@@ -362,3 +369,8 @@ keypat_pfx({HeadPat,_Gs,_}, KeyPos, #{encoding := {sext,_}}) when is_tuple(HeadP
 keypat_pfx(_, _, _) ->
     <<>>.
 
+derive_f(#{derive_obj_f := F}) when is_function(F, 1) -> F;
+derive_f(_) -> fun unit_l/1.
+
+unit_l(X) ->
+    [X].
diff --git a/test/mnesia_rocksdb_indexes_SUITE.erl b/test/mnesia_rocksdb_indexes_SUITE.erl
index b2197a4..17ee458 100644
--- a/test/mnesia_rocksdb_indexes_SUITE.erl
+++ b/test/mnesia_rocksdb_indexes_SUITE.erl
@@ -236,12 +236,22 @@ test_index_plugin(Config) ->
     if Type == rdb ->
             Res1 = lists:sort(mrdb:index_read(Tab,<<"sen">>, {pfx})),
             Res2 = lists:sort(mrdb:index_read(Tab,<<"whi">>, {pfx})),
+            ok = test_select(Tab,{pfx},[{'_', [], ['$_']}]),
             [{Tab,"foobar","sentence"}] = mrdb:index_read(
                                             Tab, <<"foo">>, {pfx});
        true ->
             ok
     end.
 
+test_select(Tab, Ix, MS) ->
+    Res = mrdb_index:select(Tab, Ix, MS),
+    ct:log("mrdb_index:select(~p, ~p, ~p) -> ~p", [Tab, Ix, MS, Res]),
+    RevRes = mrdb_index:select_reverse(Tab, Ix, MS),
+    ct:log("mrdb_index:select_reverse(~p, ~p, ~p) -> ~p", [Tab, Ix, MS, RevRes]),
+    {Res,Res} = {Res, lists:reverse(RevRes)},
+    ok.
+
+
 ixtype(T) when T==bag;
                T==ordered ->
     {{pfx}, T};