improve graph

1 files changed, 34 insertions, 16 deletions

diff --git a/xmake/core/base/graph.lua b/xmake/core/base/graph.lua
index 80d4566ec..1fb6b01dd 100644
--- a/xmake/core/base/graph.lua
+++ b/xmake/core/base/graph.lua
@@ -24,11 +24,11 @@ local object = require("base/object")
 
 -- define module
 local graph = graph or object { _init = {"_directed"} } {true}
-local edge = edge or object { _init = {"_from", "_to", "_weight"} }
+local edge = edge or object { _init = {"_from", "_to"} }
 
 -- new edge, from -> to
-function edge.new(from, to, weight)
-    return edge {from, to, weight or 1.0}
+function edge.new(from, to)
+    return edge {from, to}
 end
 
 function edge:from()
@@ -47,8 +47,8 @@ function edge:other(v)
     end
 end
 
-function edge:weight()
-    return self._weight
+function edge:__tostring()
+    return string.format("<edge:%s-%s>", self:from(), self:to())
 end
 
 -- clear graph
@@ -56,6 +56,7 @@ function graph:clear()
     self._vertices = {}
     self._edges = {}
     self._adjacent_edges = {}
+    self._edges_map = {}
 end
 
 -- is empty?
@@ -98,13 +99,19 @@ 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
         if not self:is_directed() then
             for _, w in ipairs(self:vertices()) do
                 local edges = self:adjacent_edges(w)
                 if edges then
-                    table.remove_if(edges, function (_, e) return e:other(w) == v end)
+                    table.remove_if(edges, function (_, e)
+                        if e:other(w) == v then
+                            self._edges_map[w] = nil
+                            return true
+                        end
+                    end)
                 end
             end
         end
@@ -189,8 +196,8 @@ function graph:edges()
 end
 
 -- add edge
-function graph:add_edge(from, to, weight)
-    local e = edge.new(from, to, weight)
+function graph:add_edge(from, to)
+    local e = edge.new(from, to)
     if not self:has_vertex(from) then
         table.insert(self._vertices, from)
         self._adjacent_edges[from] = {}
@@ -199,11 +206,16 @@ function graph:add_edge(from, to, weight)
         table.insert(self._vertices, to)
         self._adjacent_edges[to] = {}
     end
+    local edges_map = self._edges_map
+    edges_map[from] = edges_map[from] or {}
+    edges_map[from][to] = true
     if self:is_directed() then
-        table.insert(self._adjacent_edges[e:from()], e)
+        table.insert(self._adjacent_edges[from], e)
     else
-        table.insert(self._adjacent_edges[e:from()], e)
-        table.insert(self._adjacent_edges[e:to()], e)
+        table.insert(self._adjacent_edges[from], e)
+        table.insert(self._adjacent_edges[to], e)
+        edges_map[to] = edges_map[to] or {}
+        edges_map[to][from] = true
     end
     table.insert(self._edges, e)
 end
@@ -212,9 +224,15 @@ end
 function graph:has_edge(from, to)
     local edges = self:adjacent_edges(from)
     if edges then
-        for _, e in ipairs(edges) do
-            if e:to() == to then
-                return true
+        local edges_map = self._edges_map
+        local from_map = edges_map[from]
+        if from_map and from_map[to] then
+            return true
+        else
+            for _, e in ipairs(edges) do
+                if e:to() == to then
+                    return true
+                end
             end
         end
     end
@@ -228,7 +246,7 @@ function graph:clone()
         local edges = self:adjacent_edges(v)
         if edges then
             for _, e in ipairs(edges) do
-                gh:add_edge(e:from(), e:to(), e:weight())
+                gh:add_edge(e:from(), e:to())
             end
         end
     end
@@ -245,7 +263,7 @@ function graph:reverse()
         local edges = self:adjacent_edges(v)
         if edges then
             for _, e in ipairs(edges) do
-                gh:add_edge(e:to(), e:from(), e:weight())
+                gh:add_edge(e:to(), e:from())
             end
         end
     end