From 0d395460809ce5c7cdc3b7a400124097201d0b5c Mon Sep 17 00:00:00 2001
From: Andrew Mayorov <encube.ul@gmail.com>
Date: Mon, 20 Mar 2023 14:24:52 +0300
Subject: [PATCH] feat(wdgraph): add `fold/3` which folds over graph edges

---
 apps/emqx/src/emqx_wdgraph.erl        | 24 +++++++++++++++++++++++-
 apps/emqx/test/emqx_wdgraph_tests.erl | 20 ++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/apps/emqx/src/emqx_wdgraph.erl b/apps/emqx/src/emqx_wdgraph.erl
index 3361c52d1..bd7f58e7c 100644
--- a/apps/emqx/src/emqx_wdgraph.erl
+++ b/apps/emqx/src/emqx_wdgraph.erl
@@ -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
diff --git a/apps/emqx/test/emqx_wdgraph_tests.erl b/apps/emqx/test/emqx_wdgraph_tests.erl
index ece87b966..c159a0f49 100644
--- a/apps/emqx/test/emqx_wdgraph_tests.erl
+++ b/apps/emqx/test/emqx_wdgraph_tests.erl
@@ -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),