stylex_transform/transform/fold/
fold_var_declarator.rs

1use std::collections::hash_map::Entry;
2
3use rustc_hash::FxHashMap;
4use stylex_macros::stylex_panic;
5use swc_core::{
6  atoms::Atom,
7  common::{EqIgnoreSpan, comments::Comments},
8  ecma::{
9    ast::{Expr, KeyValueProp, Lit, ObjectLit, Prop, PropName, PropOrSpread, VarDeclarator},
10    visit::FoldWith,
11  },
12};
13
14use stylex_enums::style_vars_to_keep::{NonNullProp, NonNullProps};
15use stylex_enums::top_level_expression::TopLevelExpressionKind;
16use stylex_structures::style_vars_to_keep::StyleVarsToKeep;
17use stylex_structures::top_level_expression::TopLevelExpression;
18
19use crate::{
20  StyleXTransform,
21  shared::utils::{ast::convertors::expand_shorthand_prop, common::fill_state_declarations},
22};
23use stylex_constants::constants::messages::{
24  KEY_VALUE_EXPECTED, PROPERTY_NOT_FOUND, VAR_DECL_NAME_NOT_IDENT,
25};
26use stylex_enums::core::TransformationCycle;
27
28impl<C> StyleXTransform<C>
29where
30  C: Comments,
31{
32  pub(crate) fn fold_var_declarator_impl(
33    &mut self,
34    mut var_declarator: VarDeclarator,
35  ) -> VarDeclarator {
36    match self.state.cycle {
37      TransformationCycle::StateFilling => {
38        fill_state_declarations(&mut self.state, &var_declarator);
39
40        if let Some(Expr::Call(call)) = var_declarator.init.as_deref_mut()
41          && let Some((declaration, member)) = self.process_declaration(call)
42        {
43          let stylex_imports = self.state.stylex_import_stringified();
44
45          if let Some(declaration_string) = stylex_imports
46            .into_iter()
47            .find(|item| item.eq(declaration.0.as_str()))
48            .or_else(|| {
49              self
50                .state
51                .stylex_create_import
52                .iter()
53                .find(|decl| decl.eq_ignore_span(&&declaration.0.clone()))
54                .map(|decl| decl.to_string())
55            })
56            && self.state.cycle == TransformationCycle::StateFilling
57            && (member.as_str() == "create" || member.eq(declaration_string.as_str()))
58          {
59            self.props_declaration = var_declarator.name.as_ident().map(|ident| ident.to_id());
60          }
61        }
62
63        var_declarator.fold_children_with(self)
64      },
65      TransformationCycle::Cleaning => {
66        {
67          let mut vars_to_keep: FxHashMap<Atom, NonNullProps> = FxHashMap::default();
68
69          for StyleVarsToKeep(var_name, namespace_name, _) in
70            self.state.style_vars_to_keep.clone().into_iter()
71          {
72            match vars_to_keep.entry(var_name) {
73              Entry::Occupied(mut entry) => {
74                let entry_value = entry.get_mut();
75
76                if let NonNullProps::Vec(vec) = entry_value {
77                  match namespace_name {
78                    NonNullProp::Atom(id) => {
79                      vec.push(id);
80                    },
81                    _ => {
82                      *entry_value = NonNullProps::True;
83                    },
84                  }
85                }
86              },
87              Entry::Vacant(entry) => {
88                let value = match namespace_name {
89                  NonNullProp::Atom(namespace_name) => NonNullProps::Vec(vec![namespace_name]),
90                  NonNullProp::True => NonNullProps::True,
91                };
92
93                entry.insert(value);
94              },
95            }
96          }
97
98          for (_, var_name) in self.state.style_vars.clone().into_iter() {
99            if !var_declarator.eq_ignore_span(&var_name) {
100              continue;
101            };
102
103            let top_level_expression = self.state.top_level_expressions.clone().into_iter().find(
104              |TopLevelExpression(_, expr, _)| match var_name.init.clone() {
105                Some(init) => init.eq_ignore_span(&Box::new(expr.clone())),
106                None => {
107                  stylex_panic!(
108                    "Variable declaration must have an initializer for top-level expression lookup."
109                  )
110                },
111              },
112            );
113
114            if let Some(TopLevelExpression(kind, _, _)) = top_level_expression
115              && kind == TopLevelExpressionKind::Stmt
116              && let Some(object) = var_declarator
117                .init
118                .as_mut()
119                .and_then(|var_decl| var_decl.as_object())
120            {
121              let namespaces_to_keep = match vars_to_keep.get(&match var_name.name.as_ident() {
122                Some(i) => i.sym.clone(),
123                None => stylex_panic!("{}", VAR_DECL_NAME_NOT_IDENT),
124              }) {
125                Some(NonNullProps::Vec(vec)) => vec.clone(),
126                _ => Vec::new(),
127              };
128
129              if !namespaces_to_keep.is_empty() {
130                let mut new_object = object.clone();
131
132                let props =
133                  self.retain_object_props(&mut new_object, namespaces_to_keep, &var_name);
134                new_object.props = props;
135
136                **match var_declarator.init.as_mut() {
137                  Some(init) => init,
138                  None => stylex_panic!(
139                    "Variable declaration must have an initializer when updating an object."
140                  ),
141                } = Expr::Object(new_object);
142              }
143            }
144          }
145        }
146
147        var_declarator
148      },
149      TransformationCycle::Skip | TransformationCycle::InjectStyles => var_declarator,
150      _ => var_declarator.fold_children_with(self),
151    }
152  }
153
154  fn retain_object_props(
155    &self,
156    object: &mut ObjectLit,
157    namespace_to_keep: Vec<Atom>,
158    var_name: &VarDeclarator,
159  ) -> Vec<PropOrSpread> {
160    let mut props: Vec<PropOrSpread> = vec![];
161
162    for object_prop in object.props.iter_mut() {
163      assert!(object_prop.is_prop(), "Spread properties are not supported");
164
165      let prop = match object_prop.as_mut_prop() {
166        Some(p) => p.as_mut(),
167        None => stylex_panic!("{}", PROPERTY_NOT_FOUND),
168      };
169
170      if let Some(KeyValueProp { key, .. }) = prop.as_key_value() {
171        let key_as_ident = match key {
172          PropName::Ident(ident) => Some(ident),
173          _ => None,
174        };
175
176        if let Some(key_as_string) = key_as_ident
177          && namespace_to_keep.contains(&key_as_string.sym)
178        {
179          let var_id = &match var_name.name.as_ident() {
180            Some(i) => i,
181            None => stylex_panic!("{}", VAR_DECL_NAME_NOT_IDENT),
182          }
183          .sym;
184          let key_id = NonNullProp::Atom(
185            match key_as_ident {
186              Some(i) => i,
187              None => stylex_panic!("key_as_ident is None"),
188            }
189            .clone()
190            .sym,
191          );
192
193          let all_nulls_to_keep = self
194            .state
195            .style_vars_to_keep
196            .iter()
197            .filter_map(|top_level_expression| {
198              let StyleVarsToKeep(var, namespace_name, prop) = top_level_expression;
199
200              if var_id.eq_ignore_span(var) && namespace_name.eq(&key_id.clone()) {
201                Some(prop.clone())
202              } else {
203                None
204              }
205            })
206            .collect::<Vec<NonNullProps>>();
207
208          if !all_nulls_to_keep.contains(&NonNullProps::True) {
209            let nulls_to_keep = all_nulls_to_keep
210              .into_iter()
211              .filter_map(|item| match item {
212                NonNullProps::Vec(vec) => Some(vec),
213                NonNullProps::True => None,
214              })
215              .flatten()
216              .collect::<Vec<Atom>>();
217
218            if let Some(style_object) = match prop.as_mut_key_value() {
219              Some(kv) => kv,
220              None => stylex_panic!("{}", KEY_VALUE_EXPECTED),
221            }
222            .value
223            .as_mut_object()
224            {
225              retain_style_props(style_object, nulls_to_keep);
226            }
227          }
228
229          props.push(object_prop.clone())
230        }
231      }
232    }
233
234    props
235  }
236}
237
238fn retain_style_props(style_object: &mut ObjectLit, nulls_to_keep: Vec<Atom>) {
239  style_object.props.retain(|prop| match prop {
240    PropOrSpread::Prop(prop) => {
241      let mut prop = prop.clone();
242
243      expand_shorthand_prop(&mut prop);
244
245      if let Prop::KeyValue(key_value) = &*prop
246        && key_value
247          .value
248          .as_lit()
249          .and_then(|lit| match lit {
250            Lit::Null(_) => Some(()),
251            _ => None,
252          })
253          .is_some()
254        && matches!(key_value.key, PropName::Ident(_))
255        && let PropName::Ident(ident) = &key_value.key
256      {
257        return nulls_to_keep.contains(&ident.sym);
258      }
259
260      true
261    },
262    PropOrSpread::Spread(_) => true,
263  });
264}
stylex_transform/transform/fold/fold_var_declarator.rs

stylex_transform/transform/fold/
fold_var_declarator.rs
