From e718ba21a8e8cfdbf9d5beb1c34bf14bba1d1bee Mon Sep 17 00:00:00 2001
From: Juni May <juni_may@outlook.com>
Date: Wed, 17 Jul 2024 11:13:24 +0800
Subject: [PATCH] feat(utils): get unclosed loop values

---
 src/utils/loop_info.rs | 47 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 45 insertions(+), 2 deletions(-)

diff --git a/src/utils/loop_info.rs b/src/utils/loop_info.rs
index 9d0ac7d..39d176c 100644
--- a/src/utils/loop_info.rs
+++ b/src/utils/loop_info.rs
@@ -1,13 +1,20 @@
-use std::hash::Hash;
+use std::{
+    collections::{HashSet, VecDeque},
+    hash::Hash,
+};
 
 use rustc_hash::FxHashMap;
 
 use super::{
     cfg::{CfgInfo, CfgNode},
+    def_use::Usable,
     dominance::Dominance,
 };
 use crate::{
-    collections::storage::{ArenaAlloc, ArenaDeref, ArenaPtr, BaseArena, BaseArenaPtr},
+    collections::{
+        linked_list::{LinkedListContainerPtr, LinkedListNodePtr},
+        storage::{ArenaAlloc, ArenaDeref, ArenaPtr, BaseArena, BaseArenaPtr},
+    },
     ir,
 };
 
@@ -110,6 +117,42 @@ impl Loop<ir::Block> {
         }
         None
     }
+
+    /// Get all the values that are defined in the loop but used outside the
+    /// loop.
+    pub fn get_unclosed_values(
+        self,
+        ctx: &ir::Context,
+        loop_ctx: &LoopContext<ir::Block>,
+    ) -> HashSet<ir::Value> {
+        let header = self.header(loop_ctx);
+
+        let mut queue = VecDeque::new();
+        let mut visited = HashSet::new();
+
+        queue.push_back(header);
+
+        let mut values = HashSet::new();
+
+        while let Some(block) = queue.pop_front() {
+            if !visited.insert(block) {
+                continue;
+            }
+
+            for inst in block.iter(ctx) {
+                for result in inst.results(ctx) {
+                    for user in result.users(ctx) {
+                        let user_block = user.container(ctx).unwrap();
+                        if !loop_ctx.is_in_loop(user_block, self) {
+                            values.insert(*result);
+                        }
+                    }
+                }
+            }
+        }
+
+        values
+    }
 }
 
 pub struct LoopWithDepth<N> {