1use core::f64;
2
3use rustc_hash::FxHashMap;
4use swc_atoms::{atom, Atom};
5use swc_common::{Spanned, SyntaxContext, DUMMY_SP};
6use swc_ecma_ast::{
7 BinExpr, BinaryOp, Expr, Ident, Lit, Number, Str, TsEnumDecl, TsEnumMemberId, UnaryExpr,
8 UnaryOp,
9};
10use swc_ecma_utils::number::JsNumber;
11
12use super::{util::ast_ext::MemberPropExt, FastDts};
13
14#[derive(Debug, Clone)]
15enum ConstantValue {
16 Number(f64),
17 String(String),
18}
19
20impl FastDts {
21 pub(crate) fn transform_enum(&mut self, decl: &mut TsEnumDecl) {
22 let mut prev_init_value = Some(ConstantValue::Number(-1.0));
23 let mut prev_members = FxHashMap::default();
24 for member in &mut decl.members {
25 let value = if let Some(init_expr) = &member.init {
26 let computed_value = self.evaluate(init_expr, &decl.id.sym, &prev_members);
27 if computed_value.is_none() {
28 self.enum_member_initializers(member.id.span());
29 }
30 computed_value
31 } else if let Some(ConstantValue::Number(v)) = prev_init_value {
32 Some(ConstantValue::Number(v + 1.0))
33 } else {
34 None
35 };
36
37 prev_init_value = value.clone();
38 if let Some(value) = &value {
39 let member_name = match &member.id {
40 TsEnumMemberId::Ident(ident) => ident.sym.clone(),
41 TsEnumMemberId::Str(s) => s
42 .value
43 .clone()
44 .try_into_atom()
45 .unwrap_or_else(|wtf8| Atom::from(&*wtf8.to_string_lossy())),
46 #[cfg(swc_ast_unknown)]
47 _ => panic!("unable to access unknown nodes"),
48 };
49 prev_members.insert(member_name.clone(), value.clone());
50 }
51
52 member.init = value.map(|value| {
53 Box::new(match value {
54 ConstantValue::Number(v) => {
55 let is_neg = v < 0.0;
56 let expr = if v.is_infinite() {
57 Expr::Ident(Ident {
58 span: DUMMY_SP,
59 sym: atom!("Infinity"),
60 ctxt: SyntaxContext::empty(),
61 optional: false,
62 })
63 } else {
64 Expr::Lit(Lit::Num(Number {
65 span: DUMMY_SP,
66 value: v,
67 raw: Some(v.to_string().into()),
68 }))
69 };
70
71 if is_neg {
72 Expr::Unary(UnaryExpr {
73 span: DUMMY_SP,
74 arg: Box::new(expr),
75 op: UnaryOp::Minus,
76 })
77 } else {
78 expr
79 }
80 }
81 ConstantValue::String(s) => Expr::Lit(Lit::Str(Str {
82 span: DUMMY_SP,
83 value: s.clone().into(),
84 raw: None,
85 })),
86 })
87 });
88 }
89 }
90
91 fn evaluate(
92 &self,
93 expr: &Expr,
94 enum_name: &Atom,
95 prev_members: &FxHashMap<Atom, ConstantValue>,
96 ) -> Option<ConstantValue> {
97 match expr {
98 Expr::Lit(lit) => match lit {
99 Lit::Str(s) => Some(ConstantValue::String(s.value.to_string_lossy().to_string())),
100 Lit::Num(number) => Some(ConstantValue::Number(number.value)),
101 Lit::Null(_) | Lit::BigInt(_) | Lit::Bool(_) | Lit::Regex(_) | Lit::JSXText(_) => {
102 None
103 }
104 #[cfg(swc_ast_unknown)]
105 _ => panic!("unable to access unknown nodes"),
106 },
107 Expr::Tpl(tpl) => {
108 let mut value = String::new();
109 for part in &tpl.quasis {
110 value.push_str(&part.raw);
111 }
112 Some(ConstantValue::String(value))
113 }
114 Expr::Paren(expr) => self.evaluate(&expr.expr, enum_name, prev_members),
115 Expr::Bin(bin_expr) => self.evaluate_binary_expr(bin_expr, enum_name, prev_members),
116 Expr::Unary(unary_expr) => {
117 self.evaluate_unary_expr(unary_expr, enum_name, prev_members)
118 }
119 Expr::Ident(ident) => {
120 if ident.sym == "Infinity" {
121 Some(ConstantValue::Number(f64::INFINITY))
122 } else if ident.sym == "NaN" {
123 Some(ConstantValue::Number(f64::NAN))
124 } else {
125 prev_members.get(&ident.sym).cloned()
126 }
127 }
128 Expr::Member(member) => {
129 let ident = member.obj.as_ident()?;
130 if &ident.sym == enum_name {
131 let name = member.prop.static_name()?;
132 prev_members.get(name).cloned()
133 } else {
134 None
135 }
136 }
137 Expr::OptChain(opt_chain) => {
138 let member = opt_chain.base.as_member()?;
139 let ident = member.obj.as_ident()?;
140 if &ident.sym == enum_name {
141 let name = member.prop.static_name()?;
142 prev_members.get(name).cloned()
143 } else {
144 None
145 }
146 }
147 _ => None,
148 }
149 }
150
151 fn evaluate_unary_expr(
152 &self,
153 expr: &UnaryExpr,
154 enum_name: &Atom,
155 prev_members: &FxHashMap<Atom, ConstantValue>,
156 ) -> Option<ConstantValue> {
157 let value = self.evaluate(&expr.arg, enum_name, prev_members)?;
158 let value = match value {
159 ConstantValue::Number(n) => n,
160 ConstantValue::String(_) => {
161 let value = if expr.op == UnaryOp::Minus {
162 ConstantValue::Number(f64::NAN)
163 } else if expr.op == UnaryOp::Tilde {
164 ConstantValue::Number(-1.0)
165 } else {
166 value
167 };
168 return Some(value);
169 }
170 };
171
172 match expr.op {
173 UnaryOp::Minus => Some(ConstantValue::Number(-value)),
174 UnaryOp::Plus => Some(ConstantValue::Number(value)),
175 UnaryOp::Tilde => Some(ConstantValue::Number((!JsNumber::from(value)).into())),
176 _ => None,
177 }
178 }
179
180 fn evaluate_binary_expr(
181 &self,
182 expr: &BinExpr,
183 enum_name: &Atom,
184 prev_members: &FxHashMap<Atom, ConstantValue>,
185 ) -> Option<ConstantValue> {
186 let left = self.evaluate(&expr.left, enum_name, prev_members)?;
187 let right = self.evaluate(&expr.right, enum_name, prev_members)?;
188
189 if expr.op == BinaryOp::Add
190 && (matches!(left, ConstantValue::String(_))
191 || matches!(right, ConstantValue::String(_)))
192 {
193 let left_string = match left {
194 ConstantValue::Number(number) => number.to_string(),
195 ConstantValue::String(s) => s,
196 };
197
198 let right_string = match right {
199 ConstantValue::Number(number) => number.to_string(),
200 ConstantValue::String(s) => s,
201 };
202
203 return Some(ConstantValue::String(format!(
204 "{left_string}{right_string}"
205 )));
206 }
207
208 let left = JsNumber::from(match left {
209 ConstantValue::Number(n) => n,
210 ConstantValue::String(_) => return None,
211 });
212
213 let right = JsNumber::from(match right {
214 ConstantValue::Number(n) => n,
215 ConstantValue::String(_) => return None,
216 });
217
218 match expr.op {
219 BinaryOp::LShift => Some(left << right),
220 BinaryOp::RShift => Some(left >> right),
221 BinaryOp::ZeroFillRShift => Some(left.unsigned_shr(right)),
222 BinaryOp::Add => Some(left + right),
223 BinaryOp::Sub => Some(left - right),
224 BinaryOp::Mul => Some(left * right),
225 BinaryOp::Div => Some(left / right),
226 BinaryOp::Mod => Some(left & right),
227 BinaryOp::BitOr => Some(left | right),
228 BinaryOp::BitXor => Some(left ^ right),
229 BinaryOp::BitAnd => Some(left & right),
230 BinaryOp::Exp => Some(left.pow(right)),
231 _ => None,
232 }
233 .map(|number| ConstantValue::Number(number.into()))
234 }
235}