convert algorithms to SCC (#47866)

These places in the code can either be more efficient O(1) or more
correct using something more similar to the published SCC algorithm by
Tarjan for strongly connected components.

Author: vtjnash

src/gf.c

@@ -3377,6 +3377,7 @@ static jl_value_t *ml_matches(jl_methtable_t *mt,
                 }
             }
             // then we'll merge those numbers to assign each item in the group the same number
+            // (similar to Kosaraju's SCC algorithm?)
             uint32_t groupid = 0;
             uint32_t grouphi = 0;
             for (i = 0; i < len; i++) {

src/jitlayers.cpp

@@ -136,7 +136,11 @@ void jl_dump_llvm_opt_impl(void *s)
     **jl_ExecutionEngine->get_dump_llvm_opt_stream() = (JL_STREAM*)s;
 }
 
-static void jl_add_to_ee(orc::ThreadSafeModule &M, StringMap<orc::ThreadSafeModule*> &NewExports);
+static int jl_add_to_ee(
+        orc::ThreadSafeModule &M,
+        const StringMap<orc::ThreadSafeModule*> &NewExports,
+        DenseMap<orc::ThreadSafeModule*, int> &Queued,
+        std::vector<orc::ThreadSafeModule*> &Stack);
 static void jl_decorate_module(Module &M);
 static uint64_t getAddressForFunction(StringRef fname);
 
@@ -228,10 +232,13 @@ static jl_callptr_t _jl_compile_codeinst(
                 }
             }
         }
+        DenseMap<orc::ThreadSafeModule*, int> Queued;
+        std::vector<orc::ThreadSafeModule*> Stack;
         for (auto &def : emitted) {
             // Add the results to the execution engine now
             orc::ThreadSafeModule &M = std::get<0>(def.second);
-            jl_add_to_ee(M, NewExports);
+            jl_add_to_ee(M, NewExports, Queued, Stack);
+            assert(Queued.empty() && Stack.empty() && !M);
         }
         ++CompiledCodeinsts;
         MaxWorkqueueSize.updateMax(emitted.size());
@@ -1704,76 +1711,72 @@ static void jl_decorate_module(Module &M) {
 #endif
 }
 
+// Implements Tarjan's SCC (strongly connected components) algorithm, simplified to remove the count variable
 static int jl_add_to_ee(
         orc::ThreadSafeModule &M,
-        StringMap<orc::ThreadSafeModule*> &NewExports,
+        const StringMap<orc::ThreadSafeModule*> &NewExports,
         DenseMap<orc::ThreadSafeModule*, int> &Queued,
-        std::vector<std::vector<orc::ThreadSafeModule*>> &ToMerge,
-        int depth)
+        std::vector<orc::ThreadSafeModule*> &Stack)
 {
-    // DAG-sort (post-dominator) the compile to compute the minimum
-    // merge-module sets for linkage
+    // First check if the TSM is empty (already compiled)
     if (!M)
         return 0;
-    // First check and record if it's on the stack somewhere
+    // Next check and record if it is on the stack somewhere
     {
-        auto &Cycle = Queued[&M];
-        if (Cycle)
-            return Cycle;
-        ToMerge.push_back({});
-        Cycle = depth;
+        auto &Id = Queued[&M];
+        if (Id)
+            return Id;
+        Stack.push_back(&M);
+        Id = Stack.size();
     }
+    // Finally work out the SCC
+    int depth = Stack.size();
     int MergeUp = depth;
-    // Compute the cycle-id
+    std::vector<orc::ThreadSafeModule*> Children;
     M.withModuleDo([&](Module &m) {
         for (auto &F : m.global_objects()) {
             if (F.isDeclaration() && F.getLinkage() == GlobalValue::ExternalLinkage) {
                 auto Callee = NewExports.find(F.getName());
                 if (Callee != NewExports.end()) {
-                    auto &CM = Callee->second;
-                    int Down = jl_add_to_ee(*CM, NewExports, Queued, ToMerge, depth + 1);
-                    assert(Down <= depth);
-                    if (Down && Down < MergeUp)
-                        MergeUp = Down;
+                    auto *CM = Callee->second;
+                    if (*CM && CM != &M) {
+                        auto Down = Queued.find(CM);
+                        if (Down != Queued.end())
+                            MergeUp = std::min(MergeUp, Down->second);
+                        else
+                            Children.push_back(CM);
+                    }
                 }
             }
         }
     });
-    if (MergeUp == depth) {
-        // Not in a cycle (or at the top of it)
-        Queued.erase(&M);
-        for (auto &CM : ToMerge.at(depth - 1)) {
-            assert(Queued.find(CM)->second == depth);
-            Queued.erase(CM);
-            jl_merge_module(M, std::move(*CM));
-        }
-        jl_ExecutionEngine->addModule(std::move(M));
-        MergeUp = 0;
+    assert(MergeUp > 0);
+    for (auto *CM : Children) {
+        int Down = jl_add_to_ee(*CM, NewExports, Queued, Stack);
+        assert(Down <= (int)Stack.size());
+        if (Down)
+            MergeUp = std::min(MergeUp, Down);
     }
-    else {
-        // Add our frame(s) to the top of the cycle
-        Queued[&M] = MergeUp;
-        auto &Top = ToMerge.at(MergeUp - 1);
-        Top.push_back(&M);
-        for (auto &CM : ToMerge.at(depth - 1)) {
-            assert(Queued.find(CM)->second == depth);
-            Queued[CM] = MergeUp;
-            Top.push_back(CM);
+    if (MergeUp < depth)
+        return MergeUp;
+    while (1) {
+        // Not in a cycle (or at the top of it)
+        // remove SCC state and merge every CM from the cycle into M
+        orc::ThreadSafeModule *CM = Stack.back();
+        auto it = Queued.find(CM);
+        assert(it->second == (int)Stack.size());
+        Queued.erase(it);
+        Stack.pop_back();
+        if ((int)Stack.size() < depth) {
+            assert(&M == CM);
+            break;
         }
+        jl_merge_module(M, std::move(*CM));
     }
-    ToMerge.pop_back();
-    return MergeUp;
-}
-
-static void jl_add_to_ee(orc::ThreadSafeModule &M, StringMap<orc::ThreadSafeModule*> &NewExports)
-{
-    DenseMap<orc::ThreadSafeModule*, int> Queued;
-    std::vector<std::vector<orc::ThreadSafeModule*>> ToMerge;
-    jl_add_to_ee(M, NewExports, Queued, ToMerge, 1);
-    assert(!M);
+    jl_ExecutionEngine->addModule(std::move(M));
+    return 0;
 }
 
-
 static uint64_t getAddressForFunction(StringRef fname)
 {
     auto addr = jl_ExecutionEngine->getFunctionAddress(fname);