From 1cfb869a8c0b0a430b9ee8fd2b8830530f52fbf9 Mon Sep 17 00:00:00 2001
From: Jon Johnson <jon.johnson@chainguard.dev>
Date: Mon, 5 Jun 2023 12:26:05 -0700
Subject: [PATCH] Speed up StableTopologicalSort

The biggest savings come from just calling sort less often. Instead of
in the inner-most loop, we call it after processing the whole map of
predecessors.

The other portion is just not appending duplicate items to the queue. I
do this in a simple way by maintaining a set of queued things (similar
to the visited set).

I think the "proper" way to fix this would be to use a sorted list for
queue instead of sorting after every pass, but I'm too lazy to implement
that, and I don't think it's part of the go stdlib.

For our use case, on my machine, StableTopologicalSort took:

Before: 230.00s
After:    0.03s

Signed-off-by: Jon Johnson <jon.johnson@chainguard.dev>
---
 dag.go | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/dag.go b/dag.go
index b83f805c..3946da71 100644
--- a/dag.go
+++ b/dag.go
@@ -84,10 +84,12 @@ func StableTopologicalSort[K comparable, T any](g Graph[K, T], less func(K, K) b
 	}
 
 	queue := make([]K, 0)
+	queued := make(map[K]struct{})
 
 	for vertex, predecessors := range predecessorMap {
 		if len(predecessors) == 0 {
 			queue = append(queue, vertex)
+			queued[vertex] = struct{}{}
 		}
 	}
 
@@ -112,14 +114,21 @@ func StableTopologicalSort[K comparable, T any](g Graph[K, T], less func(K, K) b
 		for vertex, predecessors := range predecessorMap {
 			delete(predecessors, currentVertex)
 
-			if len(predecessors) == 0 {
-				queue = append(queue, vertex)
+			if len(predecessors) != 0 {
+				continue
+			}
 
-				sort.Slice(queue, func(i, j int) bool {
-					return less(queue[i], queue[j])
-				})
+			if _, ok := queued[vertex]; ok {
+				continue
 			}
+
+			queue = append(queue, vertex)
+			queued[vertex] = struct{}{}
 		}
+
+		sort.Slice(queue, func(i, j int) bool {
+			return less(queue[i], queue[j])
+		})
 	}
 
 	gOrder, err := g.Order()