Search code examples
javascriptrustbabeljsswc

Injecting Comments with a SWC Rust plugin


I am trying to inject a comment to existing javascript code.

In Babel there is the addComment helper which can be called on a AST node.

Here is a simple (but rather silly) transformation written in Babel:

console.log`Hello, ${name}!`; -> /*#__PURE__*/console.log("Hello, ", name, "!");

const babel = require("@babel/core");

function run(code) {
  const result = babel.transform(code, {
    plugins: [
      function transform({ types: t }) {
        return {
          visitor: {
            TaggedTemplateExpression(path) {
              const { tag, quasi } = path.node;

              if (
                t.isMemberExpression(tag) &&
                t.isIdentifier(tag.object, { name: "console" }) &&
                t.isIdentifier(tag.property, { name: "log" })
              ) {
                let args = [];

                quasi.quasis.forEach((element, index) => {
                  args.push(t.stringLiteral(element.value.raw));
                  if (index < quasi.expressions.length) {
                    args.push(quasi.expressions[index]);
                  }
                });

                path.replaceWith(
                  t.callExpression(
                    t.addComment(
                      t.memberExpression(
                        t.identifier("console"),
                        t.identifier("log")
                      ),
                      "leading",
                      "#__PURE__"
                    ),
                    args
                  )
                );
              }
            },
          },
        };
      },
    ],
  });
  return result.code;
}

const code = "console.log`Hello, ${name}!`;";
console.log(run(code));
// -> /*#__PURE__*/console.log("Hello, ", name, "!");

However in Rust things are a little bit more complicated as only one variable can own a data, and there's at most one mutable reference to it on top of that SWC implemented some performance tricks.

So in SWC you have to use the PluginCommentsProxy which is explained in the current SWC version 0.279.0 with this flow chart:

 Below diagram shows one reference example how guest does trampoline between
 host's memory space.
┌───────────────────────────────────────┐    ┌─────────────────────────────────────────────┐
│Host (SWC/core)                        │    │Plugin (wasm)                                │
│  ┌────────────────────────────────┐   │    │                                             │
│  │COMMENTS.with()                 │   │    │  ┌──────────────────────────────────────┐   │
│  │                                │   │    │  │PluginCommentsProxy                   │   │
│  │                                │   │    │  │                                      │   │
│  │ ┌────────────────────────────┐ │   │    │  │ ┌────────────────────────────────┐   │   │
│  │ │get_leading_comments_proxy()│◀┼───┼────┼──┼─┤get_leading()                   │   │   │
│  │ │                            │ │   │    │  │ │                                │   │   │
│  │ │                            │ │   │    │  │ │ ┌──────────────────────────┐   │   │   │
│  │ │                            │─┼───┼──┬─┼──┼─┼─▶AllocatedBytesPtr(p,len)  │   │   │   │
│  │ └────────────────────────────┘ │   │  │ │  │ │ │                          │   │   │   │
│  │                                │   │  │ │  │ │ └─────────────┬────────────┘   │   │   │
│  │                                │   │  │ │  │ │               │                │   │   │
│  │                                │   │  │ │  │ │ ┌─────────────▼────────────┐   │   │   │
│  └────────────────────────────────┘   │  │ │  │ │ │Vec<Comments>             │   │   │   │
│                                       │  └─┼──┼─┼─▶                          │   │   │   │
│                                       │    │  │ │ └──────────────────────────┘   │   │   │
│                                       │    │  │ └────────────────────────────────┘   │   │
│                                       │    │  └──────────────────────────────────────┘   │
└───────────────────────────────────────┘    └─────────────────────────────────────────────┘

 1. Plugin calls `PluginCommentsProxy::get_leading()`. PluginCommentsProxy is
 a struct constructed in plugin's memory space.
 2. `get_leading()` internally calls `__get_leading_comments_proxy`, which is
 imported fn `get_leading_comments_proxy` exists in the host.
 3. Host access necessary values in its memory space (COMMENTS)
 4. Host copies value to be returned into plugin's memory space. Memory
 allocation for the value should be manually performed.
 5. Host completes imported fn, `PluginCommentsProxy::get_leading()` now can
 read, deserialize memory host wrote.
 - In case of `get_leading`, returned value is non-deterministic vec
 (`Vec<Comments>`) guest cannot preallocate with specific length. Instead,
 guest passes a fixed size struct (AllocatedBytesPtr), once host allocates
 actual vec into guest it'll write pointer to the vec into the struct.
comments.add_leading(
  node.span.lo,
  Comment {
    kind: swc_core::common::comments::CommentKind::Block,
    span: DUMMY_SP,
    text: "#__PURE__".to_string(),
  },
)

Unfortunately I wasn't able to test it properly with swc_core::ecma::transforms::testing.

#[cfg(test)]

mod tests {
  use super::*;
  use std::path::PathBuf;
  use swc_core::ecma::transforms::testing::{test_fixture};
  use swc_ecma_transforms_testing::{FixtureTestConfig};

  #[testing::fixture("tests/fixture/**/input.tsx")]
  fn fixture(input: PathBuf) {
    test_fixture(
      Default::default(),
      &|tester| as_folder(TransformVisitor::new(&tester.comments)),
      &input,
      &input.with_file_name("output.tsx"),
      FixtureTestConfig::default(),
    );
  }
}

Unfortunately that does not work because tester.comments is of type Rc<SingleThreadedComments>.

I saw examples using <C> for example the MillionJs transformer:

fn transform_block<C>(context: &ProgramStateContext, node: &mut CallExpr, comments: C)
where
    C: Comments,
{

Ideally, tests should reflect how the code will be used in production. Adding a generic type parameter just for testing makes the code harder to read and feels wrong to me.

Is there a better way?


Solution

  • SWC author here.

    You can make your transformer generic over C: Comments, like the official pure_annotations pass. Then store C just like other generics.

    New way

    You can use PluginCommentProxy from Wasm plugin even while teting if you run tests via swc_ecma_transforms_testing or swc_core::ecma::transforms::testing, using a method like test_fixture.

    
    PluginCommentProxy.add_leading(n.span.lo, Comment {
        // ...fields
    });
    
    

    will just work while testing.

    Old way

    This was necessary before https://github.com/swc-project/swc/pull/9150, but this way still works.

    struct PureAnnotations<C>
    where
        C: Comments,
    {
        imports: AHashMap<Id, (JsWord, JsWord)>,
        comments: Option<C>,
    }
    
    

    after then, you should make the impl section generic over C.

    impl<C> VisitMut for PureAnnotations<C>
    where
        C: Comments,
    {
        noop_visit_mut_type!();
    
    }
    

    You can add proper visitor methods to achieve your goal.

    Alternatively, you can accept &dyn Comments or Option<&dyn Comments>. The official fixer pass uses this pattern, to reduce binary size. In this case, you should add '_ between impl and Fold in your constructor return type.

    pub fn fixer(comments: Option<&dyn Comments>) -> impl '_ + Fold + VisitMut {
        as_folder(Fixer {
            comments,
            ctx: Default::default(),
            span_map: Default::default(),
            in_for_stmt_head: Default::default(),
            in_opt_chain: Default::default(),
            remove_only: false,
        })
    }