feat(transformer-dts): transform enum support (#3710)

Author: Dunqing

crates/oxc_ast/src/ast_builder.rs

@@ -1899,6 +1899,16 @@ impl<'a> AstBuilder<'a> {
         )
     }
 
+    #[inline]
+    pub fn ts_enum_member(
+        self,
+        span: Span,
+        id: TSEnumMemberName<'a>,
+        initializer: Option<Expression<'a>>,
+    ) -> TSEnumMember<'a> {
+        TSEnumMember { span, id, initializer }
+    }
+
     #[inline]
     pub fn decorator(self, span: Span, expression: Expression<'a>) -> Decorator<'a> {
         Decorator { span, expression }

crates/oxc_ast/src/span.rs

@@ -732,3 +732,14 @@ impl<'a> GetSpan for JSXMemberExpressionObject<'a> {
         }
     }
 }
+
+impl<'a> GetSpan for TSEnumMemberName<'a> {
+    fn span(&self) -> Span {
+        match self {
+            TSEnumMemberName::StaticIdentifier(ident) => ident.span,
+            TSEnumMemberName::StaticStringLiteral(literal) => literal.span,
+            TSEnumMemberName::StaticNumericLiteral(literal) => literal.span,
+            expr @ match_expression!(TSEnumMemberName) => expr.to_expression().span(),
+        }
+    }
+}

crates/oxc_transformer_dts/Cargo.toml

@@ -25,7 +25,7 @@ oxc_ast         = { workspace = true }
 oxc_span        = { workspace = true }
 oxc_allocator   = { workspace = true }
 oxc_diagnostics = { workspace = true }
-oxc_syntax      = { workspace = true }
+oxc_syntax      = { workspace = true, features = ["to_js_string"] }
 
 rustc-hash = { workspace = true }
 

crates/oxc_transformer_dts/src/declaration.rs

