feat(wdgraph): add `fold/3` which folds over graph edges

This commit is contained in:
Andrew Mayorov 2023-03-20 14:24:52 +03:00 committed by Ilya Averyanov
parent 0c821cd3bd
commit 0d39546080
2 changed files with 43 additions and 1 deletions

View File

@ -27,6 +27,8 @@
-export([find_edge/3]).
-export([get_edges/2]).
-export([fold/3]).
-export([find_shortest_path/3]).
-export_type([t/0]).
@ -38,7 +40,7 @@
-type label() :: term().
-opaque t() :: t(gnode(), label()).
-opaque t(Node, Label) :: gb_trees:tree({Node}, {Node, weight(), Label}).
-opaque t(Node, Label) :: gb_trees:tree({Node}, [{Node, weight(), Label}]).
%%
@ -72,6 +74,26 @@ find_edge(From, To, G) ->
get_edges(Node, G) ->
tree_lookup({Node}, G, []).
-spec fold(FoldFun, Acc, t(Node, Label)) -> Acc when
FoldFun :: fun((Node, _Edge :: {Node, weight(), Label}, Acc) -> Acc).
fold(FoldFun, Acc, G) ->
fold_iterator(FoldFun, Acc, gb_trees:iterator(G)).
fold_iterator(FoldFun, AccIn, It) ->
case gb_trees:next(It) of
{{Node}, Edges = [_ | _], ItNext} ->
AccNext = lists:foldl(
fun(Edge = {_To, _Weight, _Label}, Acc) ->
FoldFun(Node, Edge, Acc)
end,
AccIn,
Edges
),
fold_iterator(FoldFun, AccNext, ItNext);
none ->
AccIn
end.
% Find the shortest path between two nodes, if any. If the path exists, return list
% of edge labels along that path.
% This is a Dijkstra shortest path algorithm. It is one-way right now, for

View File

@ -40,6 +40,26 @@ edges_nodes_test_() ->
?_assertEqual([{baz, 1, "cheapest"}, {foo, 0, "free"}], emqx_wdgraph:get_edges(bar, G5))
].
fold_test_() ->
G1 = emqx_wdgraph:new(),
G2 = emqx_wdgraph:insert_edge(foo, bar, 42, "fancy", G1),
G3 = emqx_wdgraph:insert_edge(bar, baz, 1, "cheapest", G2),
G4 = emqx_wdgraph:insert_edge(bar, foo, 0, "free", G3),
G5 = emqx_wdgraph:insert_edge(foo, bar, 100, "luxury", G4),
[
?_assertEqual(
% 100 + 0 + 1
101,
emqx_wdgraph:fold(fun(_From, {_, Weight, _}, Acc) -> Weight + Acc end, 0, G5)
),
?_assertEqual(
[bar, baz, foo],
lists:usort(
emqx_wdgraph:fold(fun(From, {To, _, _}, Acc) -> [From, To | Acc] end, [], G5)
)
)
].
nonexistent_nodes_path_test_() ->
G1 = emqx_wdgraph:new(),
G2 = emqx_wdgraph:insert_edge(foo, bar, 42, "fancy", G1),