From 61e91c24ae78fc8c9afe407bfaa82f353a5b6ff1 Mon Sep 17 00:00:00 2001
From: Shunsuke Shibayama <sbym1346@gmail.com>
Date: Sat, 14 Dec 2024 14:53:55 +0900
Subject: [PATCH] fix: `SharedPromises::join` deadlock (2)

---
 crates/erg_compiler/module/graph.rs   |  5 ++++-
 crates/erg_compiler/module/promise.rs | 29 ++++++++++++++++++---------
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/crates/erg_compiler/module/graph.rs b/crates/erg_compiler/module/graph.rs
index c26b06cde..a47d0afe3 100644
--- a/crates/erg_compiler/module/graph.rs
+++ b/crates/erg_compiler/module/graph.rs
@@ -127,8 +127,11 @@ impl ModuleGraph {
     /// ```erg
     /// # a.er
     /// b = import "b"
+    /// # -> a: child, b: parent
+    /// # b.er
+    /// c = import "c"
+    /// # -> ancestors(a) == {b, c}
     /// ```
-    /// -> a: child, b: parent<br>
     /// O(N)
     pub fn ancestors<'p>(&'p self, path: &'p NormalizedPathBuf) -> Set<&'p NormalizedPathBuf> {
         let mut ancestors = set! {};
diff --git a/crates/erg_compiler/module/promise.rs b/crates/erg_compiler/module/promise.rs
index 72a5a451d..26355e1d5 100644
--- a/crates/erg_compiler/module/promise.rs
+++ b/crates/erg_compiler/module/promise.rs
@@ -182,13 +182,10 @@ impl SharedPromises {
         if !self.graph.entries().contains(path) {
             return Err(Box::new(format!("not registered: {path}")));
         }
-        // prevent deadlock
-        if self.thread_id(path).is_some_and(|id| id == current().id()) {
-            return Ok(());
-        }
-        if self.graph.ancestors(path).contains(&self.root) || path == &self.root {
-            // cycle detected, `self.root` must not in the dependencies
-            // Erg analysis processes never join ancestor threads (although joining ancestors itself is allowed in Rust)
+        let current = self.current_path();
+        if self.graph.ancestors(path).contains(&current) || path == &current {
+            // cycle detected, `current` must not in the dependencies
+            // Erg analysis processes never join themselves / ancestor threads
             // self.wait_until_finished(path);
             return Ok(());
         }
@@ -198,12 +195,14 @@ impl SharedPromises {
         }
         // Suppose A depends on B and C, and B depends on C.
         // In this case, B must join C before A joins C. Otherwise, a deadlock will occur.
+        // C.children() == {A, B}
         let children = self.graph.children(path);
         for child in children.iter() {
-            if child == &self.root {
+            if child == &current {
                 continue;
-            } else if self.graph.depends_on(&self.root, child) {
-                self.wait_until_finished(path);
+            } else if self.graph.depends_on(&current, child) {
+                // A.depends_on(B) => A.wait_until_finished(B)
+                self.wait_until_finished(child);
                 return Ok(());
             }
         }
@@ -252,6 +251,16 @@ impl SharedPromises {
             .and_then(|promise| promise.thread_id())
     }
 
+    pub fn current_path(&self) -> NormalizedPathBuf {
+        let cur_id = current().id();
+        for (path, promise) in self.promises.borrow().iter() {
+            if promise.thread_id() == Some(cur_id) {
+                return path.clone();
+            }
+        }
+        self.root.clone()
+    }
+
     pub fn progress(&self) -> Progress {
         let mut total = 0;
         let mut running = 0;