improve graph

2 files changed, 111 insertions, 26 deletions

diff --git a/tests/modules/graph/test.lua b/tests/modules/graph/test.lua
index fe1d63e8d..7c17a4fec 100644
--- a/tests/modules/graph/test.lua
+++ b/tests/modules/graph/test.lua
@@ -170,6 +170,48 @@ function test_paritail_topo_sort_dynamic(t)
     end
 end
 
+function test_remove_edge_and_vertex(t)
+    local gh = graph.new(true)
+    gh:add_edge("a", "b")
+    gh:add_edge("b", "c")
+    gh:add_edge("c", "d")
+    gh:add_edge("a", "d")
+
+    t:require(gh:has_edge("a", "b"))
+    gh:remove_edge("a", "b")
+    t:require(not gh:has_edge("a", "b"))
+    t:require(gh:has_edge("a", "d"))
+
+    gh:remove_vertex("c")
+    t:require(not gh:has_edge("b", "c"))
+    t:require(not gh:has_edge("c", "d"))
+    t:are_equal(#gh:vertices(), 3)
+    local order = gh:topo_sort()
+    t:require(#order == 3)
+
+    gh:add_edge("b", "a")
+    gh:add_edge("d", "b")
+    local _, has_cycle = gh:topo_sort()
+    t:require(has_cycle)
+end
+
+function test_clone_reverse_undirected(t)
+    local ug = graph.new(false)
+    ug:add_edge(1, 2)
+    ug:add_edge(2, 3)
+    ug:add_edge(3, 1)
+
+    local clone = ug:clone()
+    t:require(#clone:edges() == #ug:edges())
+    t:require(clone:has_edge(1, 2))
+    t:require(clone:has_edge(2, 1))
+
+    local rev = ug:reverse()
+    t:require(rev:has_edge(1, 2))
+    t:require(rev:has_edge(2, 1))
+    t:require(#rev:edges() == #ug:edges())
+end
+
 function test_find_cycle(t)
     local edges = {
         {9, 1},
diff --git a/xmake/core/base/graph.lua b/xmake/core/base/graph.lua
index 64dfa3b31..6c86d1486 100644
--- a/xmake/core/base/graph.lua
+++ b/xmake/core/base/graph.lua
@@ -116,20 +116,18 @@ function graph:remove_vertex(v)
         end
     end)
     if contains then
-        self._edges_map[v] = nil
-        self._adjacent_edges[v] = nil
-        -- remove the adjacent edge with this vertex in the other vertices
-        for _, w in ipairs(self:vertices()) do
-            local edges = self:adjacent_edges(w)
-            if edges then
-                table.remove_if(edges, function (_, e)
-                    if e:other(w) == v then
-                        self._edges_map[w] = nil
-                        return true
-                    end
-                end)
+        -- remove all touching edges
+        local touching = {}
+        for _, e in ipairs(self._edges) do
+            if e:from() == v or e:to() == v then
+                table.insert(touching, {e:from(), e:to()})
             end
         end
+        for _, pair in ipairs(touching) do
+            self:remove_edge(pair[1], pair[2])
+        end
+        self._edges_map[v] = nil
+        self._adjacent_edges[v] = nil
 
         -- reset partial topological sort state since graph structure changed
         self._partial_topo_dirty = true
@@ -398,6 +396,61 @@ function graph:add_edge(from, to)
     self._partial_topo_dirty = true
 end
 
+-- remove edge
+function graph:remove_edge(from, to)
+    if not self:has_edge(from, to) then
+        return
+    end
+    local directed = self:is_directed()
+    local edges = self._edges
+    local target_index
+    local target_edge
+    for idx, e in ipairs(edges) do
+        local match = (e:from() == from and e:to() == to)
+        if not directed then
+            match = match or (e:from() == to and e:to() == from)
+        end
+        if match then
+            target_index = idx
+            target_edge = e
+            break
+        end
+    end
+    if not target_edge then
+        return
+    end
+
+    table.remove(edges, target_index)
+
+    local function remove_from_adj(vertex)
+        local adj = self._adjacent_edges[vertex]
+        if adj then
+            table.remove_if(adj, function (_, e) return e == target_edge end)
+        end
+    end
+
+    remove_from_adj(target_edge:from())
+    remove_from_adj(target_edge:to())
+
+    local edges_map = self._edges_map
+    local function clear_map(u, v)
+        local map = edges_map[u]
+        if map then
+            map[v] = nil
+            if not next(map) then
+                edges_map[u] = nil
+            end
+        end
+    end
+
+    clear_map(target_edge:from(), target_edge:to())
+    if not directed then
+        clear_map(target_edge:to(), target_edge:from())
+    end
+
+    self._partial_topo_dirty = true
+end
+
 -- has the given edge?
 function graph:has_edge(from, to)
     local edges_map = self._edges_map
@@ -413,13 +466,8 @@ end
 -- clone graph
 function graph:clone()
     local gh = graph.new(self:is_directed())
-    for _, v in ipairs(self:vertices()) do
-        local edges = self:adjacent_edges(v)
-        if edges then
-            for _, e in ipairs(edges) do
-                gh:add_edge(e:from(), e:to())
-            end
-        end
+    for _, e in ipairs(self._edges) do
+        gh:add_edge(e:from(), e:to())
     end
     return gh
 end
@@ -430,13 +478,8 @@ function graph:reverse()
         return self:clone()
     end
     local gh = graph.new(self:is_directed())
-    for _, v in ipairs(self:vertices()) do
-        local edges = self:adjacent_edges(v)
-        if edges then
-            for _, e in ipairs(edges) do
-                gh:add_edge(e:to(), e:from())
-            end
-        end
+    for _, e in ipairs(self._edges) do
+        gh:add_edge(e:to(), e:from())
     end
     return gh
 end