Improve cleaning unused var declarations

.changeset/curly-lions-allow.md

@@ -0,0 +1,6 @@
+---
+'@stylexswc/swc-plugin': patch
+'@stylexswc/nextjs-plugin': patch
+---
+
+rename ModuleCycle to TransformationCycle

.changeset/sixty-beers-end.md

@@ -0,0 +1,6 @@
+---
+'@stylexswc/swc-plugin': patch
+'@stylexswc/nextjs-plugin': patch
+---
+
+link counting: improve cleaning unused var declarations

.changeset/small-apples-exercise.md

@@ -0,0 +1,6 @@
+---
+'@stylexswc/swc-plugin': patch
+'@stylexswc/nextjs-plugin': patch
+---
+
+state: unbox atom entities

packages/swc-plugin/src/shared/enums/core.rs

@@ -1,15 +1,22 @@
 // Represents the current state of a plugin for a file.
 #[derive(Debug, PartialEq, Eq, Clone, Hash, Copy)]
-pub(crate) enum ModuleCycle {
-  // The plugin is being processed.
+pub(crate) enum TransformationCycle {
+  // The plugin is being processed
   TransformEnter,
+  // The plugin has been processed
   TransformExit,
-  // The plugin has been processed and the file is being cleaned.
+  // The plugin has been processed and the file is being cleaned
   PreCleaning,
+  // The file is being cleaned
   Cleaning,
-  // The file has been processed and the plugin is skipped.
+  // Recounting variable links
+  Recounting,
+  // The file has been processed and the plugin is skipped
   Initializing,
+  // Fill the state with expressions data before transformation
   StateFilling,
+  // Skip the plugin if import does not exist
   Skip,
+  // Inject styles metadata to the file
   InjectStyles,
 }

packages/swc-plugin/src/shared/structures/state_manager.rs

@@ -15,10 +15,13 @@ use swc_core::{
   common::{EqIgnoreSpan, FileName, DUMMY_SP},
 };
 
-use crate::shared::enums::data_structures::{
-  import_path_resolution::{ImportPathResolution, ImportPathResolutionType},
-  style_vars_to_keep::StyleVarsToKeep,
-  top_level_expression::{TopLevelExpression, TopLevelExpressionKind},
+use crate::shared::enums::{
+  core::TransformationCycle,
+  data_structures::{
+    import_path_resolution::{ImportPathResolution, ImportPathResolutionType},
+    style_vars_to_keep::StyleVarsToKeep,
+    top_level_expression::{TopLevelExpression, TopLevelExpressionKind},
+  },
 };
 use crate::shared::utils::{
   ast::factories::binding_ident_factory,
@@ -44,38 +47,41 @@ use super::{
   seen_value::SeenValue,
 };
 
+type AtomHashMap = HashMap<Atom, i16>;
+type AtomHashSet = HashSet<Atom>;
+
 #[derive(Clone, Debug)]
 pub struct StateManager {
   pub(crate) _state: Box<PluginPass>,
 
   // Imports
   pub(crate) import_paths: HashSet<String>,
-  pub(crate) stylex_import: HashSet<Box<ImportSources>>,
-  pub(crate) stylex_props_import: HashSet<Box<Atom>>,
-  pub(crate) stylex_attrs_import: HashSet<Box<Atom>>,
-  pub(crate) stylex_create_import: HashSet<Box<Atom>>,
-  pub(crate) stylex_include_import: HashSet<Box<Atom>>,
-  pub(crate) stylex_first_that_works_import: HashSet<Box<Atom>>,
-  pub(crate) stylex_keyframes_import: HashSet<Box<Atom>>,
-  pub(crate) stylex_define_vars_import: HashSet<Box<Atom>>,
-  pub(crate) stylex_create_theme_import: HashSet<Box<Atom>>,
-  pub(crate) stylex_types_import: HashSet<Box<Atom>>,
+  pub(crate) stylex_import: HashSet<ImportSources>,
+  pub(crate) stylex_props_import: AtomHashSet,
+  pub(crate) stylex_attrs_import: AtomHashSet,
+  pub(crate) stylex_create_import: AtomHashSet,
+  pub(crate) stylex_include_import: AtomHashSet,
+  pub(crate) stylex_first_that_works_import: AtomHashSet,
+  pub(crate) stylex_keyframes_import: AtomHashSet,
+  pub(crate) stylex_define_vars_import: AtomHashSet,
+  pub(crate) stylex_create_theme_import: AtomHashSet,
+  pub(crate) stylex_types_import: AtomHashSet,
   pub(crate) inject_import_inserted: Option<(Box<Ident>, Box<Ident>)>,
   pub(crate) theme_name: Option<String>,
 
   pub(crate) declarations: Vec<VarDeclarator>,
   pub(crate) top_level_expressions: Vec<TopLevelExpression>,
   pub(crate) all_call_expressions: Vec<CallExpr>,
-  pub(crate) var_decl_count_map: HashMap<Atom, i8>,
-  pub(crate) seen: HashMap<Box<Expr>, Box<SeenValue>>,
+  pub(crate) var_decl_count_map: AtomHashMap,
+  pub(crate) seen: HashMap<Box<Expr>, (Box<SeenValue>, Option<AtomHashMap>)>,
 
   // `stylex.create` calls
   pub(crate) style_map: HashMap<String, Box<StylesObjectMap>>,
   pub(crate) style_vars: HashMap<String, Box<VarDeclarator>>,
 
   // results of `stylex.create` calls that should be kept
   pub(crate) style_vars_to_keep: HashSet<Box<StyleVarsToKeep>>,
-  pub(crate) member_object_ident_count_map: HashMap<Atom, i8>,
+  pub(crate) member_object_ident_count_map: AtomHashMap,
 
   pub(crate) in_stylex_create: bool,
 
@@ -87,6 +93,8 @@ pub struct StateManager {
 
   pub(crate) injected_keyframes: IndexMap<String, Box<InjectableStyle>>,
   pub(crate) top_imports: Vec<ImportDecl>,
+
+  pub(crate) cycle: TransformationCycle,
 }
 
 impl Default for StateManager {
@@ -137,6 +145,8 @@ impl StateManager {
       prepend_import_module_items: vec![],
 
       injected_keyframes: IndexMap::new(),
+
+      cycle: TransformationCycle::Initializing,
     }
   }
 
@@ -177,7 +187,7 @@ impl StateManager {
       .stylex_import
       .clone()
       .into_iter()
-      .map(|import_source| match import_source.as_ref() {
+      .map(|import_source| match &import_source {
         ImportSources::Regular(regular) => regular.clone(),
         ImportSources::Named(named) => named.clone().r#as,
       })

packages/swc-plugin/src/shared/structures/types.rs

@@ -17,7 +17,7 @@ pub(crate) type EvaluateResultFns =
   IndexMap<String, (Vec<BindingIdent>, IndexMap<String, Box<Expr>>)>;
 pub(crate) type EvaluationCallback = Rc<dyn Fn(Vec<Option<EvaluateResultValue>>) -> Expr + 'static>;
 pub(crate) type FunctionMapMemberExpression =
-  HashMap<Box<ImportSources>, Box<HashMap<Box<Atom>, Box<FunctionConfigType>>>>;
-pub(crate) type FunctionMapIdentifiers = HashMap<Box<Atom>, Box<FunctionConfigType>>;
+  HashMap<ImportSources, Box<HashMap<Atom, Box<FunctionConfigType>>>>;
+pub(crate) type FunctionMapIdentifiers = HashMap<Atom, Box<FunctionConfigType>>;
 pub(crate) type StylesObjectMap =
   IndexMap<String, Box<IndexMap<String, Box<FlatCompiledStylesValue>>>>;

packages/swc-plugin/src/shared/utils/common.rs

@@ -1,7 +1,7 @@
 use radix_fmt::radix;
 use std::{
   any::type_name,
-  collections::HashSet,
+  collections::{hash_map::Entry, HashMap, HashSet},
   hash::{DefaultHasher, Hash, Hasher},
   ops::Deref,
   path::PathBuf,
@@ -92,10 +92,9 @@ pub(crate) fn wrap_key_in_quotes(key: &str, should_wrap_in_quotes: &bool) -> Str
 }
 
 pub fn reduce_ident_count<'a>(state: &'a mut StateManager, ident: &'a Ident) {
-  *state
-    .var_decl_count_map
-    .entry(ident.sym.clone())
-    .or_insert(0) -= 1;
+  if let Entry::Occupied(mut entry) = state.var_decl_count_map.entry(ident.sym.clone()) {
+    *entry.get_mut() -= 1;
+  }
 }
 
 pub fn increase_member_ident(state: &mut StateManager, member_obj: &MemberExpr) {
@@ -111,10 +110,12 @@ pub fn reduce_member_expression_count(state: &mut StateManager, member_expressio
 }
 
 pub fn reduce_member_ident_count(state: &mut StateManager, ident_atom: &Atom) {
-  *state
+  if let Entry::Occupied(mut entry) = state
     .member_object_ident_count_map
     .entry(ident_atom.clone())
-    .or_insert(0) -= 1;
+  {
+    *entry.get_mut() -= 1;
+  }
 }
 
 pub fn increase_ident_count(state: &mut StateManager, ident: &Ident) {
@@ -125,7 +126,7 @@ pub fn increase_member_ident_count(state: &mut StateManager, ident_atom: &Atom)
   increase_member_ident_count_by_count(state, ident_atom, 1);
 }
 
-pub fn increase_ident_count_by_count(state: &mut StateManager, ident: &Ident, count: i8) {
+pub fn increase_ident_count_by_count(state: &mut StateManager, ident: &Ident, count: i16) {
   let ident_id = &ident.sym;
 
   *state
@@ -137,7 +138,7 @@ pub fn increase_ident_count_by_count(state: &mut StateManager, ident: &Ident, co
 pub fn increase_member_ident_count_by_count(
   state: &mut StateManager,
   ident_atom: &Atom,
-  count: i8,
+  count: i16,
 ) {
   *state
     .member_object_ident_count_map
@@ -351,6 +352,74 @@ pub(crate) fn deep_merge_props(
   remove_duplicates(new_props.into_iter().rev().collect())
 }
 
+pub(crate) fn get_hash_map_difference<K, V>(
+  orig_map: &HashMap<K, V>,
+  compare_map: &HashMap<K, V>,
+) -> HashMap<K, V>
+where
+  K: Eq + Hash + Clone,
+  V: PartialEq + Clone,
+{
+  let mut diff = HashMap::new();
+
+  for (key, value) in orig_map {
+    if let Some(map2_value) = compare_map.get(key) {
+      if value != map2_value {
+        diff.insert(key.clone(), value.clone());
+      }
+    } else {
+      diff.insert(key.clone(), value.clone());
+    }
+  }
+
+  for (key, value) in compare_map {
+    if !orig_map.contains_key(key) {
+      diff.insert(key.clone(), value.clone());
+    }
+  }
+
+  diff
+}
+
+pub(crate) fn get_hash_map_value_difference(
+  orig_map: &HashMap<Atom, i16>,
+  map2: &HashMap<Atom, i16>,
+) -> HashMap<Atom, i16> {
+  let mut diff = HashMap::new();
+
+  for (key, value) in orig_map {
+    if let Some(map2_value) = map2.get(key) {
+      if value != map2_value {
+        diff.insert(key.clone(), value - map2_value);
+      }
+    } else {
+      diff.insert(key.clone(), *value);
+    }
+  }
+
+  diff
+}
+
+pub(crate) fn sum_hash_map_values(
+  orig_map: &HashMap<Atom, i16>,
+  compare_map: &HashMap<Atom, i16>,
+) -> HashMap<Atom, i16> {
+  let mut sum_map = HashMap::new();
+
+  for (key, value) in orig_map {
+    sum_map.insert(key.clone(), *value);
+  }
+
+  for (key, value) in compare_map {
+    sum_map
+      .entry(key.clone())
+      .and_modify(|e| *e += value)
+      .or_insert(*value);
+  }
+
+  sum_map
+}
+
 pub(crate) fn get_css_value(key_value: KeyValueProp) -> (Box<Expr>, Option<BaseCSSType>) {
   let Some(obj) = key_value.value.as_object() else {
     return (key_value.value, None);