refactor: call loadWorkletRuntime once in each module

.changeset/sad-terms-reply.md

@@ -0,0 +1,5 @@
+---
+"@lynx-js/react": patch
+---
+
+refactor: call loadWorkletRuntime once in each module

packages/react/transform/__test__/fixture.spec.js

@@ -1607,7 +1607,8 @@ export function bar() {
               },
               _wkltId: "da39:75a1b:1"
           };
-          loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry) && registerWorkletInternal("main-thread", "da39:75a1b:1", function(event) {
+          const __workletRuntimeLoaded = loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry);
+          __workletRuntimeLoaded && registerWorkletInternal("main-thread", "da39:75a1b:1", function(event) {
               const getCurrentDelta = lynxWorkletImpl._workletMap["da39:75a1b:1"].bind(this);
               let { foo } = this["_c"];
               "main thread";
@@ -1643,7 +1644,8 @@ export function bar() {
               },
               _wkltId: "da39:75a1b:1"
           };
-          loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry) && registerWorkletInternal("main-thread", "da39:75a1b:1", function(event) {
+          const __workletRuntimeLoaded = loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry);
+          __workletRuntimeLoaded && registerWorkletInternal("main-thread", "da39:75a1b:1", function(event) {
               const getCurrentDelta = lynxWorkletImpl._workletMap["da39:75a1b:1"].bind(this);
               let { foo } = this["_c"];
               "main thread";
@@ -1708,7 +1710,8 @@ export function foo(event) {
           },
           _wkltId: "da39:64631:1"
       };
-      loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry) && registerWorkletInternal("main-thread", "da39:64631:1", function(event) {
+      const __workletRuntimeLoaded = loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry);
+      __workletRuntimeLoaded && registerWorkletInternal("main-thread", "da39:64631:1", function(event) {
           const foo = lynxWorkletImpl._workletMap["da39:64631:1"].bind(this);
           let { bar, qux } = this["_c"];
           "main thread";
@@ -1768,12 +1771,13 @@ console.log(bar)
           _wkltId: "da39:80ef4:2"
       };
       console.log(bar);
-      loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry) && registerWorkletInternal("main-thread", "da39:80ef4:1", function() {
+      const __workletRuntimeLoaded = loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry);
+      __workletRuntimeLoaded && registerWorkletInternal("main-thread", "da39:80ef4:1", function() {
           const foo = lynxWorkletImpl._workletMap["da39:80ef4:1"].bind(this);
           "main thread";
           return null;
       });
-      loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry) && registerWorkletInternal("main-thread", "da39:80ef4:2", function() {
+      __workletRuntimeLoaded && registerWorkletInternal("main-thread", "da39:80ef4:2", function() {
           const bar = lynxWorkletImpl._workletMap["da39:80ef4:2"].bind(this);
           let { foo } = this["_c"];
           "main thread";
@@ -1821,7 +1825,8 @@ function getCurrentDelta(event) {
     expect(code).toMatchInlineSnapshot(`
       "import { loadWorkletRuntime as __loadWorkletRuntime } from "@lynx-js/react";
       var loadWorkletRuntime = __loadWorkletRuntime;
-      loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry) && registerWorkletInternal("main-thread", "da39:059d0:1", function(event) {
+      const __workletRuntimeLoaded = loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry);
+      __workletRuntimeLoaded && registerWorkletInternal("main-thread", "da39:059d0:1", function(event) {
           lynxWorkletImpl._workletMap["da39:059d0:1"].bind(this);
           let { foo, a, b } = this["_c"];
           "main thread";

packages/react/transform/crates/swc_plugin_worklet/gen_stmt.rs

@@ -11,6 +11,20 @@ use swc_plugins_shared::target::TransformTarget;
 
 pub struct StmtGen {}
 
+struct RegisterWorkletParams<'a> {
+  mode: TransformMode,
+  target: TransformTarget,
+  worklet_type: WorkletType,
+  function_name: Ident,
+  function: Box<Function>,
+  extracted_idents: Vec<Ident>,
+  extracted_js_fns: Vec<(IdentName, Box<Expr>)>,
+  hash: Expr,
+  is_class_member: bool,
+  named_imports: &'a mut HashSet<String>,
+  worklet_runtime_loaded_ident: Ident,
+}
+
 impl StmtGen {
   pub fn transform_worklet(
     mode: TransformMode,
@@ -22,6 +36,7 @@ impl StmtGen {
     ident_collector: &mut ExtractingIdentsCollector,
     is_class_member: bool,
     named_imports: &mut HashSet<String>,
+    worklet_runtime_loaded_ident: Ident,
   ) -> (Box<Expr>, Stmt) {
     let hash = Expr::Lit(hash.into());
     let extracted_value = ident_collector.take_values();
@@ -41,9 +56,9 @@ impl StmtGen {
         hash.clone(),
         named_imports,
       ),
-      StmtGen::gen_register_worklet_stmt(
+      StmtGen::gen_register_worklet_stmt(RegisterWorkletParams {
         mode,
-        if target == TransformTarget::MIXED {
+        target: if target == TransformTarget::MIXED {
           TransformTarget::LEPUS
         } else {
           target
@@ -56,7 +71,8 @@ impl StmtGen {
         hash,
         is_class_member,
         named_imports,
-      ),
+        worklet_runtime_loaded_ident,
+      }),
     )
   }
 
@@ -160,18 +176,21 @@ impl StmtGen {
   /*
    * registerWorklet($type, $hash, $function);
    */
-  fn gen_register_worklet_stmt(
-    mode: TransformMode,
-    target: TransformTarget,
-    worklet_type: WorkletType,
-    function_name: Ident,
-    function: Box<Function>,
-    extracted_idents: Vec<Ident>,
-    extracted_js_fns: Vec<(IdentName, Box<Expr>)>,
-    hash: Expr,
-    is_class_member: bool,
-    named_imports: &mut HashSet<String>,
-  ) -> Stmt {
+  fn gen_register_worklet_stmt(params: RegisterWorkletParams) -> Stmt {
+    let RegisterWorkletParams {
+      mode,
+      target,
+      worklet_type,
+      function_name,
+      function,
+      extracted_idents,
+      extracted_js_fns,
+      hash,
+      is_class_member,
+      named_imports,
+      worklet_runtime_loaded_ident,
+    } = params;
+
     let function_to_register = Box::new(StmtGen::gen_function_to_register(
       function,
       function_name,
@@ -183,7 +202,8 @@ impl StmtGen {
 
     if target == TransformTarget::LEPUS {
       named_imports.insert("loadWorkletRuntime".into());
-      quote!("loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry) && registerWorkletInternal($type_, $hash, $fn_)" as Stmt,
+      quote!("$loaded && registerWorkletInternal($type_, $hash, $fn_)" as Stmt,
+        loaded: Expr = Expr::Ident(worklet_runtime_loaded_ident.clone()),
         type_: Expr = Expr::Lit(worklet_type.type_str().into()),
         hash: Expr = hash,
         fn_: Expr = Expr::Fn(FnExpr {

packages/react/transform/crates/swc_plugin_worklet/lib.rs

@@ -15,9 +15,10 @@ use std::vec;
 use swc_core::common::util::take::Take;
 use swc_core::common::{errors::HANDLER, Span, Spanned, DUMMY_SP};
 use swc_core::ecma::ast::*;
-use swc_core::ecma::utils::prepend_stmts;
+use swc_core::ecma::utils::{prepend_stmts, private_ident};
 use swc_core::ecma::visit::VisitMutWith;
 use swc_core::ecma::visit::{noop_visit_mut_type, VisitMut};
+use swc_core::quote;
 use worklet_type::WorkletType;
 
 use swc_plugins_shared::{target::TransformTarget, transform_mode::TransformMode};
@@ -60,6 +61,8 @@ pub struct WorkletVisitor {
   named_imports: HashSet<String>,
   hasher: WorkletHash,
   shared_identifiers: FxHashSet<Id>,
+  worklet_runtime_loaded: bool,
+  worklet_runtime_loaded_ident: Ident,
 }
 
 impl Default for WorkletVisitor {
@@ -117,6 +120,7 @@ impl VisitMut for WorkletVisitor {
           &mut collector,
           true,
           &mut self.named_imports,
+          self.worklet_runtime_loaded_ident.clone(),
         );
 
         // For JS worklets, the ctx object is later converted to a function by `transformWorklet(..)`
@@ -242,6 +246,7 @@ impl VisitMut for WorkletVisitor {
           &mut collector,
           true,
           &mut self.named_imports,
+          self.worklet_runtime_loaded_ident.clone(),
         );
 
         if should_use_getter {
@@ -322,6 +327,7 @@ impl VisitMut for WorkletVisitor {
       &mut collector,
       false,
       &mut self.named_imports,
+      self.worklet_runtime_loaded_ident.clone(),
     );
 
     *n = VarDecl {
@@ -392,6 +398,7 @@ impl VisitMut for WorkletVisitor {
           &mut collector,
           false,
           &mut self.named_imports,
+          self.worklet_runtime_loaded_ident.clone(),
         );
 
         *n = *worklet_object_expr;
@@ -424,6 +431,7 @@ impl VisitMut for WorkletVisitor {
           &mut collector,
           false,
           &mut self.named_imports,
+          self.worklet_runtime_loaded_ident.clone(),
         );
 
         *n = *worklet_object_expr;
@@ -502,6 +510,7 @@ impl VisitMut for WorkletVisitor {
       &mut collector,
       false,
       &mut self.named_imports,
+      self.worklet_runtime_loaded_ident.clone(),
     );
 
     *n = ModuleDecl::ExportDefaultExpr(ExportDefaultExpr {
@@ -515,13 +524,6 @@ impl VisitMut for WorkletVisitor {
 
   fn visit_mut_module_items(&mut self, n: &mut Vec<ModuleItem>) {
     n.visit_mut_children_with(self);
-    n.extend(
-      self
-        .stmts_to_insert_at_top_level
-        .iter_mut()
-        .filter(|stmt| !stmt.is_empty())
-        .map(|stmt| stmt.take().into()),
-    );
   }
 
   fn visit_mut_module(&mut self, n: &mut Module) {
@@ -548,6 +550,15 @@ impl VisitMut for WorkletVisitor {
 
     n.visit_mut_children_with(self);
 
+    // Add global loadWorkletRuntime call if needed
+    if self.named_imports.contains("loadWorkletRuntime") && !self.worklet_runtime_loaded {
+      self.stmts_to_insert_at_top_level.insert(
+        0,
+        quote!("const $loaded = loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry)" as Stmt, loaded = self.worklet_runtime_loaded_ident.clone()),
+      );
+      self.worklet_runtime_loaded = true;
+    }
+
     let mut specifiers = self.named_imports.iter().collect::<Vec<_>>();
 
     if !specifiers.is_empty() {
@@ -621,6 +632,14 @@ impl VisitMut for WorkletVisitor {
         .into_iter(),
       );
     }
+    // Add statements to insert at top level after processing all items
+    n.body.extend(
+      self
+        .stmts_to_insert_at_top_level
+        .iter_mut()
+        .filter(|stmt| !stmt.is_empty())
+        .map(|stmt| stmt.take().into()),
+    );
   }
 }
 
@@ -686,6 +705,8 @@ impl WorkletVisitor {
       hasher: WorkletHash::new(),
       named_imports: HashSet::default(),
       shared_identifiers: FxHashSet::default(),
+      worklet_runtime_loaded: false,
+      worklet_runtime_loaded_ident: private_ident!("__workletRuntimeLoaded"),
     }
   }
 
@@ -2689,4 +2710,63 @@ class App extends Component {
       }
     "#
   );
+
+  test!(
+    module,
+    Syntax::Typescript(TsSyntax {
+      ..Default::default()
+    }),
+    |_| (
+      resolver(Mark::new(), Mark::new(), true),
+      visit_mut_pass(WorkletVisitor::new(
+        TransformMode::Test,
+        WorkletVisitorConfig {
+          filename: "index.ts".into(),
+          target: TransformTarget::LEPUS,
+          custom_global_ident_names: None,
+          runtime_pkg: "@lynx-js/react".into(),
+        }
+      )),
+      hygiene()
+    ),
+    should_add_worklet_runtime_ident_with_outer_ident,
+    r#"
+      const __workletRuntimeLoaded = false;
+      console.log(__workletRuntimeLoaded);
+
+      function foo() {
+        "main thread";
+        return 1;
+      }
+    "#
+  );
+
+  test!(
+    module,
+    Syntax::Typescript(TsSyntax {
+      ..Default::default()
+    }),
+    |_| (
+      resolver(Mark::new(), Mark::new(), true),
+      visit_mut_pass(WorkletVisitor::new(
+        TransformMode::Test,
+        WorkletVisitorConfig {
+          filename: "index.ts".into(),
+          target: TransformTarget::LEPUS,
+          custom_global_ident_names: None,
+          runtime_pkg: "@lynx-js/react".into(),
+        }
+      )),
+      hygiene()
+    ),
+    should_add_worklet_runtime_ident_with_inner_ident,
+    r#"
+      function foo() {
+        "main thread";
+        const __workletRuntimeLoaded = false;
+        console.log(__workletRuntimeLoaded);
+        return 1;
+      }
+    "#
+  );
 }

packages/react/transform/crates/swc_plugin_worklet/tests/__swc_snapshots__/lib.rs/should_add_worklet_runtime_ident_with_inner_ident.js

@@ -0,0 +1,13 @@
+import { loadWorkletRuntime as __loadWorkletRuntime } from "@lynx-js/react";
+var loadWorkletRuntime = __loadWorkletRuntime;
+let foo = {
+    _wkltId: "a123:test:1"
+};
+const __workletRuntimeLoaded = loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry);
+__workletRuntimeLoaded && registerWorkletInternal("main-thread", "a123:test:1", function() {
+    const foo = lynxWorkletImpl._workletMap["a123:test:1"].bind(this);
+    "main thread";
+    const __workletRuntimeLoaded = false;
+    console.log(__workletRuntimeLoaded);
+    return 1;
+});

packages/react/transform/crates/swc_plugin_worklet/tests/__swc_snapshots__/lib.rs/should_add_worklet_runtime_ident_with_outer_ident.js

@@ -0,0 +1,13 @@
+import { loadWorkletRuntime as __loadWorkletRuntime } from "@lynx-js/react";
+var loadWorkletRuntime = __loadWorkletRuntime;
+const __workletRuntimeLoaded = false;
+console.log(__workletRuntimeLoaded);
+let foo = {
+    _wkltId: "a123:test:1"
+};
+const __workletRuntimeLoaded1 = loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry);
+__workletRuntimeLoaded1 && registerWorkletInternal("main-thread", "a123:test:1", function() {
+    const foo = lynxWorkletImpl._workletMap["a123:test:1"].bind(this);
+    "main thread";
+    return 1;
+});

packages/react/transform/crates/swc_plugin_worklet/tests/__swc_snapshots__/lib.rs/should_capture_env_lepus.js

@@ -12,7 +12,8 @@ let Y = {
     },
     _wkltId: "a77b:test:1"
 };
-loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry) && registerWorkletInternal("main-thread", "a77b:test:1", function() {
+const __workletRuntimeLoaded = loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry);
+__workletRuntimeLoaded && registerWorkletInternal("main-thread", "a77b:test:1", function() {
     const Y = lynxWorkletImpl._workletMap["a77b:test:1"].bind(this);
     let { y1, y2, y3, y4, y5 } = this["_c"];
     "main thread";

packages/react/transform/crates/swc_plugin_worklet/tests/__swc_snapshots__/lib.rs/should_extract_ident_from_this_lepus.js

@@ -6,7 +6,8 @@ let onTapLepus = {
         a: this.a
     }
 };
-loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry) && registerWorkletInternal("main-thread", "a123:test:1", function(event: ReactLynx.Worklet.ITouchEvent) {
+const __workletRuntimeLoaded = loadWorkletRuntime(typeof globDynamicComponentEntry === 'undefined' ? undefined : globDynamicComponentEntry);
+__workletRuntimeLoaded && registerWorkletInternal("main-thread", "a123:test:1", function(event: ReactLynx.Worklet.ITouchEvent) {
     const onTapLepus = lynxWorkletImpl._workletMap["a123:test:1"].bind(this);
     "main thread";
     let a = 1;