1 files changed, 15 insertions, 4 deletions
diff --git a/init.lua b/init.lua
index 339345a..7ccec55 100644
--- a/init.lua
+++ b/init.lua
@@ -103,10 +103,11 @@ end
 
 --[[
 Topological Sort
-** This is not finished. OK for graphs with single root.
 ]]--
 function Graph:topsort()
 
+   local dummyRoot
+
    -- reverse the graph
    local rg,map = self:reverse()
    local rmap = {}
@@ -122,14 +123,24 @@ function Graph:topsort()
       error('Graph has cycles')
    end
 
-   -- run
-   for i,root in ipairs(rootnodes) do
-      root:dfs(function(node) table.insert(sortednodes,rmap[node]) end)
+   if #rootnodes > 1 then
+      dummyRoot = graph.Node('dummy_root')
+      for _, root in ipairs(rootnodes) do
+         dummyRoot:add(root)
+      end
+   else
+      dummyRoot = rootnodes[1]
    end
 
+   -- run
+   -- the trick is since the dummy node does not exist in original graph,
+   -- rmap[dummyRoot] = nil hence nothing gets inserted into the table
+   dummyRoot:dfs(function(node) table.insert(sortednodes,rmap[node]) end)
+
    if #sortednodes ~= #self.nodes then
       error('Graph has cycles')
    end
+
    return sortednodes,rg,rootnodes
 end