Performance improvement for topological sort

1 files changed, 16 insertions, 7 deletions

diff --git a/src/main/java/org/testng/internal/Graph.java b/src/main/java/org/testng/internal/Graph.java
index df121a8f..d833610e 100644
--- a/src/main/java/org/testng/internal/Graph.java
+++ b/src/main/java/org/testng/internal/Graph.java
@@ -54,10 +54,7 @@ public class Graph<T> {
       node.addPredecessor(predecessor);
       addNeighbor(tm, predecessor);
       // Remove these two nodes from the independent list
-      if (null == m_independentNodes) {
-        m_independentNodes = Maps.newHashMap();
-        m_independentNodes.putAll(m_nodes);
-      }
+      initializeIndependentNodes();
       m_independentNodes.remove(predecessor);
       m_independentNodes.remove(tm);
       ppp("  REMOVED " + predecessor + " FROM INDEPENDENT OBJECTS");
@@ -103,9 +100,7 @@ public class Graph<T> {
   public void topologicalSort() {
     ppp("================ SORTING");
     m_strictlySortedNodes = Lists.newArrayList();
-    if (null == m_independentNodes) {
-      m_independentNodes = Maps.newHashMap();
-    }
+    initializeIndependentNodes();
 
     //
     // Clone the list of nodes but only keep those that are
@@ -159,6 +154,20 @@ public class Graph<T> {
     }
   }
 
+  private void initializeIndependentNodes() {
+    if (null == m_independentNodes) {
+      List<Node<T>> list = Lists.newArrayList(m_nodes.values());
+      // Ideally, we should not have to sort this. However, due to a bug where it treats all the methods as
+      // dependent nodes. Therefore, all the nodes were mostly sorted alphabetically. So we need to keep
+      // the behavior for now.
+      Collections.sort(list);
+      m_independentNodes = Maps.newLinkedHashMap();
+      for (Node<T> node : list) {
+        m_independentNodes.put(node.getObject(), node);
+      }
+    }
+  }
+
   private void dumpSortedNodes() {
     System.out.println("====== SORTED NODES");
     for (T n : m_strictlySortedNodes) {