swc_ecma_transforms_typescript/
ts_enum.rs

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

swc_ecma_transforms_typescript/
ts_enum.rs
