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.

(cherry picked from commit b03439c)

Author: vtjnash

src/gf.c

@@ -3421,6 +3421,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

@@ -137,7 +137,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);
 
@@ -229,10 +233,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());
@@ -1715,76 +1722,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);