feat(plugin): Implement plugin api for comments (#4229)

crates/swc_common/src/plugin.rs

@@ -57,6 +57,15 @@ impl Serialized {
         Serialized { field: vec }
     }
 
+    /// Not an actual trait Into impl: simple wrapper to deserialize<T>:expect()
+    pub fn into<T>(self) -> T
+    where
+        T: rkyv::Archive,
+        T::Archived: rkyv::Deserialize<T, rkyv::Infallible>,
+    {
+        Serialized::deserialize(&self).expect("Should able to deserialize")
+    }
+
     #[allow(clippy::should_implement_trait)]
     pub fn as_ref(&self) -> &rkyv::AlignedVec {
         &self.field

crates/swc_plugin/src/lib.rs

@@ -7,7 +7,7 @@ pub use swc_common::{
 };
 
 pub mod comments {
-    pub use swc_common::comments::Comments;
+    pub use swc_common::comments::{Comment, CommentKind, Comments};
     pub use swc_plugin_comments::PluginCommentsProxy;
 }
 

crates/swc_plugin_comments/src/plugin_comments_proxy.rs

@@ -6,9 +6,22 @@ use swc_common::{
     BytePos,
 };
 
-#[cfg(target = "wasm32")]
+#[cfg(target_arch = "wasm32")]
 extern "C" {
+    fn __copy_comment_to_host_env(bytes_ptr: i32, bytes_ptr_len: i32);
+    fn __add_leading_comment_proxy(byte_pos: u32);
+    fn __add_leading_comments_proxy(byte_pos: u32);
+    fn __has_leading_comments_proxy(byte_pos: u32) -> i32;
+    fn __move_leading_comments_proxy(from_byte_pos: u32, to_byte_pos: u32);
+    fn __take_leading_comments_proxy(byte_pos: u32, allocated_ret_ptr: i32) -> i32;
     fn __get_leading_comments_proxy(byte_pos: u32, allocated_ret_ptr: i32) -> i32;
+    fn __add_trailing_comment_proxy(byte_pos: u32);
+    fn __add_trailing_comments_proxy(byte_pos: u32);
+    fn __has_trailing_comments_proxy(byte_pos: u32) -> i32;
+    fn __move_trailing_comments_proxy(from_byte_pos: u32, to_byte_pos: u32);
+    fn __take_trailing_comments_proxy(byte_pos: u32, allocated_ret_ptr: i32) -> i32;
+    fn __get_trailing_comments_proxy(byte_pos: u32, allocated_ret_ptr: i32) -> i32;
+    fn __add_pure_comment_proxy(byte_pos: u32);
 }
 
 /// A struct to exchance allocated Vec<Comment> between memory spaces.
@@ -28,28 +41,43 @@ pub struct CommentsVecPtr(pub i32, pub i32);
 pub struct PluginCommentsProxy;
 
 #[cfg(feature = "plugin-mode")]
-impl Comments for PluginCommentsProxy {
-    fn add_leading(&self, pos: BytePos, cmt: Comment) {
-        unimplemented!("not implemented yet");
+impl PluginCommentsProxy {
+    /// Copy guest memory's struct into host via CommentHostEnvironment's
+    /// comment_buffer as serialized to pass param from guest to the host for
+    /// the fn like add_leading*.
+    #[cfg_attr(not(target_arch = "wasm32"), allow(unused))]
+    fn allocate_comments_buffer_to_host<T>(&self, value: &T)
+    where
+        T: rkyv::Serialize<rkyv::ser::serializers::AllocSerializer<512>>,
+    {
+        #[cfg(target_arch = "wasm32")]
+        {
+            let serialized = Serialized::serialize(value).expect("Should able to serialize value");
+            let serialized_comment_ptr_ref = serialized.as_ref();
+            unsafe {
+                // We need to copy PluginCommentProxy's param for add_leading (Comment, or
+                // Vec<Comment>) to the host, before calling proxy to the host. This'll fill in
+                // CommentHostEnvironment's buffer, subsequent proxy call will read &
+                // deserialize it.
+                __copy_comment_to_host_env(
+                    serialized_comment_ptr_ref.as_ptr() as _,
+                    serialized_comment_ptr_ref
+                        .len()
+                        .try_into()
+                        .expect("Should able to convert ptr length"),
+                );
+            }
+        }
     }
 
-    fn add_leading_comments(&self, pos: BytePos, comments: Vec<Comment>) {
-        unimplemented!("not implemented yet");
-    }
-
-    fn has_leading(&self, pos: BytePos) -> bool {
-        unimplemented!("not implemented yet");
-    }
-
-    fn move_leading(&self, from: BytePos, to: BytePos) {
-        unimplemented!("not implemented yet");
-    }
-
-    fn take_leading(&self, pos: BytePos) -> Option<Vec<Comment>> {
-        unimplemented!("not implemented yet");
-    }
-
-    fn get_leading(&self, pos: BytePos) -> Option<Vec<Comment>> {
+    /// Utiilty wrapper to call host fn which returns a Comment or Vec<Comment>.
+    #[cfg_attr(not(target_arch = "wasm32"), allow(unused))]
+    fn read_returned_comments_from_host<F, R>(&self, f: F) -> Option<R>
+    where
+        F: FnOnce(i32) -> i32,
+        R: rkyv::Archive,
+        R::Archived: rkyv::Deserialize<R, rkyv::Infallible>,
+    {
         // Allocate CommentsVecPtr to get return value from the host
         let comments_vec_ptr = CommentsVecPtr(0, 0);
         let serialized_comments_vec_ptr = Serialized::serialize(&comments_vec_ptr)
@@ -58,18 +86,20 @@ impl Comments for PluginCommentsProxy {
         let serialized_comments_vec_ptr_raw_ptr = serialized_comments_vec_ptr_ref.as_ptr();
         let serialized_comments_vec_ptr_raw_len = serialized_comments_vec_ptr_ref.len();
 
-        #[cfg(target = "wasm32")]
+        #[cfg(target_arch = "wasm32")]
         {
-            let ret = unsafe {
-                __get_leading_comments_proxy(pos.0, serialized_comments_vec_ptr_raw_ptr as _)
-            };
+            let ret = f(serialized_comments_vec_ptr_raw_ptr as _);
 
+            // Host fn call completes: by contract in commments_proxy, if return value is 0
+            // we know there's no value to read. Otherwise, we know host filled in
+            // CommentsVecPtr to the pointer for the actual value for the
+            // results.
             if ret == 0 {
                 return None;
             }
         }
 
-        // First, reconstruct CommentsVecPtr to reveal ptr to the allocated
+        // Now reconstruct CommentsVecPtr to reveal ptr to the allocated
         // Vec<Comments>
         let comments_vec_ptr: CommentsVecPtr = unsafe {
             Serialized::deserialize_from_ptr(
@@ -81,54 +111,135 @@ impl Comments for PluginCommentsProxy {
             .expect("Should able to deserialize CommentsVecPtr")
         };
 
-        // Using CommentsVecPtr's value, reconstruct actual Vec<Comments>
+        // Using CommentsVecPtr's value, reconstruct actual Comment, or Vec<Comments>
         Some(unsafe {
             Serialized::deserialize_from_ptr(comments_vec_ptr.0 as _, comments_vec_ptr.1)
                 .expect("Returned comments should be serializable")
         })
     }
+}
+
+#[cfg(feature = "plugin-mode")]
+#[cfg_attr(not(target_arch = "wasm32"), allow(unused))]
+impl Comments for PluginCommentsProxy {
+    fn add_leading(&self, pos: BytePos, cmt: Comment) {
+        self.allocate_comments_buffer_to_host(&cmt);
+        #[cfg(target_arch = "wasm32")]
+        unsafe {
+            __add_leading_comment_proxy(pos.0);
+        }
+    }
+
+    fn add_leading_comments(&self, pos: BytePos, comments: Vec<Comment>) {
+        self.allocate_comments_buffer_to_host(&comments);
+        #[cfg(target_arch = "wasm32")]
+        unsafe {
+            __add_leading_comments_proxy(pos.0);
+        }
+    }
+
+    fn has_leading(&self, pos: BytePos) -> bool {
+        #[cfg(target_arch = "wasm32")]
+        {
+            if unsafe { __has_leading_comments_proxy(pos.0) } == 0 {
+                false
+            } else {
+                true
+            }
+        }
+        #[cfg(not(target_arch = "wasm32"))]
+        false
+    }
+
+    fn move_leading(&self, from: BytePos, to: BytePos) {
+        #[cfg(target_arch = "wasm32")]
+        unsafe {
+            __move_leading_comments_proxy(from.0, to.0)
+        }
+    }
+
+    fn take_leading(&self, pos: BytePos) -> Option<Vec<Comment>> {
+        #[cfg(target_arch = "wasm32")]
+        return self.read_returned_comments_from_host(|serialized_ptr| unsafe {
+            __take_leading_comments_proxy(pos.0, serialized_ptr)
+        });
+
+        #[cfg(not(target_arch = "wasm32"))]
+        None
+    }
+
+    fn get_leading(&self, pos: BytePos) -> Option<Vec<Comment>> {
+        #[cfg(target_arch = "wasm32")]
+        return self.read_returned_comments_from_host(|serialized_ptr| unsafe {
+            __get_leading_comments_proxy(pos.0, serialized_ptr)
+        });
+
+        #[cfg(not(target_arch = "wasm32"))]
+        None
+    }
 
     fn add_trailing(&self, pos: BytePos, cmt: Comment) {
-        unimplemented!("not implemented yet");
+        self.allocate_comments_buffer_to_host(&cmt);
+        #[cfg(target_arch = "wasm32")]
+        unsafe {
+            __add_trailing_comment_proxy(pos.0);
+        }
     }
 
     fn add_trailing_comments(&self, pos: BytePos, comments: Vec<Comment>) {
-        unimplemented!("not implemented yet");
+        self.allocate_comments_buffer_to_host(&comments);
+        #[cfg(target_arch = "wasm32")]
+        unsafe {
+            __add_trailing_comments_proxy(pos.0);
+        }
     }
 
     fn has_trailing(&self, pos: BytePos) -> bool {
-        unimplemented!("not implemented yet");
+        #[cfg(target_arch = "wasm32")]
+        {
+            if unsafe { __has_trailing_comments_proxy(pos.0) } == 0 {
+                false
+            } else {
+                true
+            }
+        }
+        #[cfg(not(target_arch = "wasm32"))]
+        false
     }
 
     fn move_trailing(&self, from: BytePos, to: BytePos) {
-        unimplemented!("not implemented yet");
+        #[cfg(target_arch = "wasm32")]
+        unsafe {
+            __move_trailing_comments_proxy(from.0, to.0)
+        }
     }
 
     fn take_trailing(&self, pos: BytePos) -> Option<Vec<Comment>> {
-        unimplemented!("not implemented yet");
+        #[cfg(target_arch = "wasm32")]
+        return self.read_returned_comments_from_host(|serialized_ptr| unsafe {
+            __take_trailing_comments_proxy(pos.0, serialized_ptr)
+        });
+
+        #[cfg(not(target_arch = "wasm32"))]
+        None
     }
 
     fn get_trailing(&self, pos: BytePos) -> Option<Vec<Comment>> {
-        unimplemented!("not implemented yet");
+        #[cfg(target_arch = "wasm32")]
+        return self.read_returned_comments_from_host(|serialized_ptr| unsafe {
+            __get_trailing_comments_proxy(pos.0, serialized_ptr)
+        });
+
+        #[cfg(not(target_arch = "wasm32"))]
+        None
     }
 
     fn add_pure_comment(&self, pos: BytePos) {
-        unimplemented!("not implemented yet");
-    }
-
-    fn with_leading<F, Ret>(&self, pos: BytePos, f: F) -> Ret
-    where
-        Self: Sized,
-        F: FnOnce(&[Comment]) -> Ret,
-    {
-        unimplemented!("not implemented yet");
-    }
-
-    fn with_trailing<F, Ret>(&self, pos: BytePos, f: F) -> Ret
-    where
-        Self: Sized,
-        F: FnOnce(&[Comment]) -> Ret,
-    {
-        unimplemented!("not implemented yet");
+        #[cfg(target_arch = "wasm32")]
+        {
+            unsafe {
+                __add_pure_comment_proxy(pos.0);
+            }
+        }
     }
 }

crates/swc_plugin_runner/src/host_environment.rs

@@ -1,29 +1,25 @@
-use std::sync::Arc;
+use wasmer::{LazyInit, Memory};
 
-use parking_lot::Mutex;
-use wasmer::{LazyInit, Memory, NativeFunc};
-
-#[derive(wasmer::WasmerEnv, Clone)]
 /// An external enviornment state imported (declared in host, injected into
-/// guest) fn can access. This'll allow host access updated state via plugin's
-/// transform.
+/// guest) fn can access. This'll allow host to read from updated state from
+/// guest.
+///
+/// This is `base` environment exposes guest's memory space only. For other
+/// calls requires additional data to be set in the host, separate
+/// hostenvironments are decalred. Refer `CommentsHostEnvironment` for an
+/// example.
+///
 /// ref: https://docs.wasmer.io/integrations/examples/host-functions#declaring-the-data
-pub struct HostEnvironment {
+#[derive(wasmer::WasmerEnv, Clone)]
+pub struct BaseHostEnvironment {
     #[wasmer(export)]
     pub memory: wasmer::LazyInit<Memory>,
-    pub transform_result: Arc<Mutex<Vec<u8>>>,
-    /// Attached imported fn `__alloc` to the hostenvironment to allow any other
-    /// imported fn can allocate guest's memory space from host runtime.
-    #[wasmer(export(name = "__alloc"))]
-    pub alloc_guest_memory: LazyInit<NativeFunc<u32, i32>>,
 }
 
-impl HostEnvironment {
-    pub fn new(transform_result: &Arc<Mutex<Vec<u8>>>) -> HostEnvironment {
-        HostEnvironment {
-            memory: LazyInit::default(),
-            transform_result: transform_result.clone(),
-            alloc_guest_memory: LazyInit::default(),
+impl BaseHostEnvironment {
+    pub fn new() -> BaseHostEnvironment {
+        BaseHostEnvironment {
+            memory: LazyInit::new(),
         }
     }
 }

crates/swc_plugin_runner/src/imported_fn/comments.rs

@@ -1,7 +1,216 @@
-use swc_common::{comments::Comments, plugin::Serialized, BytePos};
-use swc_plugin_comments::{CommentsVecPtr, HostCommentsStorage, COMMENTS};
+use std::sync::Arc;
 
-use crate::{host_environment::HostEnvironment, memory_interop::write_into_memory_view};
+use parking_lot::Mutex;
+use swc_common::{
+    comments::{Comments, SingleThreadedComments},
+    plugin::Serialized,
+    BytePos,
+};
+use swc_plugin_comments::{CommentsVecPtr, COMMENTS};
+use wasmer::{LazyInit, Memory, NativeFunc};
+
+use crate::memory_interop::{copy_bytes_into_host, write_into_memory_view};
+
+/// External environment state for imported (declared in host, injected into
+/// guest) fn for comments proxy.
+#[derive(wasmer::WasmerEnv, Clone)]
+pub struct CommentHostEnvironment {
+    #[wasmer(export)]
+    pub memory: wasmer::LazyInit<Memory>,
+    /// Attached imported fn `__alloc` to the hostenvironment to allow any other
+    /// imported fn can allocate guest's memory space from host runtime.
+    #[wasmer(export(name = "__alloc"))]
+    pub alloc_guest_memory: LazyInit<NativeFunc<u32, i32>>,
+    /// A buffer to `Comment`, or `Vec<Comment>` plugin need to pass to the host
+    /// to perform mutable comment operations like `add_leading, or
+    /// add_leading_comments`. This is vec to serialized bytes, doesn't
+    /// distinguish if it's Vec<Comment> or not. Host should perform
+    /// typed deserialization accordingly.
+    pub mutable_comment_buffer: Arc<Mutex<Vec<u8>>>,
+}
+
+impl CommentHostEnvironment {
+    pub fn new(mutable_comment_buffer: &Arc<Mutex<Vec<u8>>>) -> CommentHostEnvironment {
+        CommentHostEnvironment {
+            memory: LazyInit::default(),
+            alloc_guest_memory: LazyInit::default(),
+            mutable_comment_buffer: mutable_comment_buffer.clone(),
+        }
+    }
+}
+
+/// Copy given serialized byte into host's comment buffer, subsequent proxy call
+/// in the host can read it.
+pub fn copy_comment_to_host_env(env: &CommentHostEnvironment, bytes_ptr: i32, bytes_ptr_len: i32) {
+    if let Some(memory) = env.memory_ref() {
+        (*env.mutable_comment_buffer.lock()) =
+            copy_bytes_into_host(memory, bytes_ptr, bytes_ptr_len);
+    }
+}
+
+/// Utility fn to unwrap necessary values for the comments fn operation when fn
+/// needs to return values.
+fn unwrap_comments_storage_or_default<F, R>(f: F, default: R) -> R
+where
+    F: FnOnce(&SingleThreadedComments) -> R,
+{
+    if !COMMENTS.is_set() {
+        return default;
+    }
+
+    COMMENTS.with(|storage| {
+        if let Some(comments) = &storage.inner {
+            f(comments)
+        } else {
+            default
+        }
+    })
+}
+
+/// Utility fn to unwrap necessary values for the comments fn operation when fn
+/// does not need to return values.
+fn unwrap_comments_storage<F>(f: F)
+where
+    F: FnOnce(&SingleThreadedComments),
+{
+    unwrap_comments_storage_or_default(
+        |comments| {
+            f(comments);
+            None
+        },
+        None as Option<i32>,
+    );
+}
+
+/// Utility fn to unwrap necessary values for the commments, as well as host
+/// environment's state.
+fn unwrap_comments_storage_with_env<F, R>(env: &CommentHostEnvironment, f: F, default: R) -> R
+where
+    F: FnOnce(&SingleThreadedComments, &Memory, &NativeFunc<u32, i32>) -> R,
+{
+    if let Some(memory) = env.memory_ref() {
+        if let Some(alloc_guest_memory) = env.alloc_guest_memory_ref() {
+            return unwrap_comments_storage_or_default(
+                |comments| f(comments, memory, alloc_guest_memory),
+                default,
+            );
+        }
+    }
+    default
+}
+
+/// Set `return` value to pass into guest from functions returning values with
+/// non-deterministic size like `Vec<Comment>`. Guest pre-allocates a struct to
+/// contain ptr to the value, host in here allocates guest memory for the actual
+/// value then returns its ptr with length to the preallocated struct.
+fn allocate_return_values_into_guest(
+    memory: &Memory,
+    alloc_guest_memory: &NativeFunc<u32, i32>,
+    allocated_ret_ptr: i32,
+    serialized_bytes: &Serialized,
+) {
+    let serialized_bytes_len = serialized_bytes.as_ref().len();
+
+    let (allocated_ptr, allocated_ptr_len) =
+        write_into_memory_view(memory, serialized_bytes, |_| {
+            // In most cases our host-plugin tranmpoline works in a way that
+            // plugin pre-allocates
+            // memory before calling host imported fn. But in case of
+            // comments return value is Vec<Comments> which
+            // guest cannot predetermine size to allocate, intead
+            // let host allocate by calling guest's alloc via attached
+            // hostenvironment.
+            alloc_guest_memory
+                .call(
+                    serialized_bytes_len
+                        .try_into()
+                        .expect("Should be able to convert size"),
+                )
+                .expect("Should able to allocate memory in the plugin")
+        });
+
+    // Retuning (allocated_ptr, len) into caller (plugin)
+    let comment_ptr_serialized =
+        Serialized::serialize(&CommentsVecPtr(allocated_ptr, allocated_ptr_len))
+            .expect("Should be serializable");
+
+    write_into_memory_view(memory, &comment_ptr_serialized, |_| allocated_ret_ptr);
+}
+
+/// Common logics for add_*_comment/comments.
+fn add_comments_inner<F>(env: &CommentHostEnvironment, byte_pos: u32, f: F)
+where
+    F: FnOnce(&SingleThreadedComments, BytePos, Serialized),
+{
+    unwrap_comments_storage(|comments| {
+        let byte_pos = BytePos(byte_pos);
+        // PluginCommentProxy in the guest should've copied buffer already
+        let comment_byte = &mut (*env.mutable_comment_buffer.lock());
+        let serialized = Serialized::new_for_plugin(
+            comment_byte,
+            comment_byte
+                .len()
+                .try_into()
+                .expect("Should able to convert ptr length"),
+        );
+
+        f(comments, byte_pos, serialized);
+
+        // This is not strictly required, but will help to ensure to access previous
+        // call's buffer when we make some mistakes for the execution order.
+        comment_byte.clear();
+    });
+}
+
+pub fn add_leading_comment_proxy(env: &CommentHostEnvironment, byte_pos: u32) {
+    add_comments_inner(env, byte_pos, |comments, byte_pos, serialized| {
+        comments.add_leading(byte_pos, serialized.into());
+    });
+}
+
+pub fn add_leading_comments_proxy(env: &CommentHostEnvironment, byte_pos: u32) {
+    add_comments_inner(env, byte_pos, |comments, byte_pos, serialized| {
+        comments.add_leading_comments(byte_pos, serialized.into());
+    });
+}
+
+pub fn has_leading_comments_proxy(byte_pos: u32) -> i32 {
+    unwrap_comments_storage_or_default(|comments| comments.has_leading(BytePos(byte_pos)) as i32, 0)
+}
+
+pub fn move_leading_comments_proxy(from_byte_pos: u32, to_byte_pos: u32) {
+    unwrap_comments_storage(|comments| {
+        comments.move_leading(BytePos(from_byte_pos), BytePos(to_byte_pos))
+    });
+}
+
+pub fn take_leading_comments_proxy(
+    env: &CommentHostEnvironment,
+    byte_pos: u32,
+    allocated_ret_ptr: i32,
+) -> i32 {
+    unwrap_comments_storage_with_env(
+        env,
+        |comments, memory, alloc_guest_memory| {
+            let leading_comments = comments.take_leading(BytePos(byte_pos));
+            if let Some(leading_comments) = leading_comments {
+                let serialized_leading_comments_vec_bytes =
+                    Serialized::serialize(&leading_comments).expect("Should be serializable");
+
+                allocate_return_values_into_guest(
+                    memory,
+                    alloc_guest_memory,
+                    allocated_ret_ptr,
+                    &serialized_leading_comments_vec_bytes,
+                );
+                1
+            } else {
+                0
+            }
+        },
+        0,
+    )
+}
 
 /// Ask to get leading_comments from currently scoped comments held by
 /// HostCommentsStorage.
@@ -9,64 +218,114 @@ use crate::{host_environment::HostEnvironment, memory_interop::write_into_memory
 /// Returns 1 if operation success with Some(Vec<Comments>), 0 otherwise.
 /// Allocated results should be read through CommentsPtr.
 pub fn get_leading_comments_proxy(
-    env: &HostEnvironment,
+    env: &CommentHostEnvironment,
     byte_pos: u32,
     allocated_ret_ptr: i32,
 ) -> i32 {
-    if let Some(memory) = env.memory_ref() {
-        if let Some(alloc_guest_memory) = env.alloc_guest_memory_ref() {
-            if !COMMENTS.is_set() {
-                return 0;
-            }
+    unwrap_comments_storage_with_env(
+        env,
+        |comments, memory, alloc_guest_memory| {
+            let leading_comments = comments.get_leading(BytePos(byte_pos));
+            if let Some(leading_comments) = leading_comments {
+                let serialized_leading_comments_vec_bytes =
+                    Serialized::serialize(&leading_comments).expect("Should be serializable");
 
-            return COMMENTS.with(|storage: &HostCommentsStorage| {
-                if let Some(comments) = &storage.inner {
-                    let leading_comments = comments.get_leading(BytePos(byte_pos));
-                    if let Some(leading_comments) = leading_comments {
-                        let serialized_leading_comments_vec_bytes =
-                            Serialized::serialize(&leading_comments)
-                                .expect("Should be serializable");
-
-                        let serialized_bytes_len =
-                            serialized_leading_comments_vec_bytes.as_ref().len();
-
-                        let (allocated_ptr, allocated_ptr_len) = write_into_memory_view(
-                            memory,
-                            &serialized_leading_comments_vec_bytes,
-                            |_| {
-                                // In most cases our host-plugin tranmpoline works in a way that
-                                // plugin pre-allocates
-                                // memory before calling host imported fn. But in case of
-                                // comments return value is Vec<Comments> which
-                                // guest cannot predetermine size to allocate, intead
-                                // let host allocate by calling guest's alloc via attached
-                                // hostenvironment.
-                                alloc_guest_memory
-                                    .call(
-                                        serialized_bytes_len
-                                            .try_into()
-                                            .expect("Should be able to convert size"),
-                                    )
-                                    .expect("Should able to allocate memory in the plugin")
-                            },
-                        );
-
-                        // Retuning (allocated_ptr, len) into caller (plugin)
-                        let comment_ptr_serialized = Serialized::serialize(&CommentsVecPtr(
-                            allocated_ptr,
-                            allocated_ptr_len,
-                        ))
-                        .expect("Should be serializable");
-
-                        write_into_memory_view(memory, &comment_ptr_serialized, |_| {
-                            allocated_ret_ptr
-                        });
-                        return 1;
-                    }
-                }
+                allocate_return_values_into_guest(
+                    memory,
+                    alloc_guest_memory,
+                    allocated_ret_ptr,
+                    &serialized_leading_comments_vec_bytes,
+                );
+                1
+            } else {
                 0
-            });
-        }
-    }
-    0
+            }
+        },
+        0,
+    )
+}
+
+pub fn add_trailing_comment_proxy(env: &CommentHostEnvironment, byte_pos: u32) {
+    add_comments_inner(env, byte_pos, |comments, byte_pos, serialized| {
+        comments.add_trailing(byte_pos, serialized.into());
+    });
+}
+
+pub fn add_trailing_comments_proxy(env: &CommentHostEnvironment, byte_pos: u32) {
+    add_comments_inner(env, byte_pos, |comments, byte_pos, serialized| {
+        comments.add_trailing_comments(byte_pos, serialized.into());
+    });
+}
+
+pub fn has_trailing_comments_proxy(byte_pos: u32) -> i32 {
+    unwrap_comments_storage_or_default(
+        |comments| comments.has_trailing(BytePos(byte_pos)) as i32,
+        0,
+    )
+}
+
+pub fn move_trailing_comments_proxy(from_byte_pos: u32, to_byte_pos: u32) {
+    unwrap_comments_storage(|comments| {
+        comments.move_trailing(BytePos(from_byte_pos), BytePos(to_byte_pos))
+    });
+}
+
+pub fn take_trailing_comments_proxy(
+    env: &CommentHostEnvironment,
+    byte_pos: u32,
+    allocated_ret_ptr: i32,
+) -> i32 {
+    unwrap_comments_storage_with_env(
+        env,
+        |comments, memory, alloc_guest_memory| {
+            let trailing_comments = comments.take_trailing(BytePos(byte_pos));
+            if let Some(leading_comments) = trailing_comments {
+                let serialized_leading_comments_vec_bytes =
+                    Serialized::serialize(&leading_comments).expect("Should be serializable");
+
+                allocate_return_values_into_guest(
+                    memory,
+                    alloc_guest_memory,
+                    allocated_ret_ptr,
+                    &serialized_leading_comments_vec_bytes,
+                );
+                1
+            } else {
+                0
+            }
+        },
+        0,
+    )
+}
+
+pub fn get_trailing_comments_proxy(
+    env: &CommentHostEnvironment,
+    byte_pos: u32,
+    allocated_ret_ptr: i32,
+) -> i32 {
+    unwrap_comments_storage_with_env(
+        env,
+        |comments, memory, alloc_guest_memory| {
+            let trailing_comments = comments.get_trailing(BytePos(byte_pos));
+            if let Some(leading_comments) = trailing_comments {
+                let serialized_leading_comments_vec_bytes =
+                    Serialized::serialize(&leading_comments).expect("Should be serializable");
+
+                allocate_return_values_into_guest(
+                    memory,
+                    alloc_guest_memory,
+                    allocated_ret_ptr,
+                    &serialized_leading_comments_vec_bytes,
+                );
+                1
+            } else {
+                0
+            }
+        },
+        0,
+    )
+}
+
+pub fn add_pure_comment_proxy(byte_pos: u32) {
+    unwrap_comments_storage(|comments| comments.add_pure_comment(BytePos(byte_pos)));
 }

crates/swc_plugin_runner/src/imported_fn/handler.rs

@@ -3,9 +3,9 @@ use swc_common::{
     plugin::Serialized,
 };
 
-use crate::{host_environment::HostEnvironment, memory_interop::copy_bytes_into_host};
+use crate::{host_environment::BaseHostEnvironment, memory_interop::copy_bytes_into_host};
 
-pub fn emit_diagnostics(env: &HostEnvironment, bytes_ptr: i32, bytes_ptr_len: i32) {
+pub fn emit_diagnostics(env: &BaseHostEnvironment, bytes_ptr: i32, bytes_ptr_len: i32) {
     if let Some(memory) = env.memory_ref() {
         if HANDLER.is_set() {
             HANDLER.with(|handler| {

crates/swc_plugin_runner/src/imported_fn/hygiene.rs

@@ -1,6 +1,6 @@
 use swc_common::{hygiene::MutableMarkContext, plugin::Serialized, Mark, SyntaxContext};
 
-use crate::{host_environment::HostEnvironment, memory_interop::write_into_memory_view};
+use crate::{host_environment::BaseHostEnvironment, memory_interop::write_into_memory_view};
 
 /// A proxy to Mark::fresh() that can be used in plugin.
 /// This it not direcly called by plugin, instead `impl Mark` will selectively
@@ -27,7 +27,7 @@ pub fn mark_set_builtin_proxy(self_mark: u32, is_builtin: u32) {
 /// Inside of guest context, once this host function returns it'll assign params
 /// with return value accordingly.
 pub fn mark_is_descendant_of_proxy(
-    env: &HostEnvironment,
+    env: &BaseHostEnvironment,
     self_mark: u32,
     ancestor: u32,
     allocated_ptr: i32,
@@ -49,7 +49,7 @@ pub fn mark_is_descendant_of_proxy(
     }
 }
 
-pub fn mark_least_ancestor_proxy(env: &HostEnvironment, a: u32, b: u32, allocated_ptr: i32) {
+pub fn mark_least_ancestor_proxy(env: &BaseHostEnvironment, a: u32, b: u32, allocated_ptr: i32) {
     let a = Mark::from_u32(a);
     let b = Mark::from_u32(b);
 
@@ -70,7 +70,11 @@ pub fn syntax_context_apply_mark_proxy(self_syntax_context: u32, mark: u32) -> u
         .as_u32()
 }
 
-pub fn syntax_context_remove_mark_proxy(env: &HostEnvironment, self_mark: u32, allocated_ptr: i32) {
+pub fn syntax_context_remove_mark_proxy(
+    env: &BaseHostEnvironment,
+    self_mark: u32,
+    allocated_ptr: i32,
+) {
     let mut self_mark = SyntaxContext::from_u32(self_mark);
 
     let return_value = self_mark.remove_mark();

crates/swc_plugin_runner/src/imported_fn/mod.rs

@@ -46,7 +46,19 @@ use std::sync::Arc;
 use parking_lot::Mutex;
 use wasmer::{imports, Function, ImportObject, Module};
 
-use crate::{host_environment::HostEnvironment, imported_fn::comments::get_leading_comments_proxy};
+use crate::{
+    host_environment::BaseHostEnvironment,
+    imported_fn::{
+        comments::{
+            add_leading_comment_proxy, add_leading_comments_proxy, add_pure_comment_proxy,
+            add_trailing_comment_proxy, add_trailing_comments_proxy, copy_comment_to_host_env,
+            get_leading_comments_proxy, get_trailing_comments_proxy, has_leading_comments_proxy,
+            has_trailing_comments_proxy, move_leading_comments_proxy, move_trailing_comments_proxy,
+            take_leading_comments_proxy, take_trailing_comments_proxy, CommentHostEnvironment,
+        },
+        set_transform_result::{set_transform_result, TransformResultHostEnvironment},
+    },
+};
 
 mod comments;
 mod handler;
@@ -55,7 +67,6 @@ mod set_transform_result;
 
 use handler::*;
 use hygiene::*;
-use set_transform_result::*;
 
 /// Create an ImportObject includes functions to be imported from host to the
 /// plugins.
@@ -68,16 +79,13 @@ pub(crate) fn build_import_object(
     // transfrom_result
     let set_transform_result_fn_decl = Function::new_native_with_env(
         wasmer_store,
-        HostEnvironment::new(transform_result),
+        TransformResultHostEnvironment::new(transform_result),
         set_transform_result,
     );
 
     // handler
-    let emit_diagnostics_fn_decl = Function::new_native_with_env(
-        wasmer_store,
-        HostEnvironment::new(transform_result),
-        emit_diagnostics,
-    );
+    let emit_diagnostics_fn_decl =
+        Function::new_native_with_env(wasmer_store, BaseHostEnvironment::new(), emit_diagnostics);
 
     // hygiene
     let mark_fresh_fn_decl = Function::new_native(wasmer_store, mark_fresh_proxy);
@@ -86,13 +94,13 @@ pub(crate) fn build_import_object(
     let mark_set_builtin_fn_decl = Function::new_native(wasmer_store, mark_set_builtin_proxy);
     let mark_is_descendant_of_fn_decl = Function::new_native_with_env(
         wasmer_store,
-        HostEnvironment::new(transform_result),
+        BaseHostEnvironment::new(),
         mark_is_descendant_of_proxy,
     );
 
     let mark_least_ancestor_fn_decl = Function::new_native_with_env(
         wasmer_store,
-        HostEnvironment::new(transform_result),
+        BaseHostEnvironment::new(),
         mark_least_ancestor_proxy,
     );
 
@@ -100,19 +108,91 @@ pub(crate) fn build_import_object(
         Function::new_native(wasmer_store, syntax_context_apply_mark_proxy);
     let syntax_context_remove_mark_fn_decl = Function::new_native_with_env(
         wasmer_store,
-        HostEnvironment::new(transform_result),
+        BaseHostEnvironment::new(),
         syntax_context_remove_mark_proxy,
     );
     let syntax_context_outer_fn_decl =
         Function::new_native(wasmer_store, syntax_context_outer_proxy);
 
     // comments
+    let comment_buffer = Arc::new(Mutex::new(vec![]));
+
+    let copy_comment_to_host_env_fn_decl = Function::new_native_with_env(
+        wasmer_store,
+        CommentHostEnvironment::new(&comment_buffer),
+        copy_comment_to_host_env,
+    );
+
+    let add_leading_comment_fn_decl = Function::new_native_with_env(
+        wasmer_store,
+        CommentHostEnvironment::new(&comment_buffer),
+        add_leading_comment_proxy,
+    );
+
+    let add_leading_comments_fn_decl = Function::new_native_with_env(
+        wasmer_store,
+        CommentHostEnvironment::new(&comment_buffer),
+        add_leading_comments_proxy,
+    );
+
+    let has_leading_comments_fn_decl =
+        Function::new_native(wasmer_store, has_leading_comments_proxy);
+
+    let move_leading_comments_fn_decl =
+        Function::new_native(wasmer_store, move_leading_comments_proxy);
+
+    let take_leading_comments_fn_decl = Function::new_native_with_env(
+        wasmer_store,
+        // take_* doesn't need to share buffer to pass values from plugin to the host - do not
+        // clone buffer here.
+        CommentHostEnvironment::new(&Default::default()),
+        take_leading_comments_proxy,
+    );
+
     let get_leading_comments_fn_decl = Function::new_native_with_env(
         wasmer_store,
-        HostEnvironment::new(transform_result),
+        // get_* doesn't need to share buffer to pass values from plugin to the host - do not clone
+        // buffer here.
+        CommentHostEnvironment::new(&Default::default()),
         get_leading_comments_proxy,
     );
 
+    let add_trailing_comment_fn_decl = Function::new_native_with_env(
+        wasmer_store,
+        CommentHostEnvironment::new(&comment_buffer),
+        add_trailing_comment_proxy,
+    );
+
+    let add_trailing_comments_fn_decl = Function::new_native_with_env(
+        wasmer_store,
+        CommentHostEnvironment::new(&comment_buffer),
+        add_trailing_comments_proxy,
+    );
+
+    let has_trailing_comments_fn_decl =
+        Function::new_native(wasmer_store, has_trailing_comments_proxy);
+
+    let move_trailing_comments_fn_decl =
+        Function::new_native(wasmer_store, move_trailing_comments_proxy);
+
+    let take_trailing_comments_fn_decl = Function::new_native_with_env(
+        wasmer_store,
+        // take_* doesn't need to share buffer to pass values from plugin to the host - do not
+        // clone buffer here.
+        CommentHostEnvironment::new(&Default::default()),
+        take_trailing_comments_proxy,
+    );
+
+    let get_trailing_comments_fn_decl = Function::new_native_with_env(
+        wasmer_store,
+        // get_* doesn't need to share buffer to pass values from plugin to the host - do not clone
+        // buffer here.
+        CommentHostEnvironment::new(&Default::default()),
+        get_trailing_comments_proxy,
+    );
+
+    let add_pure_comment_fn_decl = Function::new_native(wasmer_store, add_pure_comment_proxy);
+
     imports! {
         "env" => {
             // transform
@@ -130,7 +210,20 @@ pub(crate) fn build_import_object(
             "__syntax_context_remove_mark_proxy" => syntax_context_remove_mark_fn_decl,
             "__syntax_context_outer_proxy" => syntax_context_outer_fn_decl,
             // comments
+            "__copy_comment_to_host_env" => copy_comment_to_host_env_fn_decl,
+            "__add_leading_comment_proxy" => add_leading_comment_fn_decl,
+            "__add_leading_comments_proxy" => add_leading_comments_fn_decl,
+            "__has_leading_comments_proxy" => has_leading_comments_fn_decl,
+            "__move_leading_comments_proxy" => move_leading_comments_fn_decl,
+            "__take_leading_comments_proxy" => take_leading_comments_fn_decl,
             "__get_leading_comments_proxy" => get_leading_comments_fn_decl,
+            "__add_trailing_comment_proxy" => add_trailing_comment_fn_decl,
+            "__add_trailing_comments_proxy" => add_trailing_comments_fn_decl,
+            "__has_trailing_comments_proxy" => has_trailing_comments_fn_decl,
+            "__move_trailing_comments_proxy" => move_trailing_comments_fn_decl,
+            "__take_trailing_comments_proxy" => take_trailing_comments_fn_decl,
+            "__get_trailing_comments_proxy" => get_trailing_comments_fn_decl,
+            "__add_pure_comment_proxy" => add_pure_comment_fn_decl,
         }
     }
 }

crates/swc_plugin_runner/src/imported_fn/set_transform_result.rs

@@ -1,10 +1,40 @@
-use crate::{host_environment::HostEnvironment, memory_interop::copy_bytes_into_host};
+use std::sync::Arc;
+
+use parking_lot::Mutex;
+use wasmer::{LazyInit, Memory};
+
+use crate::memory_interop::copy_bytes_into_host;
+
+/// Environment states allow to return guest's transform result back to the
+/// host, using a buffer `transform_result` attached to the environment.
+///
+/// When plugin performs its transform it'll fill in `transform_result` with
+/// serialized result. Host will reconstruct AST from those value.
+#[derive(wasmer::WasmerEnv, Clone)]
+pub struct TransformResultHostEnvironment {
+    #[wasmer(export)]
+    pub memory: wasmer::LazyInit<Memory>,
+    pub transform_result: Arc<Mutex<Vec<u8>>>,
+}
+
+impl TransformResultHostEnvironment {
+    pub fn new(transform_result: &Arc<Mutex<Vec<u8>>>) -> TransformResultHostEnvironment {
+        TransformResultHostEnvironment {
+            memory: LazyInit::default(),
+            transform_result: transform_result.clone(),
+        }
+    }
+}
 
 /// Set plugin's transformed result into host's enviroment.
 /// This is an `imported` fn - when we instantiate plugin module, we inject this
 /// fn into pluging's export space. Once transform completes, plugin will call
 /// this to set its result back to host.
-pub fn set_transform_result(env: &HostEnvironment, bytes_ptr: i32, bytes_ptr_len: i32) {
+pub fn set_transform_result(
+    env: &TransformResultHostEnvironment,
+    bytes_ptr: i32,
+    bytes_ptr_len: i32,
+) {
     if let Some(memory) = env.memory_ref() {
         (*env.transform_result.lock()) = copy_bytes_into_host(memory, bytes_ptr, bytes_ptr_len);
     }

tests/rust-plugins/swc_internal_plugin/Cargo.lock

@@ -63,15 +63,6 @@ dependencies = [
  "scoped-tls",
 ]
 
-[[package]]
-name = "better_scoped_tls"
-version = "0.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b73e8ecdec39e98aa3b19e8cd0b8ed8f77ccb86a6b0b2dc7cd86d105438a2123"
-dependencies = [
- "scoped-tls",
-]
-
 [[package]]
 name = "bitflags"
 version = "1.3.2"
@@ -621,12 +612,12 @@ dependencies = [
 
 [[package]]
 name = "swc_common"
-version = "0.17.18"
+version = "0.17.19"
 dependencies = [
  "ahash",
  "anyhow",
  "ast_node",
- "better_scoped_tls 0.1.0",
+ "better_scoped_tls",
  "cfg-if",
  "debug_unreachable",
  "either",
@@ -647,7 +638,7 @@ dependencies = [
 
 [[package]]
 name = "swc_ecma_ast"
-version = "0.73.0"
+version = "0.73.1"
 dependencies = [
  "is-macro",
  "num-bigint",
@@ -701,7 +692,7 @@ dependencies = [
 
 [[package]]
 name = "swc_plugin"
-version = "0.39.0"
+version = "0.41.0"
 dependencies = [
  "swc_atoms",
  "swc_common",
@@ -715,14 +706,14 @@ dependencies = [
 name = "swc_plugin_comments"
 version = "0.1.0"
 dependencies = [
- "better_scoped_tls 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "better_scoped_tls",
  "rkyv",
  "swc_common",
 ]
 
 [[package]]
 name = "swc_plugin_macro"
-version = "0.3.1"
+version = "0.4.0"
 dependencies = [
  "proc-macro2",
  "quote",