fix edge index pruning

Author: zblanco

lib/graph.ex

@@ -1528,7 +1528,6 @@ defmodule Graph do
   defp prune_edge_index(
          %__MODULE__{
            multigraph: true,
-           edge_index: edge_index,
            edges: meta,
            partition_by: partition_by
          } = g,
@@ -1548,14 +1547,14 @@ defmodule Graph do
         v1_key = {v1_id, edge_p}
         v2_key = {v2_id, edge_p}
 
-        edge_index =
-          edge_index
+        updated_edge_index =
+          acc.edge_index
           |> Map.delete(v1_key)
           |> Map.delete(v2_key)
 
         %__MODULE__{
           acc
-          | edge_index: edge_index
+          | edge_index: updated_edge_index
         }
       end)
     end)
@@ -1564,7 +1563,6 @@ defmodule Graph do
   defp prune_edge_index(
          %__MODULE__{
            multigraph: true,
-           edge_index: edge_index,
            edges: meta,
            partition_by: partition_by
          } = g,
@@ -1584,25 +1582,47 @@ defmodule Graph do
 
     edge_partitions = partition_by.(edge)
 
+    edge_key = {v1_id, v2_id}
+
     Enum.reduce(edge_partitions, g, fn edge_p, acc ->
       partition =
-        edge_index
+        acc.edge_index
         |> Map.get(edge_p, %{})
-        |> Map.reject(fn {k, v} ->
-          (k == v1_id and MapSet.member?(v, {v1_id, v2_id})) or
-            (k == v2_id and MapSet.member?(v, {v1_id, v2_id}))
+        |> Enum.reduce(%{}, fn {k, v}, new_partition ->
+          cond do
+            k == v1_id ->
+              remaining = MapSet.delete(v, edge_key)
+
+              if MapSet.size(remaining) > 0 do
+                Map.put(new_partition, k, remaining)
+              else
+                new_partition
+              end
+
+            k == v2_id ->
+              remaining = MapSet.delete(v, edge_key)
+
+              if MapSet.size(remaining) > 0 do
+                Map.put(new_partition, k, remaining)
+              else
+                new_partition
+              end
+
+            true ->
+              Map.put(new_partition, k, v)
+          end
         end)
 
-      edge_index =
+      updated_edge_index =
         if not Enum.empty?(partition) do
-          Map.put(edge_index, edge_p, partition)
+          Map.put(acc.edge_index, edge_p, partition)
         else
-          Map.delete(edge_index, edge_p)
+          Map.delete(acc.edge_index, edge_p)
         end
 
       %__MODULE__{
         acc
-        | edge_index: edge_index
+        | edge_index: updated_edge_index
       }
     end)
   end

test/graph_test.exs

@@ -156,30 +156,27 @@ defmodule GraphTest do
       assert Map.has_key?(g.edge_index[:baz], g.vertex_identifier.(:a))
     end
 
-    # test "BFS traversal using multigraph partitions" do
-    #   graph =
-    #     Graph.new(multigraph: true)
-    #     |> Graph.add_edges([
-    #       {:a, :b, label: :foo},
-    #       {:a, :b, label: :bar},
-    #       {:b, :c, weight: 3},
-    #       {:b, :a, label: {:complex, :label}}
-    #     ])
-
-    #   assert Graph.bfs(graph, :a) == [:a, :b, :c]
-    # end
-
-    # test "DFS traversal using multigraph partitions" do
-    # end
-
-    # test "Dijkstra traversal using multigraph partitions" do
-    # end
-
-    # test "A* traversal using multigraph partitions" do
-    # end
-
-    # test "Bellman-Ford traversal using multigraph partitions" do
-    # end
+    test "update_labelled_edge preserves sibling edges in partition index" do
+      g =
+        Graph.new(multigraph: true)
+        |> Graph.add_edge(:fact, :join, label: :runnable)
+        |> Graph.add_edge(:fact, :step_a, label: :runnable)
+        |> Graph.add_edge(:fact, :step_b, label: :runnable)
+
+      assert length(Graph.edges(g, by: [:runnable])) == 3
+
+      g = Graph.update_labelled_edge(g, :fact, :join, :runnable, label: :ran)
+
+      ran_edges = Graph.edges(g, by: [:ran])
+      assert length(ran_edges) == 1
+      assert hd(ran_edges).v1 == :fact and hd(ran_edges).v2 == :join
+
+      runnable_edges = Graph.edges(g, by: [:runnable])
+      assert length(runnable_edges) == 2
+
+      runnable_targets = Enum.map(runnable_edges, & &1.v2) |> Enum.sort()
+      assert runnable_targets == [:step_a, :step_b]
+    end
   end
 
   describe "edge properties" do