Use DependencyGraph in integration sorting

This updates the SmithyIntegration sorting to use the generic
DependencyGraph. It also adds some benchmarking configurations,
which show a minor speed bump for highly dependent integrations
and a major speed bump for highly independent integrations.

Author: JordonPhillips

smithy-codegen-core/build.gradle.kts

@@ -4,6 +4,7 @@
  */
 plugins {
     id("smithy.module-conventions")
+    id("smithy.profiling-conventions")
 }
 
 description = "This module provides a code generation framework for generating clients, " +
@@ -14,6 +15,7 @@ extra["moduleName"] = "software.amazon.smithy.codegen.core"
 
 dependencies {
     api(project(":smithy-utils"))
+    jmh(project(":smithy-utils"))
     api(project(":smithy-model"))
     api(project(":smithy-build"))
 }

smithy-codegen-core/src/jmh/java/software/amazon/smithy/codegen/core/jmh/SmithyIntegrations.java

@@ -1,4 +1,225 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
 package software.amazon.smithy.codegen.core.jmh;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import software.amazon.smithy.codegen.core.CodegenContext;
+import software.amazon.smithy.codegen.core.ImportContainer;
+import software.amazon.smithy.codegen.core.SmithyIntegration;
+import software.amazon.smithy.codegen.core.Symbol;
+import software.amazon.smithy.codegen.core.SymbolWriter;
+import software.amazon.smithy.utils.DependencyGraph;
+
+@Warmup(iterations = 3)
+@Measurement(iterations = 3, timeUnit = TimeUnit.MICROSECONDS)
+@BenchmarkMode(Mode.AverageTime)
+@Fork(1)
 public class SmithyIntegrations {
+
+    @State(Scope.Thread)
+    public static class SmithyIntegrationsState {
+        public Map<String, SmithyIntegration<?, ?, ?>> independentIntegrations10;
+        public Map<String, SmithyIntegration<?, ?, ?>> dependentIntegrations10;
+        public Map<String, SmithyIntegration<?, ?, ?>> independentIntegrations100;
+        public Map<String, SmithyIntegration<?, ?, ?>> dependentIntegrations100;
+        public Map<String, SmithyIntegration<?, ?, ?>> independentIntegrations1000;
+        public Map<String, SmithyIntegration<?, ?, ?>> dependentIntegrations1000;
+
+        @Setup
+        public void setup() {
+            independentIntegrations10 = new LinkedHashMap<>(10);
+            independentIntegrations100 = new LinkedHashMap<>(100);
+            independentIntegrations1000 = new LinkedHashMap<>(1000);
+            for (int i = 0; i < 1000; i++) {
+                String name = "integration" + i;
+                TestIntegration integration = new TestIntegration(name);
+                if (i < 10) {
+                    independentIntegrations10.put(name, integration);
+                }
+                if (i < 100) {
+                    independentIntegrations100.put(name, integration);
+                }
+                independentIntegrations1000.put(name, integration);
+            }
+            independentIntegrations100 = Collections.unmodifiableMap(independentIntegrations100);
+
+            dependentIntegrations10 = new LinkedHashMap<>(10);
+            dependentIntegrations100 = new LinkedHashMap<>(100);
+            dependentIntegrations1000 = new LinkedHashMap<>(1000);
+            for (int i = 0; i < 1000; i++) {
+                String name = "integration" + i;
+                String next = "integration" + (i + 1);
+                String afterNext = "integration" + (i + 2);
+
+
+                if (i < 10) {
+                    List<String> dependencies = new ArrayList<>();
+                    if (i < 9) {
+                        dependencies.add(next);
+                    }
+                    if (i < 8) {
+                        dependencies.add(afterNext);
+                    }
+                    dependentIntegrations10.put(name, new TestIntegration(name, (byte) 0, Collections.emptyList(), dependencies));
+                }
+
+                if (i < 100) {
+                    List<String> dependencies = new ArrayList<>();
+                    if (i < 99)  {
+                        dependencies.add(next);
+                    }
+                    if (i < 98) {
+                        dependencies.add(afterNext);
+                    }
+                    dependentIntegrations100.put(name, new TestIntegration(name, (byte) 0, Collections.emptyList(), dependencies));
+                }
+
+                List<String> dependencies = new ArrayList<>();
+                if (i < 999) {
+                    dependencies.add(next);
+                }
+                if (i < 998) {
+                    dependencies.add(afterNext);
+                }
+                dependentIntegrations1000.put(name,
+                        new TestIntegration(name, (byte) 0, Collections.emptyList(), dependencies));
+            }
+        }
+    }
+
+    @Benchmark
+    public List<SmithyIntegration<?, ?, ?>> sortIndependentIntegrations10(SmithyIntegrationsState state) {
+        return SmithyIntegration.sort(state.independentIntegrations10.values());
+    }
+
+    @Benchmark
+    public List<SmithyIntegration<?, ?, ?>> sortIndependentIntegrations100(SmithyIntegrationsState state) {
+        return SmithyIntegration.sort(state.independentIntegrations100.values());
+    }
+
+    @Benchmark
+    public List<SmithyIntegration<?, ?, ?>> sortIndependentIntegrations1000(SmithyIntegrationsState state) {
+        return SmithyIntegration.sort(state.independentIntegrations1000.values());
+    }
+
+    @Benchmark
+    public List<SmithyIntegration<?, ?, ?>> sortDependentIntegrations10(SmithyIntegrationsState state) {
+        return SmithyIntegration.sort(state.dependentIntegrations10.values());
+    }
+
+    @Benchmark
+    public List<SmithyIntegration<?, ?, ?>> sortDependentIntegrations100(SmithyIntegrationsState state) {
+        return SmithyIntegration.sort(state.dependentIntegrations100.values());
+    }
+
+    @Benchmark
+    public List<SmithyIntegration<?, ?, ?>> sortDependentIntegrations1000(SmithyIntegrationsState state) {
+        return SmithyIntegration.sort(state.dependentIntegrations1000.values());
+    }
+
+    private static class IntegrationComparator implements Comparator<String> {
+
+        private final Map<String, SmithyIntegration<?, ?, ?>> lookup;
+
+        IntegrationComparator(Map<String, SmithyIntegration<?, ?, ?>> lookup) {
+            this.lookup = lookup;
+        }
+
+        @Override
+        public int compare(String o1, String o2) {
+            SmithyIntegration<?, ?, ?> left = lookup.get(o1);
+            SmithyIntegration<?, ?, ?> right = lookup.get(o2);
+            if (left == null || right == null) {
+                return 0;
+            }
+            return Byte.compare(left.priority(), right.priority());
+        }
+    }
+
+    private static class TestImportContainer implements ImportContainer {
+        @Override
+        public void importSymbol(Symbol symbol, String alias) {
+
+        }
+    }
+
+    private static class TestSymbolWriter extends SymbolWriter<TestSymbolWriter, TestImportContainer> {
+        public TestSymbolWriter(TestImportContainer importContainer) {
+            super(importContainer);
+        }
+    }
+
+    private static class TestSettings {}
+
+    private static class TestIntegration implements SmithyIntegration<
+            TestSettings,
+            TestSymbolWriter,
+            CodegenContext<TestSettings, TestSymbolWriter, TestIntegration>> {
+
+        private final String name;
+        private final byte priority;
+        private final List<String> runBefore;
+        private final List<String> runAfter;
+
+        TestIntegration(String name) {
+            this(name, (byte) 0);
+        }
+
+        TestIntegration(String name, byte priority) {
+            this(name, priority, Collections.emptyList());
+        }
+
+        TestIntegration(String name, byte priority, List<String> runBefore) {
+            this(name, priority, runBefore, Collections.emptyList());
+        }
+
+        TestIntegration(String name, byte priority, List<String> runBefore, List<String> runAfter) {
+            this.name = name;
+            this.priority = priority;
+            this.runBefore = runBefore;
+            this.runAfter = runAfter;
+        }
+
+        @Override
+        public String name() {
+            return name;
+        }
+
+        @Override
+        public byte priority() {
+            return priority;
+        }
+
+        @Override
+        public List<String> runBefore() {
+            return runBefore;
+        }
+
+        @Override
+        public List<String> runAfter() {
+            return runAfter;
+        }
+
+        @Override
+        public String toString() {
+            return name();
+        }
+    }
 }

smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/IntegrationTopologicalSort.java

@@ -5,65 +5,37 @@
 package software.amazon.smithy.codegen.core;
 
 import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
 import java.util.LinkedHashMap;
-import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.PriorityQueue;
-import java.util.Queue;
-import java.util.Set;
 import java.util.logging.Logger;
+import software.amazon.smithy.utils.DependencyGraph;
 
 final class IntegrationTopologicalSort<I extends SmithyIntegration<?, ?, ?>> {
 
     private static final Logger LOGGER = Logger.getLogger(IntegrationTopologicalSort.class.getName());
 
     private final Map<String, I> integrationLookup = new LinkedHashMap<>();
-    private final Map<String, Integer> insertionOrder = new HashMap<>();
-    private final Map<String, Set<String>> forwardDependencies = new LinkedHashMap<>();
-    private final Map<String, Set<String>> reverseDependencies = new LinkedHashMap<>();
-
-    private final Queue<String> satisfied = new PriorityQueue<>((left, right) -> {
-        I leftIntegration = integrationLookup.get(left);
-        I rightIntegration = integrationLookup.get(right);
-        // Priority order is used to sort first.
-        int byteResult = Byte.compare(rightIntegration.priority(), leftIntegration.priority());
-        // If priority is a tie, then sort based on insertion order of integrations.
-        // This makes the order deterministic.
-        return byteResult == 0
-                ? Integer.compare(insertionOrder.get(left), insertionOrder.get(right))
-                : byteResult;
-    });
+    private final DependencyGraph<String> dependencyGraph = new DependencyGraph<>();
 
     IntegrationTopologicalSort(Iterable<I> integrations) {
-        // Validate name conflicts and register integrations with the lookup table + insertion order table.
         for (I integration : integrations) {
             addIntegration(integration);
         }
 
         // Validate missing dependencies and add found dependencies.
         for (I integration : integrations) {
             for (String before : getValidatedDependencies("before", integration.name(), integration.runBefore())) {
-                addDependency(before, integration.name());
+                dependencyGraph.addDependency(before, integration.name());
             }
             for (String after : getValidatedDependencies("after", integration.name(), integration.runAfter())) {
-                addDependency(integration.name(), after);
-            }
-        }
-
-        // Offer satisfied dependencies.
-        for (I integration : integrations) {
-            if (!forwardDependencies.containsKey(integration.name())) {
-                satisfied.offer(integration.name());
+                dependencyGraph.addDependency(integration.name(), after);
             }
         }
     }
 
     private void addIntegration(I integration) {
         I previous = this.integrationLookup.put(integration.name(), integration);
-        insertionOrder.put(integration.name(), insertionOrder.size());
         if (previous != null) {
             throw new IllegalArgumentException(String.format(
                     "Conflicting SmithyIntegration names detected for '%s': %s and %s",
@@ -91,33 +63,14 @@ private List<String> getValidatedDependencies(String descriptor, String what, Li
         }
     }
 
-    private void addDependency(String what, String dependsOn) {
-        forwardDependencies.computeIfAbsent(what, n -> new LinkedHashSet<>()).add(dependsOn);
-        reverseDependencies.computeIfAbsent(dependsOn, n -> new LinkedHashSet<>()).add(what);
-    }
-
     List<I> sort() {
-        List<I> result = new ArrayList<>();
-
-        while (!satisfied.isEmpty()) {
-            String current = satisfied.poll();
-            forwardDependencies.remove(current);
-            result.add(integrationLookup.get(current));
-
-            for (String dependent : reverseDependencies.getOrDefault(current, Collections.emptySet())) {
-                Set<String> dependentDependencies = forwardDependencies.get(dependent);
-                dependentDependencies.remove(current);
-                if (dependentDependencies.isEmpty()) {
-                    satisfied.offer(dependent);
-                }
-            }
+        List<I> result = new ArrayList<>(dependencyGraph.size());
+        List<String> sortedNames = dependencyGraph.toSortedList((left, right) -> {
+            return Byte.compare(integrationLookup.get(left).priority(), integrationLookup.get(right).priority());
+        });
+        for (String name : sortedNames) {
+            result.add(integrationLookup.get(name));
         }
-
-        if (!forwardDependencies.isEmpty()) {
-            throw new IllegalArgumentException("SmithyIntegration cycles detected among "
-                    + forwardDependencies.keySet());
-        }
-
         return result;
     }
 }