swc_ecma_transforms_typescript/
ts_enum.rs

1use rustc_hash::FxHashMap;
2use swc_atoms::{atom, Atom, Wtf8Atom};
3use swc_common::{SyntaxContext, DUMMY_SP};
4use swc_ecma_ast::*;
5use swc_ecma_utils::{
6    number::{JsNumber, ToJsString},
7    ExprFactory,
8};
9use swc_ecma_visit::{noop_visit_mut_type, VisitMut, VisitMutWith};
10
11#[inline]
12fn atom_from_wtf8_atom(value: &Wtf8Atom) -> Atom {
13    value
14        .as_str()
15        .map(Atom::from)
16        .unwrap_or_else(|| Atom::from(value.to_string_lossy()))
17}
18
19#[derive(Debug, Clone, PartialEq, Eq, Hash)]
20pub(crate) struct TsEnumRecordKey {
21    pub enum_id: Id,
22    pub member_name: Atom,
23}
24
25pub(crate) type TsEnumRecord = FxHashMap<TsEnumRecordKey, TsEnumRecordValue>;
26
27#[derive(Debug, Clone)]
28pub(crate) enum TsEnumRecordValue {
29    String(Atom),
30    Number(JsNumber),
31    Opaque(Box<Expr>),
32    Void,
33}
34
35impl TsEnumRecordValue {
36    pub fn inc(&self) -> Self {
37        match self {
38            Self::Number(num) => Self::Number((**num + 1.0).into()),
39            _ => Self::Void,
40        }
41    }
42
43    pub fn is_const(&self) -> bool {
44        matches!(
45            self,
46            TsEnumRecordValue::String(..) | TsEnumRecordValue::Number(..)
47        )
48    }
49
50    pub fn is_string(&self) -> bool {
51        matches!(self, TsEnumRecordValue::String(..))
52    }
53
54    pub fn has_value(&self) -> bool {
55        !matches!(self, TsEnumRecordValue::Void)
56    }
57}
58
59impl From<TsEnumRecordValue> for Expr {
60    fn from(value: TsEnumRecordValue) -> Self {
61        match value {
62            TsEnumRecordValue::String(string) => Lit::Str(string.into()).into(),
63            TsEnumRecordValue::Number(num) if num.is_nan() => Ident {
64                span: DUMMY_SP,
65                sym: atom!("NaN"),
66                ..Default::default()
67            }
68            .into(),
69            TsEnumRecordValue::Number(num) if num.is_infinite() => {
70                let value: Expr = Ident {
71                    span: DUMMY_SP,
72                    sym: atom!("Infinity"),
73                    ..Default::default()
74                }
75                .into();
76
77                if num.is_sign_negative() {
78                    UnaryExpr {
79                        span: DUMMY_SP,
80                        op: op!(unary, "-"),
81                        arg: value.into(),
82                    }
83                    .into()
84                } else {
85                    value
86                }
87            }
88            TsEnumRecordValue::Number(num) => Lit::Num(Number {
89                span: DUMMY_SP,
90                value: *num,
91                raw: None,
92            })
93            .into(),
94            TsEnumRecordValue::Void => *Expr::undefined(DUMMY_SP),
95            TsEnumRecordValue::Opaque(expr) => *expr,
96        }
97    }
98}
99
100impl From<f64> for TsEnumRecordValue {
101    fn from(value: f64) -> Self {
102        Self::Number(value.into())
103    }
104}
105
106pub(crate) struct EnumValueComputer<'a> {
107    pub enum_id: &'a Id,
108    pub unresolved_ctxt: SyntaxContext,
109    pub record: &'a TsEnumRecord,
110}
111
112/// https://github.com/microsoft/TypeScript/pull/50528
113impl EnumValueComputer<'_> {
114    pub fn compute(&mut self, expr: Box<Expr>) -> TsEnumRecordValue {
115        let mut expr = self.compute_rec(expr);
116        if let TsEnumRecordValue::Opaque(expr) = &mut expr {
117            expr.visit_mut_with(self);
118        }
119        expr
120    }
121
122    fn compute_rec(&self, expr: Box<Expr>) -> TsEnumRecordValue {
123        match *expr {
124            Expr::Lit(Lit::Str(s)) => TsEnumRecordValue::String(atom_from_wtf8_atom(&s.value)),
125            Expr::Lit(Lit::Num(n)) => TsEnumRecordValue::Number(n.value.into()),
126            Expr::Ident(Ident { ctxt, sym, .. })
127                if &*sym == "NaN" && ctxt == self.unresolved_ctxt =>
128            {
129                TsEnumRecordValue::Number(f64::NAN.into())
130            }
131            Expr::Ident(Ident { ctxt, sym, .. })
132                if &*sym == "Infinity" && ctxt == self.unresolved_ctxt =>
133            {
134                TsEnumRecordValue::Number(f64::INFINITY.into())
135            }
136            Expr::Ident(ref ident) => self
137                .record
138                .get(&TsEnumRecordKey {
139                    enum_id: self.enum_id.clone(),
140                    member_name: ident.sym.clone(),
141                })
142                .cloned()
143                .map(|value| match value {
144                    TsEnumRecordValue::String(..) | TsEnumRecordValue::Number(..) => value,
145                    _ => TsEnumRecordValue::Opaque(
146                        self.enum_id
147                            .clone()
148                            .make_member(ident.clone().into())
149                            .into(),
150                    ),
151                })
152                .unwrap_or_else(|| TsEnumRecordValue::Opaque(expr)),
153            Expr::Paren(e) => self.compute_rec(e.expr),
154            Expr::Unary(e) => self.compute_unary(e),
155            Expr::Bin(e) => self.compute_bin(e),
156            Expr::Member(e) => self.compute_member(e),
157            Expr::Tpl(e) => self.compute_tpl(e),
158            _ => TsEnumRecordValue::Opaque(expr),
159        }
160    }
161
162    fn compute_unary(&self, expr: UnaryExpr) -> TsEnumRecordValue {
163        if !matches!(expr.op, op!(unary, "+") | op!(unary, "-") | op!("~")) {
164            return TsEnumRecordValue::Opaque(expr.into());
165        }
166
167        let inner = self.compute_rec(expr.arg);
168
169        let TsEnumRecordValue::Number(num) = inner else {
170            return TsEnumRecordValue::Opaque(
171                UnaryExpr {
172                    span: expr.span,
173                    op: expr.op,
174                    arg: Box::new(inner.into()),
175                }
176                .into(),
177            );
178        };
179
180        match expr.op {
181            op!(unary, "+") => TsEnumRecordValue::Number(num),
182            op!(unary, "-") => TsEnumRecordValue::Number(-num),
183            op!("~") => TsEnumRecordValue::Number(!num),
184            _ => unreachable!(),
185        }
186    }
187
188    fn compute_bin(&self, expr: BinExpr) -> TsEnumRecordValue {
189        let origin_expr = expr.clone();
190        if !matches!(
191            expr.op,
192            op!(bin, "+")
193                | op!(bin, "-")
194                | op!("*")
195                | op!("/")
196                | op!("%")
197                | op!("**")
198                | op!("<<")
199                | op!(">>")
200                | op!(">>>")
201                | op!("|")
202                | op!("&")
203                | op!("^"),
204        ) {
205            return TsEnumRecordValue::Opaque(origin_expr.into());
206        }
207
208        let left = self.compute_rec(expr.left);
209        let right = self.compute_rec(expr.right);
210
211        match (left, right, expr.op) {
212            (TsEnumRecordValue::Number(left), TsEnumRecordValue::Number(right), op) => {
213                let value = match op {
214                    op!(bin, "+") => left + right,
215                    op!(bin, "-") => left - right,
216                    op!("*") => left * right,
217                    op!("/") => left / right,
218                    op!("%") => left % right,
219                    op!("**") => left.pow(right),
220                    op!("<<") => left << right,
221                    op!(">>") => left >> right,
222                    op!(">>>") => left.unsigned_shr(right),
223                    op!("|") => left | right,
224                    op!("&") => left & right,
225                    op!("^") => left ^ right,
226                    _ => unreachable!(),
227                };
228
229                TsEnumRecordValue::Number(value)
230            }
231            (TsEnumRecordValue::String(left), TsEnumRecordValue::String(right), op!(bin, "+")) => {
232                TsEnumRecordValue::String(format!("{left}{right}").into())
233            }
234            (TsEnumRecordValue::Number(left), TsEnumRecordValue::String(right), op!(bin, "+")) => {
235                let left = left.to_js_string();
236
237                TsEnumRecordValue::String(format!("{left}{right}").into())
238            }
239            (TsEnumRecordValue::String(left), TsEnumRecordValue::Number(right), op!(bin, "+")) => {
240                let right = right.to_js_string();
241
242                TsEnumRecordValue::String(format!("{left}{right}").into())
243            }
244            (left, right, _) => {
245                let mut origin_expr = origin_expr;
246
247                if left.is_const() {
248                    origin_expr.left = Box::new(left.into());
249                }
250
251                if right.is_const() {
252                    origin_expr.right = Box::new(right.into());
253                }
254
255                TsEnumRecordValue::Opaque(origin_expr.into())
256            }
257        }
258    }
259
260    fn compute_member(&self, expr: MemberExpr) -> TsEnumRecordValue {
261        if matches!(expr.prop, MemberProp::PrivateName(..)) {
262            return TsEnumRecordValue::Opaque(expr.into());
263        }
264
265        let opaque_expr = TsEnumRecordValue::Opaque(expr.clone().into());
266
267        let member_name = match expr.prop {
268            MemberProp::Ident(ident) => ident.sym,
269            MemberProp::Computed(ComputedPropName { expr, .. }) => {
270                let Expr::Lit(Lit::Str(s)) = *expr else {
271                    return opaque_expr;
272                };
273
274                atom_from_wtf8_atom(&s.value)
275            }
276            _ => return opaque_expr,
277        };
278
279        let Expr::Ident(ident) = *expr.obj else {
280            return opaque_expr;
281        };
282
283        self.record
284            .get(&TsEnumRecordKey {
285                enum_id: ident.to_id(),
286                member_name,
287            })
288            .cloned()
289            .filter(TsEnumRecordValue::has_value)
290            .unwrap_or(opaque_expr)
291    }
292
293    fn compute_tpl(&self, expr: Tpl) -> TsEnumRecordValue {
294        let opaque_expr = TsEnumRecordValue::Opaque(expr.clone().into());
295
296        let Tpl { exprs, quasis, .. } = expr;
297
298        let mut quasis_iter = quasis.into_iter();
299
300        let Some(mut string) = quasis_iter.next().map(|q| q.raw.to_string()) else {
301            return opaque_expr;
302        };
303
304        for (q, expr) in quasis_iter.zip(exprs) {
305            let expr = self.compute_rec(expr);
306
307            let expr = match expr {
308                TsEnumRecordValue::String(s) => s.to_string(),
309                TsEnumRecordValue::Number(n) => n.to_js_string(),
310                _ => return opaque_expr,
311            };
312
313            string.push_str(&expr);
314            string.push_str(&q.raw);
315        }
316
317        TsEnumRecordValue::String(string.into())
318    }
319}
320
321impl VisitMut for EnumValueComputer<'_> {
322    noop_visit_mut_type!();
323
324    fn visit_mut_expr(&mut self, expr: &mut Expr) {
325        expr.visit_mut_children_with(self);
326
327        let Expr::Ident(ident) = expr else { return };
328
329        if self.record.contains_key(&TsEnumRecordKey {
330            enum_id: self.enum_id.clone(),
331            member_name: ident.sym.clone(),
332        }) {
333            *expr = self
334                .enum_id
335                .clone()
336                .make_member(ident.clone().into())
337                .into();
338        }
339    }
340}
swc_ecma_transforms_typescript/ts_enum.rs

swc_ecma_transforms_typescript/
ts_enum.rs