@@ -213,7 +213,7 @@ impl<'a> TransformerDts<'a> {
             }
             Declaration::TSEnumDeclaration(enum_decl) => {
                 if !check_binding || self.scope.has_reference(&enum_decl.id.name) {
-                    Some(self.ctx.ast.copy(decl))
+                    self.transform_ts_enum_declaration(enum_decl)
                 } else {
                     None
                 }

crates/oxc_transformer_dts/src/diagnostics.rs

@@ -31,3 +31,8 @@ pub fn signature_computed_property_name(span: Span) -> OxcDiagnostic {
     OxcDiagnostic::error("Computed properties must be number or string literals, variables or dotted expressions with --isolatedDeclarations.")
         .with_label(span)
 }
+
+pub fn enum_member_initializers(span: Span) -> OxcDiagnostic {
+    OxcDiagnostic::error("Enum member initializers must be computable without references to external symbols with --isolatedDeclarations.")
+        .with_label(span)
+}

crates/oxc_transformer_dts/src/enum.rs

@@ -0,0 +1,282 @@
+#[allow(clippy::wildcard_imports)]
+use oxc_ast::ast::*;
+
+use oxc_span::{Atom, GetSpan, SPAN};
+use oxc_syntax::{
+    number::{NumberBase, ToJsInt32, ToJsString},
+    operator::{BinaryOperator, UnaryOperator},
+};
+use rustc_hash::FxHashMap;
+
+use crate::{diagnostics::enum_member_initializers, TransformerDts};
+
+#[derive(Debug, Clone)]
+enum ConstantValue {
+    Number(f64),
+    String(String),
+}
+
+impl<'a> TransformerDts<'a> {
+    pub fn transform_ts_enum_declaration(
+        &mut self,
+        decl: &TSEnumDeclaration<'a>,
+    ) -> Option<Declaration<'a>> {
+        let mut members = self.ctx.ast.new_vec();
+        let mut prev_initializer_value = Some(ConstantValue::Number(0.0));
+        let mut prev_members = FxHashMap::default();
+        for member in &decl.members {
+            let value = if let Some(initializer) = &member.initializer {
+                let computed_value =
+                    self.computed_constant_value(initializer, &decl.id.name, &prev_members);
+
+                if computed_value.is_none() {
+                    self.ctx.error(enum_member_initializers(member.id.span()));
+                }
+
+                computed_value
+            } else if let Some(ConstantValue::Number(v)) = prev_initializer_value {
+                Some(ConstantValue::Number(v + 1.0))
+            } else {
+                None
+            };
+
+            prev_initializer_value.clone_from(&value);
+
+            if let Some(value) = &value {
+                let member_name = match &member.id {
+                    TSEnumMemberName::StaticIdentifier(id) => &id.name,
+                    TSEnumMemberName::StaticStringLiteral(str) => &str.value,
+                    #[allow(clippy::unnested_or_patterns)] // Clippy is wrong
+                    TSEnumMemberName::StaticNumericLiteral(_)
+                    | match_expression!(TSEnumMemberName) => {
+                        unreachable!()
+                    }
+                };
+                prev_members.insert(member_name.clone(), value.clone());
+            }
+
+            let member = self.ctx.ast.ts_enum_member(
+                member.span,
+                self.ctx.ast.copy(&member.id),
+                value.map(|v| match v {
+                    ConstantValue::Number(v) => {
+                        let is_negative = v < 0.0;
+
+                        // Infinity
+                        let expr = if v.is_infinite() {
+                            let ident =
+                                IdentifierReference::new(SPAN, self.ctx.ast.new_atom("Infinity"));
+                            self.ctx.ast.identifier_reference_expression(ident)
+                        } else {
+                            let value = if is_negative { -v } else { v };
+                            self.ctx.ast.literal_number_expression(NumericLiteral {
+                                span: SPAN,
+                                value,
+                                raw: self.ctx.ast.new_str(&value.to_string()),
+                                base: NumberBase::Decimal,
+                            })
+                        };
+
+                        if is_negative {
+                            self.ctx.ast.unary_expression(SPAN, UnaryOperator::UnaryNegation, expr)
+                        } else {
+                            expr
+                        }
+                    }
+                    ConstantValue::String(v) => self
+                        .ctx
+                        .ast
+                        .literal_string_expression(self.ctx.ast.string_literal(SPAN, &v)),
+                }),
+            );
+
+            members.push(member);
+        }
+        Some(self.ctx.ast.ts_enum_declaration(
+            decl.span,
+            self.ctx.ast.copy(&decl.id),
+            members,
+            self.modifiers_declare(),
+        ))
+    }
+
+    /// Evaluate the expression to a constant value.
+    /// Refer to [babel](https://github.com/babel/babel/blob/610897a9a96c5e344e77ca9665df7613d2f88358/packages/babel-plugin-transform-typescript/src/enum.ts#L241C1-L394C2)
+    fn computed_constant_value(
+        &self,
+        expr: &Expression<'a>,
+        enum_name: &Atom<'a>,
+        prev_members: &FxHashMap<Atom<'a>, ConstantValue>,
+    ) -> Option<ConstantValue> {
+        self.evaluate(expr, enum_name, prev_members)
+    }
+
+    #[allow(clippy::unused_self)]
+    fn evaluate_ref(
+        &self,
+        expr: &Expression<'a>,
+        enum_name: &Atom<'a>,
+        prev_members: &FxHashMap<Atom<'a>, ConstantValue>,
+    ) -> Option<ConstantValue> {
+        match expr {
+            match_member_expression!(Expression) => {
+                let expr = expr.to_member_expression();
+                let Expression::Identifier(ident) = expr.object() else { return None };
+                if ident.name == enum_name {
+                    let property = expr.static_property_name()?;
+                    prev_members.get(property).cloned()
+                } else {
+                    None
+                }
+            }
+            Expression::Identifier(ident) => {
+                if ident.name == "Infinity" {
+                    return Some(ConstantValue::Number(f64::INFINITY));
+                } else if ident.name == "NaN" {
+                    return Some(ConstantValue::Number(f64::NAN));
+                }
+
+                if let Some(value) = prev_members.get(&ident.name) {
+                    return Some(value.clone());
+                }
+
+                None
+            }
+            _ => None,
+        }
+    }
+
+    fn evaluate(
+        &self,
+        expr: &Expression<'a>,
+        enum_name: &Atom<'a>,
+        prev_members: &FxHashMap<Atom<'a>, ConstantValue>,
+    ) -> Option<ConstantValue> {
+        match expr {
+            Expression::Identifier(_)
+            | Expression::ComputedMemberExpression(_)
+            | Expression::StaticMemberExpression(_)
+            | Expression::PrivateFieldExpression(_) => {
+                self.evaluate_ref(expr, enum_name, prev_members)
+            }
+            Expression::BinaryExpression(expr) => {
+                self.eval_binary_expression(expr, enum_name, prev_members)
+            }
+            Expression::UnaryExpression(expr) => {
+                self.eval_unary_expression(expr, enum_name, prev_members)
+            }
+            Expression::NumericLiteral(lit) => Some(ConstantValue::Number(lit.value)),
+            Expression::StringLiteral(lit) => Some(ConstantValue::String(lit.value.to_string())),
+            Expression::TemplateLiteral(lit) => {
+                let mut value = String::new();
+                for part in &lit.quasis {
+                    value.push_str(&part.value.raw);
+                }
+                Some(ConstantValue::String(value))
+            }
+            Expression::ParenthesizedExpression(expr) => {
+                self.evaluate(&expr.expression, enum_name, prev_members)
+            }
+            _ => None,
+        }
+    }
+
+    #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss, clippy::cast_sign_loss)]
+    fn eval_binary_expression(
+        &self,
+        expr: &BinaryExpression<'a>,
+        enum_name: &Atom<'a>,
+        prev_members: &FxHashMap<Atom<'a>, ConstantValue>,
+    ) -> Option<ConstantValue> {
+        let left = self.evaluate(&expr.left, enum_name, prev_members)?;
+        let right = self.evaluate(&expr.right, enum_name, prev_members)?;
+
+        if matches!(expr.operator, BinaryOperator::Addition)
+            && (matches!(left, ConstantValue::String(_))
+                || matches!(right, ConstantValue::String(_)))
+        {
+            let left_string = match left {
+                ConstantValue::String(str) => str,
+                ConstantValue::Number(v) => v.to_js_string(),
+            };
+
+            let right_string = match right {
+                ConstantValue::String(str) => str,
+                ConstantValue::Number(v) => v.to_js_string(),
+            };
+
+            return Some(ConstantValue::String(format!("{left_string}{right_string}")));
+        }
+
+        let left = match left {
+            ConstantValue::Number(v) => v,
+            ConstantValue::String(_) => return None,
+        };
+
+        let right = match right {
+            ConstantValue::Number(v) => v,
+            ConstantValue::String(_) => return None,
+        };
+
+        match expr.operator {
+            BinaryOperator::ShiftRight => Some(ConstantValue::Number(f64::from(
+                left.to_js_int_32().wrapping_shr(right.to_js_int_32() as u32),
+            ))),
+            BinaryOperator::ShiftRightZeroFill => Some(ConstantValue::Number(f64::from(
+                (left.to_js_int_32() as u32).wrapping_shr(right.to_js_int_32() as u32),
+            ))),
+            BinaryOperator::ShiftLeft => Some(ConstantValue::Number(f64::from(
+                left.to_js_int_32().wrapping_shl(right.to_js_int_32() as u32),
+            ))),
+            BinaryOperator::BitwiseXOR => {
+                Some(ConstantValue::Number(f64::from(left.to_js_int_32() ^ right.to_js_int_32())))
+            }
+            BinaryOperator::BitwiseOR => {
+                Some(ConstantValue::Number(f64::from(left.to_js_int_32() | right.to_js_int_32())))
+            }
+            BinaryOperator::BitwiseAnd => {
+                Some(ConstantValue::Number(f64::from(left.to_js_int_32() & right.to_js_int_32())))
+            }
+            BinaryOperator::Multiplication => Some(ConstantValue::Number(left * right)),
+            BinaryOperator::Division => Some(ConstantValue::Number(left / right)),
+            BinaryOperator::Addition => Some(ConstantValue::Number(left + right)),
+            BinaryOperator::Subtraction => Some(ConstantValue::Number(left - right)),
+            BinaryOperator::Remainder => Some(ConstantValue::Number(left % right)),
+            BinaryOperator::Exponential => Some(ConstantValue::Number(left.powf(right))),
+            _ => None,
+        }
+    }
+
+    #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
+    fn eval_unary_expression(
+        &self,
+        expr: &UnaryExpression<'a>,
+        enum_name: &Atom<'a>,
+        prev_members: &FxHashMap<Atom<'a>, ConstantValue>,
+    ) -> Option<ConstantValue> {
+        let value = self.evaluate(&expr.argument, enum_name, prev_members)?;
+
+        let value = match value {
+            ConstantValue::Number(value) => value,
+            ConstantValue::String(_) => {
+                let value = if expr.operator == UnaryOperator::UnaryNegation {
+                    ConstantValue::Number(f64::NAN)
+                } else if expr.operator == UnaryOperator::BitwiseNot {
+                    ConstantValue::Number(-1.0)
+                } else {
+                    value
+                };
+                return Some(value);
+            }
+        };
+
+        match expr.operator {
+            UnaryOperator::UnaryPlus => Some(ConstantValue::Number(value)),
+            UnaryOperator::UnaryNegation => Some(ConstantValue::Number(-value)),
+            UnaryOperator::BitwiseNot => {
+                Some(ConstantValue::Number(f64::from(!value.to_js_int_32())))
+            }
+            _ => None,
+        }
+    }
+}

crates/oxc_transformer_dts/src/lib.rs

@@ -9,6 +9,7 @@ mod class;
 mod context;
 mod declaration;
 mod diagnostics;
+mod r#enum;
 mod function;
 mod inferrer;
 mod module;