From 3e131907fcea9a13da872ba402f67248ffe5fcf4 Mon Sep 17 00:00:00 2001
From: Matthew Dailis <matthew.l.dailis@jpl.nasa.gov>
Date: Thu, 18 Jul 2024 18:16:26 -0700
Subject: [PATCH] Implement topological sort for spans

---
 .../PostgresResultsCellRepository.java        | 44 ++++++++++++++++---
 1 file changed, 37 insertions(+), 7 deletions(-)

diff --git a/merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/remotes/postgres/PostgresResultsCellRepository.java b/merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/remotes/postgres/PostgresResultsCellRepository.java
index 75aefa1247..3bb18f400c 100644
--- a/merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/remotes/postgres/PostgresResultsCellRepository.java
+++ b/merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/remotes/postgres/PostgresResultsCellRepository.java
@@ -31,12 +31,14 @@
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.time.Instant;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.SortedMap;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 public final class PostgresResultsCellRepository implements ResultsCellRepository {
@@ -424,13 +426,7 @@ private static void postActivities(
       // Sorts the map by SpanRecord parent ID to ensure foreign key constraints are met.
       // Entries with null parent IDs are placed first to avoid foreign key violations
       // for the "span_has_parent_span" constraint.
-      final var sortedAllActivityRecords = new LinkedHashMap<Long,SpanRecord>();
-      sortedAllActivityRecords.putAll(allActivityRecords.entrySet().stream()
-                                                        .filter(entry -> entry.getValue().parentId().isEmpty())
-                                                        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
-      sortedAllActivityRecords.putAll(allActivityRecords.entrySet().stream()
-                                                        .filter(entry -> !entry.getValue().parentId().isEmpty())
-                                                        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
+      final var sortedAllActivityRecords = topoSort(allActivityRecords, $ -> $.parentId().stream().toList());
 
       postActivitiesAction.apply(
           datasetId,
@@ -439,6 +435,40 @@ private static void postActivities(
     }
   }
 
+  /**
+   * Take an unsorted map representing a directed acyclic graph and produce a sorted LinkedHashMap where nodes always
+   * come after their dependencies. The edges must represent an acyclic graph.
+   * @param nodes a map from keys to values - the keys are used to define dependencies
+   * @param dependencies - for a given value, what are the keys of its dependencies?
+   * @return a sorted LinkedHashMap where nodes always come after their dependencies
+   */
+  private static <K, V> LinkedHashMap<K, V> topoSort(Map<K, V> nodes, Function<V, List<K>> dependencies) {
+    final var worklist = new ArrayList<>(nodes.keySet());
+    final var output = new LinkedHashMap<K, V>();
+    while (!worklist.isEmpty()) {
+      var cycleDetected = true;
+      for (int i = worklist.size() - 1; i >= 0; i--) {
+        final var key = worklist.get(i);
+        final var node = nodes.get(key);
+        // A node is ready to be added to the output if all of its dependencies are already in the output
+        if (dependencies.apply(node).stream().allMatch(output::containsKey)) {
+          output.put(key, node);
+          worklist.remove(i);
+          cycleDetected = false;
+        }
+      }
+      // If no nodes were added to the output in this round, there must be a cycle in the remaining nodes
+      if (cycleDetected) {
+        final var cycle = new LinkedHashMap<>();
+        for (final var key : worklist) {
+          cycle.put(key, nodes.get(key));
+        }
+        throw new IllegalArgumentException("Cycle detected in input to topoSort:" + cycle);
+      }
+    }
+    return output;
+  }
+
   private static SpanRecord simulatedActivityToRecord(final SimulatedActivity activity) {
     return new SpanRecord(
         activity.type(),