Skip to main content

stylex_transform/transform/fold/
fold_module_items.rs

1use stylex_macros::stylex_panic;
2use swc_core::ecma::ast::{ExportDecl, ImportDecl, Str};
3use swc_core::{
4  common::{DUMMY_SP, comments::Comments},
5  ecma::{
6    ast::{Decl, Expr, Ident, Lit, ModuleDecl, ModuleItem, Pat, Stmt, VarDeclarator},
7    visit::FoldWith,
8  },
9};
10
11use crate::{
12  StyleXTransform,
13  shared::utils::{
14    ast::convertors::convert_atom_to_string,
15    common::{fill_state_declarations, stable_hash},
16  },
17};
18use stylex_ast::ast::factories::create_binding_ident;
19use stylex_constants::constants::messages::VAR_DECL_INIT_REQUIRED;
20use stylex_enums::core::TransformationCycle;
21use stylex_regex::regex::STYLEX_CONSTS_IMPORT_REGEX;
22
23impl<C> StyleXTransform<C>
24where
25  C: Comments,
26{
27  pub(crate) fn fold_module_items(&mut self, module_items: Vec<ModuleItem>) -> Vec<ModuleItem> {
28    match self.state.cycle {
29      TransformationCycle::Skip => module_items,
30      TransformationCycle::Initializing => {
31        let mut transformed_module_items = module_items.fold_children_with(self);
32
33        if self.state.import_paths.is_empty() {
34          self.state.cycle = TransformationCycle::Skip;
35
36          return transformed_module_items;
37        }
38
39        if self.state.options.inject_stylex_side_effects {
40          let mut side_effect_imports = Vec::new();
41
42          for module_item in &transformed_module_items {
43            if let ModuleItem::ModuleDecl(ModuleDecl::Import(import_decl)) = module_item {
44              let source_path = convert_atom_to_string(&import_decl.src.value);
45              if STYLEX_CONSTS_IMPORT_REGEX
46                .is_match(&source_path)
47                .unwrap_or(false)
48              {
49                side_effect_imports.push(ModuleItem::ModuleDecl(ModuleDecl::Import(ImportDecl {
50                  span: DUMMY_SP,
51                  specifiers: vec![],
52                  src: Box::new(Str {
53                    span: DUMMY_SP,
54                    value: import_decl.src.value.clone(),
55                    raw: None,
56                  }),
57                  type_only: false,
58                  with: None,
59                  phase: Default::default(),
60                })));
61              }
62            }
63          }
64
65          transformed_module_items.extend(side_effect_imports);
66        }
67
68        transformed_module_items
69      },
70      TransformationCycle::StateFilling => {
71        module_items.iter().for_each(|module_item| {
72          if let ModuleItem::Stmt(Stmt::Decl(Decl::Var(var_decl))) = module_item {
73            var_decl.decls.iter().for_each(|decl| {
74              if let Pat::Ident(_) = &decl.name {
75                fill_state_declarations(&mut self.state, decl);
76              }
77            });
78          }
79        });
80
81        module_items.fold_children_with(self)
82      },
83      TransformationCycle::TransformEnter
84      | TransformationCycle::PreCleaning
85      | TransformationCycle::Recounting => module_items.fold_children_with(self),
86      TransformationCycle::TransformExit => {
87        if self.state.hoisted_module_items.is_empty() {
88          return module_items.fold_children_with(self);
89        }
90
91        let import_end = module_items
92          .iter()
93          .enumerate()
94          .skip(1)
95          .find(|(_, item)| !matches!(item, ModuleItem::ModuleDecl(ModuleDecl::Import(_))))
96          .map(|(idx, _)| idx)
97          .unwrap_or(0);
98
99        let mut module_items = module_items;
100
101        let total_capacity = module_items.len() + self.state.hoisted_module_items.len();
102
103        let mut result_module_items = Vec::with_capacity(total_capacity);
104
105        result_module_items.extend(module_items.drain(..import_end));
106
107        result_module_items.extend(self.state.hoisted_module_items.iter().cloned());
108
109        result_module_items.extend(module_items);
110
111        result_module_items.fold_children_with(self)
112      },
113      TransformationCycle::InjectStyles => {
114        let mut result_module_items: Vec<ModuleItem> =
115          self.state.prepend_include_module_items.clone();
116
117        result_module_items.extend(self.state.prepend_import_module_items.clone());
118
119        let mut items_to_skip: usize = 0;
120
121        if let Some(first) = module_items
122          .first()
123          .and_then(|first| first.as_stmt())
124          .and_then(|stmp| stmp.as_expr())
125          && let Some(Lit::Str(_)) = first.expr.as_lit()
126        {
127          result_module_items.insert(
128            0,
129            match module_items.first() {
130              Some(item) => item.clone(),
131              None => stylex_panic!("Module items list is unexpectedly empty."),
132            },
133          );
134          items_to_skip = 1;
135        }
136
137        for module_item in module_items.iter().skip(items_to_skip) {
138          if let Some(decls) = match &module_item {
139            ModuleItem::ModuleDecl(decl) => match decl {
140              ModuleDecl::ExportDecl(export_decl) => export_decl.decl.as_var().map(|var_decl| {
141                var_decl
142                  .decls
143                  .iter() // Use iter() to avoid cloning the entire collection
144                  .filter(|decl| {
145                    decl
146                      .init
147                      .as_ref()
148                      .is_some_and(|init| init.is_object() || init.is_lit())
149                  })
150                  .cloned() // Clone only the filtered elements
151                  .collect::<Vec<VarDeclarator>>()
152              }),
153              ModuleDecl::ExportDefaultExpr(export_default_expr) => {
154                export_default_expr.expr.as_object().map(|obj| {
155                  vec![VarDeclarator {
156                    definite: true,
157                    span: DUMMY_SP,
158                    name: Pat::Ident(create_binding_ident(Ident::from("default"))),
159                    init: Some(Box::new(Expr::from(obj.clone()))),
160                  }]
161                })
162              },
163              _ => None,
164            },
165            ModuleItem::Stmt(Stmt::Decl(Decl::Var(var_decl))) => Some(
166              var_decl
167                .decls
168                .iter()
169                .filter(|decl| {
170                  decl
171                    .init
172                    .as_ref()
173                    .is_some_and(|init| init.is_object() || init.is_lit())
174                })
175                .cloned()
176                .collect::<Vec<VarDeclarator>>(),
177            ),
178            _ => None,
179          } {
180            for decl in decls {
181              let key = match decl.init.clone() {
182                Some(k) => k,
183                None => stylex_panic!("{}", VAR_DECL_INIT_REQUIRED),
184              };
185
186              let key_hash = stable_hash(key.as_ref());
187
188              if let Some(metadata_items) = self.state.styles_to_inject.get(&key_hash) {
189                for module_item in metadata_items.iter() {
190                  result_module_items.push(module_item.clone());
191                }
192              }
193            }
194          }
195          result_module_items.push(module_item.clone());
196        }
197
198        result_module_items
199      },
200      TransformationCycle::Cleaning => {
201        // We need it twice for a clear dead code after declaration transforms
202        let mut module_items = module_items.fold_children_with(self);
203
204        // We remove `Stmt::Empty` from the statement list.
205        // This is optional, but it's required if you don't want extra `;` in output.
206        module_items.retain(|module_item| {
207          !matches!(module_item, ModuleItem::Stmt(Stmt::Empty(..)))
208            && !matches!(module_item, ModuleItem::ModuleDecl(ModuleDecl::ExportDecl(ExportDecl { decl: Decl::Var(var), .. })) if var.decls.is_empty())
209        });
210
211        module_items
212      },
213    }
214  }
215}