WIP: Use a prefix tree for the reverse keep-while condition index

src/khepri_prefix_tree.erl

@@ -0,0 +1,105 @@
+-module(khepri_prefix_tree).
+
+-include("src/khepri_payload.hrl").
+
+-type children() :: #{khepri_path:node_id() => tree()}.
+
+-record(?MODULE, {children = #{},
+                  payload = ?NO_PAYLOAD}).
+
+-type tree(Payload) :: #?MODULE{children :: children(),
+                                payload :: Payload | ?NO_PAYLOAD}.
+
+-type tree() :: tree(term()).
+
+-export_type([tree/1]).
+
+-export([new/0,
+         fold_prefixes_of/4,
+         find_path/2,
+         update/3]).
+
+-spec new() -> tree().
+
+new() -> #?MODULE{}.
+
+-spec fold_prefixes_of(Fun, Acc, Path, Tree) -> Ret when
+      Fun :: fun((Payload, Acc) -> Acc1),
+      Acc :: term(),
+      Acc1 :: term(),
+      Path :: khepri_path:native_path(),
+      Tree :: khepri_prefix_tree:tree(Payload),
+      Payload :: term(),
+      Ret :: Acc1.
+
+fold_prefixes_of(Fun, Acc, Path, Tree) when is_function(Fun, 2) ->
+    do_fold_prefixes_of(Fun, Acc, Path, Tree).
+
+do_fold_prefixes_of(
+  Fun, Acc, [], #?MODULE{payload = Payload, children = Children}) ->
+    Acc1 = case Payload of
+               ?NO_PAYLOAD ->
+                   Acc;
+               _ ->
+                   Fun(Payload, Acc)
+           end,
+    maps:fold(
+      fun(_Component, Subtree, Acc2) ->
+              do_fold_prefixes_of(Fun, Acc2, [], Subtree)
+      end, Acc1, Children);
+do_fold_prefixes_of(
+  Fun, Acc, [Component | Rest], #?MODULE{children = Children}) ->
+    case maps:find(Component, Children) of
+        {ok, Subtree} ->
+            do_fold_prefixes_of(Fun, Acc, Rest, Subtree);
+        error ->
+            Acc
+    end.
+
+-spec find_path(Path, Tree) -> Ret when
+      Path :: khepri_path:native_path(),
+      Tree :: khepri_prefix_tree:tree(Payload),
+      Payload :: term(),
+      Ret :: {ok, Payload} | error.
+
+find_path(Path, Tree) ->
+    do_find_path(Path, Tree).
+
+do_find_path([], #?MODULE{payload = Payload}) ->
+    case Payload of
+        ?NO_PAYLOAD ->
+            error;
+        _ ->
+            {ok, Payload}
+    end;
+do_find_path([Component | Rest], #?MODULE{children = Children}) ->
+    case maps:find(Component, Children) of
+        {ok, Subtree} ->
+            do_find_path(Rest, Subtree);
+        error ->
+            error
+    end.
+
+-spec update(Fun, Path, Tree) -> Ret when
+      Fun :: fun((Payload | ?NO_PAYLOAD) -> Payload | ?NO_PAYLOAD),
+      Path :: khepri_path:native_path(),
+      Tree :: khepri_prefix_tree:tree(Payload),
+      Payload :: term(),
+      Ret :: khepri_prefix_tree:tree().
+
+update(Fun, Path, Tree) ->
+    do_update(Fun, Path, Tree).
+
+do_update(Fun, [], #?MODULE{payload = Payload} = Tree) ->
+    Tree#?MODULE{payload = Fun(Payload)};
+do_update(Fun, [Component | Rest], #?MODULE{children = Children} = Tree) ->
+    Subtree = maps:get(Component, Children, khepri_prefix_tree:new()),
+    Children1 = case do_update(Fun, Rest, Subtree) of
+                    #?MODULE{payload = ?NO_PAYLOAD, children = C}
+                      when C =:= #{} ->
+                        %% Drop unused branches.
+                        maps:remove(Component, Children);
+                    Subtree1 ->
+                        maps:put(Component, Subtree1, Children)
+                end,
+    Tree#?MODULE{children = Children1}.

src/khepri_tree.erl

@@ -40,8 +40,8 @@
                                   khepri_condition:native_keep_while()}.
 %% Per-node `keep_while' conditions.
 
--type keep_while_conds_revidx() :: #{khepri_path:native_path() =>
-                                     #{khepri_path:native_path() => ok}}.
+-type keep_while_conds_revidx() :: khepri_prefix_tree:tree(
+                                     #{khepri_path:native_path() => ok}).
 %% Internal reverse index of the keep_while conditions. If node A depends on a
 %% condition on node B, then this reverse index will have a "node B => node A"
 %% entry.
@@ -315,19 +315,25 @@ update_keep_while_conds_revidx(
     OldWatcheds = maps:get(Watcher, KeepWhileConds, #{}),
     KeepWhileCondsRevIdx1 = maps:fold(
                           fun(Watched, _, KWRevIdx) ->
-                                  Watchers = maps:get(Watched, KWRevIdx),
-                                  Watchers1 = maps:remove(Watcher, Watchers),
-                                  case maps:size(Watchers1) of
-                                      0 -> maps:remove(Watched, KWRevIdx);
-                                      _ -> KWRevIdx#{Watched => Watchers1}
-                                  end
+                                  khepri_prefix_tree:update(
+                                    fun(Watchers) ->
+                                            Watchers1 = maps:remove(
+                                                          Watcher, Watchers),
+                                            case maps:size(Watchers1) of
+                                                0 -> ?NO_PAYLOAD;
+                                                _ -> Watchers1
+                                            end
+                                    end, Watched, KWRevIdx)
                           end, KeepWhileCondsRevIdx, OldWatcheds),
     %% Then, record the watched paths.
     KeepWhileCondsRevIdx2 = maps:fold(
                           fun(Watched, _, KWRevIdx) ->
-                                  Watchers = maps:get(Watched, KWRevIdx, #{}),
-                                  Watchers1 = Watchers#{Watcher => ok},
-                                  KWRevIdx#{Watched => Watchers1}
+                                  khepri_prefix_tree:update(
+                                    fun (?NO_PAYLOAD) ->
+                                            #{Watcher => ok};
+                                        (Watchers) ->
+                                            Watchers#{Watcher => ok}
+                                    end, Watched, KWRevIdx)
                           end, KeepWhileCondsRevIdx1, KeepWhile),
     Tree#tree{keep_while_conds_revidx = KeepWhileCondsRevIdx2}.
 
@@ -1290,22 +1296,17 @@ eval_keep_while_conditions(
     maps:fold(
       fun
           (RemovedPath, delete, ToDelete) ->
-              maps:fold(
-                fun(Path, Watchers, ToDelete1) ->
-                        case lists:prefix(RemovedPath, Path) of
-                            true ->
-                                eval_keep_while_conditions_after_removal(
-                                  Tree, Watchers, ToDelete1);
-                            false ->
-                                ToDelete1
-                        end
-                end, ToDelete, KeepWhileCondsRevIdx);
+              khepri_prefix_tree:fold_prefixes_of(
+                fun(Watchers, ToDelete1) ->
+                        eval_keep_while_conditions_after_removal(
+                          Tree, Watchers, ToDelete1)
+                end, ToDelete, RemovedPath, KeepWhileCondsRevIdx);
           (UpdatedPath, NodeProps, ToDelete) ->
-              case KeepWhileCondsRevIdx of
-                  #{UpdatedPath := Watchers} ->
+              case khepri_prefix_tree:find_path(UpdatedPath, KeepWhileCondsRevIdx) of
+                  {ok, Watchers} ->
                       eval_keep_while_conditions_after_update(
                         Tree, UpdatedPath, NodeProps, Watchers, ToDelete);
-                  _ ->
+                  error ->
                       ToDelete
               end
       end, #{}, AppliedChanges).

src/khepri_tree.hrl

@@ -24,5 +24,5 @@
 
 -record(tree, {root = #node{} :: khepri_tree:tree_node(),
                keep_while_conds = #{} :: khepri_tree:keep_while_conds_map(),
-               keep_while_conds_revidx = #{} ::
+               keep_while_conds_revidx = khepri_prefix_tree:new() ::
                khepri_tree:keep_while_conds_revidx()}).

test/keep_while_conditions.erl

@@ -70,7 +70,6 @@ insert_when_keep_while_true_test() ->
     {S1, Ret, SE} = khepri_machine:apply(?META, Command, S0),
     Root = khepri_machine:get_root(S1),
     KeepWhileConds = khepri_machine:get_keep_while_conds(S1),
-    KeepWhileCondsRevIdx = khepri_machine:get_keep_while_conds_revidx(S1),
 
     ?assertEqual(
        #node{
@@ -90,9 +89,6 @@ insert_when_keep_while_true_test() ->
     ?assertEqual(
        #{[baz] => KeepWhile},
        KeepWhileConds),
-    ?assertEqual(
-       #{[foo] => #{[baz] => ok}},
-       KeepWhileCondsRevIdx),
     ?assertEqual({ok, #{[baz] => #{}}}, Ret),
     ?assertEqual([], SE).