feat(plugin): Perform actual transforms in plugins (#3220)

Cargo.lock

@@ -2158,9 +2158,9 @@ checksum = "11000e6ba5020e53e7cc26f73b91ae7d5496b4977851479edb66b694c0675c21"
 
 [[package]]
 name = "rkyv"
-version = "0.7.28"
+version = "0.7.29"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "631f7d2a2854abb66724f492ce5256e79685a673dc210ac022194cedd5c914d3"
+checksum = "49a37de5dfc60bae2d94961dacd03c7b80e426b66a99fa1b17799570dbdd8f96"
 dependencies = [
  "bytecheck",
  "hashbrown",
@@ -2172,9 +2172,9 @@ dependencies = [
 
 [[package]]
 name = "rkyv_derive"
-version = "0.7.28"
+version = "0.7.29"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c067e650861a749720952aed722fb344449bc95de33e6456d426f5c7d44f71c0"
+checksum = "719d447dd0e84b23cee6cb5b32d97e21efb112a3e3c636c8da36647b938475a1"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -3410,6 +3410,7 @@ dependencies = [
  "anyhow",
  "serde",
  "serde_json",
+ "swc_atoms",
  "swc_common",
  "swc_ecma_ast",
  "swc_ecma_visit",
@@ -3423,12 +3424,14 @@ dependencies = [
  "libloading",
  "once_cell",
  "parking_lot",
+ "rkyv",
  "serde",
  "serde_json",
  "swc_atoms",
  "swc_common",
  "swc_ecma_ast",
  "swc_ecma_parser",
+ "swc_ecma_visit",
  "testing",
  "wasmer",
  "wasmer-cache",

crates/swc_plugin/Cargo.toml

@@ -17,3 +17,4 @@ serde_json = "1.0.64"
 swc_common = {version = "0.16.0", path = "../swc_common", features = ["plugin-mode"]}
 swc_ecma_ast = {version = "0.63.0", path = "../swc_ecma_ast", features = ["rkyv-impl"]}
 swc_ecma_visit = {version = "0.49.0", path = "../swc_ecma_visit"}
+swc_atoms = {version = "0.2.0", path = "../swc_atoms"}

crates/swc_plugin/src/lib.rs

@@ -1,7 +1,10 @@
 // Reexports
-pub use swc_common::chain;
-pub use swc_ecma_ast::*;
-pub use swc_ecma_visit::*;
+pub use swc_common::{chain, DUMMY_SP};
+pub mod ast {
+    pub use swc_atoms::*;
+    pub use swc_ecma_ast::*;
+    pub use swc_ecma_visit::*;
+}
 
 #[doc(hidden)]
 #[cfg(target_arch = "wasm32")]

crates/swc_plugin_runner/Cargo.toml

@@ -14,6 +14,7 @@ anyhow = "1.0.42"
 libloading = "0.7.0"
 once_cell = "1.8.0"
 parking_lot = "0.11"
+rkyv = "0.7.29"
 serde = {version = "1.0.126", features = ["derive"]}
 serde_json = "1.0.64"
 swc_atoms = {version = "0.2.7", path = '../swc_atoms'}
@@ -25,5 +26,6 @@ wasmer-cache = "2.1.1"
 
 [dev-dependencies]
 testing = {version = "0.17.0", path = "../testing"}
+swc_ecma_visit = { version = "0.49.0", path = "../swc_ecma_visit" }
 
 [features]

crates/swc_plugin_runner/src/lib.rs

@@ -7,9 +7,10 @@ use anyhow::{Context, Error};
 use once_cell::sync::Lazy;
 use parking_lot::Mutex;
 use resolve::PluginCache;
+use rkyv::{AlignedVec, Deserialize};
 use swc_common::{collections::AHashMap, plugin::serialize_for_plugin};
 use swc_ecma_ast::Program;
-use wasmer::{imports, Instance, Memory, MemoryType, Module, Store};
+use wasmer::{imports, Array, Instance, Memory, MemoryType, Module, Store, WasmPtr};
 use wasmer_cache::{Cache, Hash};
 
 pub mod resolve;
@@ -133,6 +134,10 @@ fn copy_memory_to_instance(instance: &Instance, program: &Program) -> Result<(i3
     Ok((alloc_ptr, serialized_len.try_into().unwrap()))
 }
 
+// TODO
+// - alloc / free memories
+// - resolve trait bounds compilation errors for swc_plugin::{se,dese}rialize fn
+// - ergonomic wrappers for the transform logics with error handlings
 pub fn apply_js_plugin(
     _plugin_name: &str,
     path: &Path,
@@ -143,14 +148,39 @@ pub fn apply_js_plugin(
     (|| -> Result<_, Error> {
         let instance = load_plugin(path, cache)?;
 
-        let plugin_process = instance
+        let plugin_process_wasm_exported_fn = instance
             .exports
             .get_native_function::<(i32, u32), (i32, i32)>("process")?;
 
         let (alloc_ptr, len) = copy_memory_to_instance(&instance, &program)?;
-        let (_returned_ptr, _returned_len) = plugin_process.call(alloc_ptr, len)?;
+        let (returned_ptr, returned_len) = plugin_process_wasm_exported_fn.call(alloc_ptr, len)?;
+        let returned_len = returned_len.try_into().unwrap();
 
-        Ok(program)
+        // Reconstruct AlignedVec from ptr returned by plugin
+        let memory = instance.exports.get_memory("memory")?;
+        let ptr: WasmPtr<u8, Array> = WasmPtr::new(returned_ptr as _);
+
+        let mut transformed_serialized = AlignedVec::with_capacity(returned_len);
+
+        // Deref & read through plguin's wasm memory space via returned ptr
+        let derefed_ptr = ptr
+            .deref(memory, 0, returned_len.try_into().unwrap())
+            .unwrap();
+
+        let transformed_raw_bytes = derefed_ptr
+            .iter()
+            .enumerate()
+            .take(returned_len)
+            .map(|(_size, cell)| cell.get())
+            .collect::<Vec<u8>>();
+
+        transformed_serialized.extend_from_slice(&transformed_raw_bytes);
+
+        // Deserialize into Program from reconstructed raw bytes
+        let archived = unsafe { rkyv::archived_root::<Program>(&transformed_serialized[..]) };
+        let transformed_program: Program = archived.deserialize(&mut rkyv::Infallible).unwrap();
+
+        Ok(transformed_program)
     })()
     .with_context(|| {
         format!(

crates/swc_plugin_runner/tests/integration.rs

@@ -5,8 +5,9 @@ use std::{
     process::{Command, Stdio},
 };
 use swc_common::FileName;
-use swc_ecma_ast::EsVersion;
+use swc_ecma_ast::{CallExpr, EsVersion, Expr, ExprOrSuper, Lit, MemberExpr, Str};
 use swc_ecma_parser::{lexer::Lexer, EsConfig, Parser, StringInput, Syntax};
+use swc_ecma_visit::{Visit, VisitWith};
 
 /// Returns the path to the built plugin
 fn build_plugin(dir: &Path) -> Result<PathBuf, Error> {
@@ -41,6 +42,29 @@ fn build_plugin(dir: &Path) -> Result<PathBuf, Error> {
     Err(anyhow!("Could not find built plugin"))
 }
 
+struct TestVisitor {
+    pub plugin_transform_found: bool,
+}
+
+impl Visit for TestVisitor {
+    fn visit_call_expr(&mut self, call: &CallExpr) {
+        if let ExprOrSuper::Expr(expr) = &call.callee {
+            if let Expr::Member(MemberExpr { obj, .. }) = &**expr {
+                if let ExprOrSuper::Expr(expr) = obj {
+                    if let Expr::Ident(ident) = &**expr {
+                        if ident.sym == *"console" {
+                            let args = &*(call.args[0].expr);
+                            if let Expr::Lit(Lit::Str(Str { value, .. })) = args {
+                                self.plugin_transform_found = value == "changed_via_plugin";
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
 #[test]
 fn internal() -> Result<(), Error> {
     let path = build_plugin(
@@ -67,11 +91,19 @@ fn internal() -> Result<(), Error> {
 
         let program = parser.parse_program().unwrap();
 
-        let _program =
+        let program =
             swc_plugin_runner::apply_js_plugin("internal-test", &path, &mut None, "{}", program)
-                .unwrap();
+                .expect("Plugin should apply transform");
 
-        Ok(())
+        let mut visitor = TestVisitor {
+            plugin_transform_found: false,
+        };
+        program.visit_with(&mut visitor);
+
+        visitor
+            .plugin_transform_found
+            .then(|| visitor.plugin_transform_found)
+            .ok_or(())
     })
     .unwrap();
 

tests/rust-plugins/swc_internal_plugin/src/lib.rs

@@ -1,11 +1,28 @@
 use rkyv::{AlignedVec, Deserialize};
-use swc_plugin::{ModuleItem, Program, VisitMut};
+use swc_plugin::{ast::*, DUMMY_SP};
 
-struct Dummy;
+struct ConsoleOutputReplacer;
 
-impl VisitMut for Dummy {
-    fn visit_mut_module_item(&mut self, _module: &mut ModuleItem) {
-        // noop
+/// An example plugin replaces any `console.log(${text})` into
+/// `console.log('changed_via_plugin')`.
+impl VisitMut for ConsoleOutputReplacer {
+    fn visit_mut_call_expr(&mut self, call: &mut CallExpr) {
+        if let ExprOrSuper::Expr(expr) = &call.callee {
+            if let Expr::Member(MemberExpr { obj, .. }) = &**expr {
+                if let ExprOrSuper::Expr(expr) = obj {
+                    if let Expr::Ident(ident) = &**expr {
+                        if ident.sym == *"console" {
+                            call.args[0].expr = Box::new(Expr::Lit(Lit::Str(Str {
+                                span: DUMMY_SP,
+                                has_escape: false,
+                                kind: StrKind::default(),
+                                value: JsWord::from("changed_via_plugin"),
+                            })));
+                        }
+                    }
+                }
+            }
+        }
     }
 }
 
@@ -16,13 +33,22 @@ impl VisitMut for Dummy {
 // - better developer experience: println / dbg!() doesn't emit anything
 // - typed json config instead of str, which requires additional deserialization
 pub fn process(ast_ptr: *mut u8, len: u32) -> (i32, i32) {
-    let mut vec = AlignedVec::with_capacity(len.try_into().unwrap());
-    let v = unsafe { std::slice::from_raw_parts(ast_ptr, len.try_into().unwrap()) };
-    vec.extend_from_slice(v);
+    // Read raw serialized bytes from wasm's memory space. Host (SWC) should allocated memory, copy bytes and pass ptr to it.
+    let raw_serialized_bytes = unsafe { std::slice::from_raw_parts(ast_ptr, len.try_into().unwrap()) };
+
+    // Reconstruct AlignedVec from raw bytes
+    let mut serialized_program = AlignedVec::with_capacity(len.try_into().unwrap());
+    serialized_program.extend_from_slice(raw_serialized_bytes);
 
     // TODO: trait bound complaining in deserialize_for_plugin<T>
-    let archived = unsafe { rkyv::archived_root::<Program>(&vec[..]) };
-    let _v: Program = archived.deserialize(&mut rkyv::Infallible).unwrap();
+    // Actual deserialization
+    let archived = unsafe { rkyv::archived_root::<Program>(&serialized_program[..]) };
+    let program: Program = archived.deserialize(&mut rkyv::Infallible).unwrap();
 
-    (0, 0)
+    // Actual plugin transformation
+    let transformed_program = program.fold_with(&mut as_folder(ConsoleOutputReplacer));
+
+    // Serialize transformed result, return back to the host.
+    let serialized_result = rkyv::to_bytes::<_, 512>(&transformed_program).expect("Should serializable");
+    (serialized_result.as_ptr() as _, serialized_result.len().try_into().unwrap())
 }