refactor/fix: exhaustive logic (#45)

Author: kermanx

crates/tree_shaker/src/builtins/react/context.rs

@@ -98,7 +98,7 @@ fn create_react_context_provider_impl<'a>(
         let dep = data.dep;
 
         let should_consume =
-          analyzer.request_exhaustive_callbacks(true, ExhaustiveDepId::Object(object_id));
+          analyzer.request_exhaustive_callbacks(ExhaustiveDepId::Object(object_id));
 
         if should_consume {
           analyzer.consume(context_id);

crates/tree_shaker/src/scope/cf_scope.rs

@@ -5,21 +5,14 @@ use crate::{
 };
 use oxc::{ast::ast::LabeledStatement, span::Atom};
 use oxc_index::define_index_type;
-use rustc_hash::FxHashSet;
 use std::mem;
 
-use super::exhaustive::ExhaustiveDepId;
+use super::exhaustive::{ExhaustiveData, ExhaustiveDepId};
 
 define_index_type! {
   pub struct CfScopeId = u32;
 }
 
-#[derive(Debug, Default)]
-pub struct ExhaustiveData {
-  pub clean: bool,
-  pub deps: FxHashSet<ExhaustiveDepId>,
-}
-
 #[derive(Debug)]
 pub enum CfScopeKind<'a> {
   Root,
@@ -117,33 +110,32 @@ impl<'a> CfScope<'a> {
     }
   }
 
-  pub fn mark_exhaustive_read(&mut self, id: ExhaustiveDepId) {
-    if let Some(data) = self.exhaustive_data_mut() {
-      if data.clean {
-        data.deps.insert(id);
-      }
-    }
-  }
-
   pub fn mark_exhaustive_write(&mut self, id: ExhaustiveDepId) -> bool {
     if let Some(data) = self.exhaustive_data_mut() {
-      if data.clean && data.deps.contains(&id) {
-        data.clean = false;
+      if data.clean {
+        if let Some(temp_deps) = &data.temp_deps {
+          if temp_deps.contains(&id) {
+            data.clean = false;
+          }
+        }
       }
       true
     } else {
       false
     }
   }
 
-  pub fn iterate_exhaustively(&mut self) -> bool {
+  pub fn post_exhaustive_iterate(&mut self) -> bool {
     let exited = self.must_exited();
     let data = self.exhaustive_data_mut().unwrap();
-    let clean = data.clean;
-    data.clean = true;
-    if !clean && !exited {
-      data.deps.clear();
-      true
+    if !data.clean && !exited {
+      if let Some(temp_deps) = &mut data.temp_deps {
+        temp_deps.clear();
+        data.clean = true;
+        true
+      } else {
+        false
+      }
     } else {
       false
     }

crates/tree_shaker/src/scope/exhaustive.rs

@@ -19,14 +19,21 @@ pub enum ExhaustiveDepId {
   ObjectProperty(ObjectId, ObjectPropertyId),
 }
 
+#[derive(Debug)]
+pub struct ExhaustiveData {
+  pub clean: bool,
+  pub temp_deps: Option<FxHashSet<ExhaustiveDepId>>,
+  pub register_deps: Option<FxHashSet<ExhaustiveDepId>>,
+}
+
 #[derive(Clone)]
 pub struct ExhaustiveCallback<'a> {
   pub handler: Rc<dyn Fn(&mut Analyzer<'a>) + 'a>,
-  pub once: bool,
+  pub drain: bool,
 }
 impl<'a> PartialEq for ExhaustiveCallback<'a> {
   fn eq(&self, other: &Self) -> bool {
-    self.once == other.once && Rc::ptr_eq(&self.handler, &other.handler)
+    self.drain == other.drain && Rc::ptr_eq(&self.handler, &other.handler)
   }
 }
 impl<'a> Eq for ExhaustiveCallback<'a> {}
@@ -40,7 +47,7 @@ impl<'a> Analyzer<'a> {
   pub fn exec_loop(&mut self, runner: impl Fn(&mut Analyzer<'a>) + 'a) {
     let runner = Rc::new(runner);
 
-    self.exec_exhaustively("loop", runner.clone(), false);
+    self.exec_exhaustively("loop", true, false, runner.clone());
 
     let cf_scope = self.cf_scope();
     if cf_scope.referred_state != ReferredState::ReferredClean && cf_scope.deps.may_not_referred() {
@@ -66,51 +73,55 @@ impl<'a> Analyzer<'a> {
       }
       analyzer.pop_cf_scope();
     });
-    let deps = self.exec_exhaustively(kind, runner.clone(), false);
-    self.register_exhaustive_callbacks(false, runner, deps);
+    self.exec_exhaustively(kind, true, true, runner);
   }
 
   pub fn exec_async_or_generator_fn(&mut self, runner: impl Fn(&mut Analyzer<'a>) + 'a) {
-    let runner = Rc::new(runner);
-    let deps = self.exec_exhaustively("async/generator", runner.clone(), true);
-    self.register_exhaustive_callbacks(true, runner, deps);
+    self.exec_exhaustively("async/generator", false, true, Rc::new(runner));
   }
 
   fn exec_exhaustively(
     &mut self,
     _kind: &str,
+    drain: bool,
+    register: bool,
     runner: Rc<dyn Fn(&mut Analyzer<'a>) + 'a>,
-    once: bool,
-  ) -> FxHashSet<ExhaustiveDepId> {
-    self.push_cf_scope(CfScopeKind::Exhaustive(Default::default()), Some(false));
+  ) {
+    self.push_cf_scope(
+      CfScopeKind::Exhaustive(ExhaustiveData {
+        clean: true,
+        temp_deps: drain.then(FxHashSet::default),
+        register_deps: register.then(Default::default),
+      }),
+      Some(false),
+    );
     let mut round_counter = 0;
-    while self.cf_scope_mut().iterate_exhaustively() {
+    loop {
       #[cfg(feature = "flame")]
       let _scope_guard = flame::start_guard(format!(
         "!{_kind}@{:06X} x{}",
         (Rc::as_ptr(&runner) as *const () as usize) & 0xFFFFFF,
         round_counter
       ));
-
       runner(self);
       round_counter += 1;
-      if once {
-        let data = self.cf_scope_mut().exhaustive_data_mut().unwrap();
-        data.clean = true;
-        break;
-      }
       if round_counter > 1000 {
         unreachable!("Exhaustive loop is too deep");
       }
+      if !self.cf_scope_mut().post_exhaustive_iterate() {
+        break;
+      }
     }
     let id = self.pop_cf_scope();
     let data = self.scoping.cf.get_mut(id).exhaustive_data_mut().unwrap();
-    mem::take(&mut data.deps)
+    if let Some(register_deps) = data.register_deps.take() {
+      self.register_exhaustive_callbacks(drain, runner, register_deps);
+    }
   }
 
   fn register_exhaustive_callbacks(
     &mut self,
-    once: bool,
+    drain: bool,
     handler: Rc<dyn Fn(&mut Analyzer<'a>) + 'a>,
     deps: FxHashSet<ExhaustiveDepId>,
   ) {
@@ -119,43 +130,58 @@ impl<'a> Analyzer<'a> {
         .exhaustive_callbacks
         .entry(id)
         .or_default()
-        .insert(ExhaustiveCallback { handler: handler.clone(), once });
+        .insert(ExhaustiveCallback { handler: handler.clone(), drain });
     }
   }
 
   pub fn mark_exhaustive_read(&mut self, id: ExhaustiveDepId, target: usize) {
-    for depth in target..self.scoping.cf.stack.len() {
-      self.scoping.cf.get_mut_from_depth(depth).mark_exhaustive_read(id);
+    let mut registered = false;
+    for depth in (target..self.scoping.cf.stack.len()).rev() {
+      let scope = self.scoping.cf.get_mut_from_depth(depth);
+      if let Some(data) = scope.exhaustive_data_mut() {
+        if data.clean {
+          if let Some(temp_deps) = data.temp_deps.as_mut() {
+            temp_deps.insert(id);
+          }
+        }
+        if !registered {
+          if let Some(register_deps) = data.register_deps.as_mut() {
+            registered = true;
+            register_deps.insert(id);
+          }
+        }
+      }
     }
   }
 
   pub fn mark_exhaustive_write(&mut self, id: ExhaustiveDepId, target: usize) -> (bool, bool) {
-    let mut should_consume = false;
+    let mut exhaustive = false;
     let mut indeterminate = false;
+    let mut need_mark = true;
     for depth in target..self.scoping.cf.stack.len() {
       let scope = self.scoping.cf.get_mut_from_depth(depth);
-      if !should_consume {
-        should_consume |= scope.mark_exhaustive_write(id);
-      }
       indeterminate |= scope.is_indeterminate();
+      if let Some(data) = scope.exhaustive_data_mut() {
+        exhaustive = true;
+        if (need_mark || data.register_deps.is_some()) && data.clean {
+          if let Some(temp_deps) = &data.temp_deps {
+            if temp_deps.contains(&id) {
+              data.clean = false;
+            }
+            need_mark = false;
+          }
+        }
+      }
     }
-    (should_consume, indeterminate)
+    (exhaustive, indeterminate)
   }
 
-  pub fn request_exhaustive_callbacks(
-    &mut self,
-    should_consume: bool,
-    id: ExhaustiveDepId,
-  ) -> bool {
+  pub fn request_exhaustive_callbacks(&mut self, id: ExhaustiveDepId) -> bool {
     if let Some(runners) = self.exhaustive_callbacks.get_mut(&id) {
       if runners.is_empty() {
         false
       } else {
-        if should_consume {
-          self.pending_deps.extend(runners.drain());
-        } else {
-          self.pending_deps.extend(runners.iter().cloned());
-        }
+        self.pending_deps.extend(runners.drain());
         true
       }
     } else {
@@ -171,9 +197,8 @@ impl<'a> Analyzer<'a> {
       let runners = mem::take(&mut self.pending_deps);
       for runner in runners {
         // let old_count = self.referred_deps.debug_count();
-        let ExhaustiveCallback { handler: runner, once } = runner;
-        let deps = self.exec_exhaustively("dep", runner.clone(), once);
-        self.register_exhaustive_callbacks(once, runner, deps);
+        let ExhaustiveCallback { handler: runner, drain } = runner;
+        self.exec_exhaustively("dep", drain, true, runner.clone());
         // let new_count = self.referred_deps.debug_count();
         // self.debug += 1;
       }

crates/tree_shaker/src/scope/utils.rs

@@ -30,7 +30,7 @@ impl<'a> Analyzer<'a> {
       }
     }
 
-    self.request_exhaustive_callbacks(true, ExhaustiveDepId::Object(object_id));
+    self.request_exhaustive_callbacks(ExhaustiveDepId::Object(object_id));
 
     (has_exhaustive, indeterminate, exec_deps)
   }
@@ -50,6 +50,6 @@ impl<'a> Analyzer<'a> {
       }
       mem::take(&mut scope.deps).consume_all(self);
     }
-    self.request_exhaustive_callbacks(true, ExhaustiveDepId::Object(object_id));
+    self.request_exhaustive_callbacks(ExhaustiveDepId::Object(object_id));
   }
 }

crates/tree_shaker/src/scope/variable_scope.rs

@@ -103,7 +103,7 @@ impl<'a> Analyzer<'a> {
       }));
       self.scoping.variable.get_mut(id).variables.insert(symbol, variable);
       if has_fn_value {
-        self.request_exhaustive_callbacks(false, ExhaustiveDepId::Variable(id, symbol));
+        self.request_exhaustive_callbacks(ExhaustiveDepId::Variable(id, symbol));
       }
     }
   }
@@ -130,7 +130,7 @@ impl<'a> Analyzer<'a> {
     } else {
       variable.value =
         Some(self.factory.computed(value.unwrap_or(self.factory.undefined), init_node));
-      self.request_exhaustive_callbacks(false, ExhaustiveDepId::Variable(id, symbol));
+      self.request_exhaustive_callbacks(ExhaustiveDepId::Variable(id, symbol));
     }
   }
 
@@ -226,7 +226,7 @@ impl<'a> Analyzer<'a> {
           };
           drop(variable_ref);
 
-          self.request_exhaustive_callbacks(should_consume, ExhaustiveDepId::Variable(id, symbol));
+          self.request_exhaustive_callbacks(ExhaustiveDepId::Variable(id, symbol));
         }
       }
       true