perf(transformer/jsx): no string comparisons generating pragma expression (#7620)

#7553 introduced support for some unusual JSX pragmas e.g. `this.foo` and `import.meta.foo`. We want to support these to pass tests, but they're very unlikely to be used in practice.

Identify these strange patterns when parsing the pragma (which happens only once per file), and encode them as an enum. The removes expensive string comparisons from `Pragma::create_expression` (which is called for every JSX element), and keeps the path for common cases fast.

crates/oxc_transformer/src/jsx/jsx_impl.rs

@@ -312,9 +312,20 @@ fn get_import_source(jsx_runtime_importer: &str, react_importer_len: u32) -> Ato
     Atom::from(&jsx_runtime_importer[..react_importer_len as usize])
 }
 
-/// Pragma used in classic mode
-struct Pragma<'a> {
-    ident_parts: Vec<Atom<'a>>,
+/// Pragma used in classic mode.
+///
+/// `Double` is first as it's most common.
+enum Pragma<'a> {
+    /// `React.createElement`
+    Double(Atom<'a>, Atom<'a>),
+    /// `createElement`
+    Single(Atom<'a>),
+    /// `foo.bar.qux`
+    Multiple(Vec<Atom<'a>>),
+    /// `this`, `this.foo`, `this.foo.bar.qux`
+    This(Vec<Atom<'a>>),
+    /// `import.meta`, `import.meta.foo`, `import.meta.foo.bar.qux`
+    ImportMeta(Vec<Atom<'a>>),
 }
 
 impl<'a> Pragma<'a> {
@@ -329,43 +340,64 @@ impl<'a> Pragma<'a> {
     ) -> Self {
         if let Some(pragma) = pragma {
             if pragma.is_empty() {
-                return Self::invalid(default_property_name, ctx);
+                ctx.error(diagnostics::invalid_pragma());
+            } else {
+                return Self::parse_impl(pragma, ast);
             }
+        }
 
-            let ident_parts = pragma.split('.').map(|item| ast.atom(item)).collect::<Vec<_>>();
-            Self { ident_parts }
-        } else {
-            Self::default(default_property_name)
+        Self::Double(Atom::from("React"), Atom::from(default_property_name))
+    }
+
+    fn parse_impl(pragma: &str, ast: AstBuilder<'a>) -> Self {
+        let strs_to_atoms = |parts: &[&str]| parts.iter().map(|part| ast.atom(part)).collect();
+
+        let parts = pragma.split('.').collect::<Vec<_>>();
+        match &parts[..] {
+            [] => unreachable!(),
+            ["this", rest @ ..] => Self::This(strs_to_atoms(rest)),
+            ["import", "meta", rest @ ..] => Self::ImportMeta(strs_to_atoms(rest)),
+            [first, second] => Self::Double(ast.atom(first), ast.atom(second)),
+            [only] => Self::Single(ast.atom(only)),
+            parts => Self::Multiple(strs_to_atoms(parts)),
         }
     }
 
-    fn invalid(default_property_name: &'static str, ctx: &TransformCtx<'a>) -> Self {
-        ctx.error(diagnostics::invalid_pragma());
-        Self::default(default_property_name)
-    }
-
-    fn default(default_property_name: &'static str) -> Self {
-        Self { ident_parts: vec![Atom::from("React"), Atom::from(default_property_name)] }
-    }
-
     fn create_expression(&self, ctx: &mut TraverseCtx<'a>) -> Expression<'a> {
-        let (mut expr, parts) = match &self.ident_parts[..] {
-            [] => unreachable!(),
-            [atom, rest @ ..] if atom == "this" => (ctx.ast.expression_this(SPAN), rest),
-            [first, second, rest @ ..] if first == "import" && second == "meta" => (
-                ctx.ast.expression_meta_property(
-                    SPAN,
-                    ctx.ast.identifier_name(SPAN, first),
-                    ctx.ast.identifier_name(SPAN, second),
-                ),
-                rest,
-            ),
-            [first, rest @ ..] => {
+        let (object, parts) = match self {
+            Self::Double(first, second) => {
                 let object = get_read_identifier_reference(SPAN, first.clone(), ctx);
-                let expr = Expression::Identifier(ctx.alloc(object));
-                (expr, rest)
+                return Expression::from(ctx.ast.member_expression_static(
+                    SPAN,
+                    object,
+                    ctx.ast.identifier_name(SPAN, second.clone()),
+                    false,
+                ));
+            }
+            Self::Single(single) => {
+                return get_read_identifier_reference(SPAN, single.clone(), ctx);
+            }
+            Self::Multiple(parts) => {
+                let mut parts = parts.iter();
+                let first = parts.next().unwrap().clone();
+                let object = get_read_identifier_reference(SPAN, first, ctx);
+                (object, parts)
+            }
+            Self::This(parts) => {
+                let object = ctx.ast.expression_this(SPAN);
+                (object, parts.iter())
+            }
+            Self::ImportMeta(parts) => {
+                let object = ctx.ast.expression_meta_property(
+                    SPAN,
+                    ctx.ast.identifier_name(SPAN, Atom::from("import")),
+                    ctx.ast.identifier_name(SPAN, Atom::from("meta")),
+                );
+                (object, parts.iter())
             }
         };
+
+        let mut expr = object;
         for item in parts {
             let name = ctx.ast.identifier_name(SPAN, item.clone());
             expr = ctx.ast.member_expression_static(SPAN, expr, name, false).into();
@@ -1049,10 +1081,11 @@ fn get_read_identifier_reference<'a>(
     span: Span,
     name: Atom<'a>,
     ctx: &mut TraverseCtx<'a>,
-) -> IdentifierReference<'a> {
+) -> Expression<'a> {
     let reference_id =
         ctx.create_reference_in_current_scope(name.to_compact_str(), ReferenceFlags::Read);
-    ctx.ast.identifier_reference_with_reference_id(span, name, reference_id)
+    let ident = ctx.ast.alloc_identifier_reference_with_reference_id(span, name, reference_id);
+    Expression::Identifier(ident)
 }
 
 fn create_static_member_expression<'a>(
@@ -1076,6 +1109,7 @@ mod test {
     use oxc_allocator::Allocator;
     use oxc_ast::ast::Expression;
     use oxc_semantic::{ScopeTree, SymbolTable};
+    use oxc_syntax::{node::NodeId, scope::ScopeFlags};
     use oxc_traverse::ReusableTraverseCtx;
 
     use crate::{TransformCtx, TransformOptions};
@@ -1085,18 +1119,91 @@ mod test {
     macro_rules! setup {
         ($traverse_ctx:ident, $transform_ctx:ident) => {
             let allocator = Allocator::default();
-            let traverse_ctx =
-                ReusableTraverseCtx::new(ScopeTree::default(), SymbolTable::default(), &allocator);
+
+            let mut scopes = ScopeTree::default();
+            scopes.add_scope(None, NodeId::DUMMY, ScopeFlags::Top);
+
+            let traverse_ctx = ReusableTraverseCtx::new(scopes, SymbolTable::default(), &allocator);
             // SAFETY: Macro user only gets a `&mut TraverseCtx`, which cannot be abused
             let mut traverse_ctx = unsafe { traverse_ctx.unwrap() };
             let $traverse_ctx = &mut traverse_ctx;
+
             let $transform_ctx =
                 TransformCtx::new(Path::new("test.jsx"), &TransformOptions::default());
         };
     }
 
     #[test]
-    fn this_expr_pragma() {
+    fn default_pragma() {
+        setup!(traverse_ctx, transform_ctx);
+
+        let pragma = None;
+        let pragma = Pragma::parse(pragma, "createElement", traverse_ctx.ast, &transform_ctx);
+        let expr = pragma.create_expression(traverse_ctx);
+
+        let Expression::StaticMemberExpression(member) = &expr else { panic!() };
+        let Expression::Identifier(object) = &member.object else { panic!() };
+        assert_eq!(object.name, "React");
+        assert_eq!(member.property.name, "createElement");
+    }
+
+    #[test]
+    fn single_part_pragma() {
+        setup!(traverse_ctx, transform_ctx);
+
+        let pragma = Some("single");
+        let pragma = Pragma::parse(pragma, "createElement", traverse_ctx.ast, &transform_ctx);
+        let expr = pragma.create_expression(traverse_ctx);
+
+        let Expression::Identifier(ident) = &expr else { panic!() };
+        assert_eq!(ident.name, "single");
+    }
+
+    #[test]
+    fn two_part_pragma() {
+        setup!(traverse_ctx, transform_ctx);
+
+        let pragma = Some("first.second");
+        let pragma = Pragma::parse(pragma, "createElement", traverse_ctx.ast, &transform_ctx);
+        let expr = pragma.create_expression(traverse_ctx);
+
+        let Expression::StaticMemberExpression(member) = &expr else { panic!() };
+        let Expression::Identifier(object) = &member.object else { panic!() };
+        assert_eq!(object.name, "first");
+        assert_eq!(member.property.name, "second");
+    }
+
+    #[test]
+    fn multi_part_pragma() {
+        setup!(traverse_ctx, transform_ctx);
+
+        let pragma = Some("first.second.third");
+        let pragma = Pragma::parse(pragma, "createElement", traverse_ctx.ast, &transform_ctx);
+        let expr = pragma.create_expression(traverse_ctx);
+
+        let Expression::StaticMemberExpression(outer_member) = &expr else { panic!() };
+        let Expression::StaticMemberExpression(inner_member) = &outer_member.object else {
+            panic!()
+        };
+        let Expression::Identifier(object) = &inner_member.object else { panic!() };
+        assert_eq!(object.name, "first");
+        assert_eq!(inner_member.property.name, "second");
+        assert_eq!(outer_member.property.name, "third");
+    }
+
+    #[test]
+    fn this_pragma() {
+        setup!(traverse_ctx, transform_ctx);
+
+        let pragma = Some("this");
+        let pragma = Pragma::parse(pragma, "createElement", traverse_ctx.ast, &transform_ctx);
+        let expr = pragma.create_expression(traverse_ctx);
+
+        assert!(matches!(&expr, Expression::ThisExpression(_)));
+    }
+
+    #[test]
+    fn this_prop_pragma() {
         setup!(traverse_ctx, transform_ctx);
 
         let pragma = Some("this.a.b");
@@ -1116,6 +1223,19 @@ mod test {
     fn import_meta_pragma() {
         setup!(traverse_ctx, transform_ctx);
 
+        let pragma = Some("import.meta");
+        let pragma = Pragma::parse(pragma, "createElement", traverse_ctx.ast, &transform_ctx);
+        let expr = pragma.create_expression(traverse_ctx);
+
+        let Expression::MetaProperty(meta_prop) = &expr else { panic!() };
+        assert_eq!(&meta_prop.meta.name, "import");
+        assert_eq!(&meta_prop.property.name, "meta");
+    }
+
+    #[test]
+    fn import_meta_prop_pragma() {
+        setup!(traverse_ctx, transform_ctx);
+
         let pragma = Some("import.meta.prop");
         let pragma = Pragma::parse(pragma, "createElement", traverse_ctx.ast, &transform_ctx);
         let expr = pragma.create_expression(traverse_ctx);