Bind handlers' this in Rust

Author: mrbbot

README.md

@@ -18,13 +18,9 @@ import { HTMLRewriter } from "html-rewriter-wasm";
 const encoder = new TextEncoder();
 const decoder = new TextDecoder();
 
+let output = "";
 const rewriter = new HTMLRewriter((outputChunk) => {
-  if (outputChunk.length === 0) {
-    // Remember to free memory on last chunk
-    queueMicrotask(() => rewriter.free());
-  } else {
-    console.log(decoder.decode(outputChunk)); // <p>new</p>
-  }
+  output += decoder.decode(outputChunk);
 });
 
 rewriter.on("p", {
@@ -33,8 +29,13 @@ rewriter.on("p", {
   },
 });
 
-await rewriter.write(encoder.encode("<p>old</p>"));
-await rewriter.end();
+try {
+  await rewriter.write(encoder.encode("<p>old</p>"));
+  await rewriter.end();
+  console.log(output); // <p>new</p>
+} finally {
+  rewriter.free(); // Remember to free memory
+}
 ```
 
 See [test/index.ts](./test/index.ts) for a more traditional `HTMLRewriter`
@@ -85,30 +86,6 @@ and output to strings.
   await rewriter.write(encoder.encode("<p>2</p>"));
   ```
 
-- If using handler classes, you must bind their methods to the class first:
-
-  ```js
-  class Handler {
-    constructor(value) {
-      this.value = value;
-    }
-
-    element(element) {
-      element.setInnerContent(this.value);
-    }
-  }
-  const rewriter = new HTMLRewriter(...);
-  const handler = new Handler("new");
-
-  // ❌
-  rewriter.on("p", handler);
-
-  // ✅
-  rewriter.on("p", {
-    element: handler.element.bind(handler)
-  })
-  ```
-
 ## Building
 
 You can build the package by running `npm run build`. You must do this prior to

src/handlers.rs

@@ -10,6 +10,7 @@ use lol_html::{
     ElementContentHandlers as NativeElementContentHandlers,
 };
 use std::mem;
+use std::rc::Rc;
 use thiserror::Error;
 use wasm_bindgen::JsCast;
 
@@ -29,13 +30,12 @@ extern "C" {
 }
 
 macro_rules! make_handler {
-    ($handler:ident, $JsArgType:ident, $stack_ptr:ident) => {
+    ($handler:ident, $JsArgType:ident, $this:ident, $stack_ptr:ident) => {
         move |arg: &mut _| {
             let (js_arg, anchor) = $JsArgType::from_native(arg);
-            let this = JsValue::NULL;
             let js_arg = JsValue::from(js_arg);
 
-            let res = match $handler.call1(&this, &js_arg) {
+            let res = match $handler.call1(&$this, &js_arg) {
                 Ok(res) => {
                     if let Some(promise) = res.dyn_ref::<JsPromise>() {
                         await_promise($stack_ptr, promise);
@@ -72,18 +72,22 @@ extern "C" {
 
 impl IntoNativeHandlers<NativeElementContentHandlers<'static>> for ElementContentHandlers {
     fn into_native(self, stack_ptr: *mut u8) -> NativeElementContentHandlers<'static> {
+        let handlers: Rc<JsValue> = Rc::new((&self).into());
         let mut native = NativeElementContentHandlers::default();
 
         if let Some(handler) = self.element() {
-            native = native.element(make_handler!(handler, Element, stack_ptr));
+            let this = handlers.clone();
+            native = native.element(make_handler!(handler, Element, this, stack_ptr));
         }
 
         if let Some(handler) = self.comments() {
-            native = native.comments(make_handler!(handler, Comment, stack_ptr));
+            let this = handlers.clone();
+            native = native.comments(make_handler!(handler, Comment, this, stack_ptr));
         }
 
         if let Some(handler) = self.text() {
-            native = native.text(make_handler!(handler, TextChunk, stack_ptr));
+            let this = handlers.clone();
+            native = native.text(make_handler!(handler, TextChunk, this, stack_ptr));
         }
 
         native
@@ -109,22 +113,27 @@ extern "C" {
 
 impl IntoNativeHandlers<NativeDocumentContentHandlers<'static>> for DocumentContentHandlers {
     fn into_native(self, stack_ptr: *mut u8) -> NativeDocumentContentHandlers<'static> {
+        let handlers: Rc<JsValue> = Rc::new((&self).into());
         let mut native = NativeDocumentContentHandlers::default();
 
         if let Some(handler) = self.doctype() {
-            native = native.doctype(make_handler!(handler, Doctype, stack_ptr));
+            let this = handlers.clone();
+            native = native.doctype(make_handler!(handler, Doctype, this, stack_ptr));
         }
 
         if let Some(handler) = self.comments() {
-            native = native.comments(make_handler!(handler, Comment, stack_ptr));
+            let this = handlers.clone();
+            native = native.comments(make_handler!(handler, Comment, this, stack_ptr));
         }
 
         if let Some(handler) = self.text() {
-            native = native.text(make_handler!(handler, TextChunk, stack_ptr));
+            let this = handlers.clone();
+            native = native.text(make_handler!(handler, TextChunk, this, stack_ptr));
         }
 
         if let Some(handler) = self.end() {
-            native = native.end(make_handler!(handler, DocumentEnd, stack_ptr));
+            let this = handlers.clone();
+            native = native.end(make_handler!(handler, DocumentEnd, this, stack_ptr));
         }
 
         native

test/index.ts

@@ -17,42 +17,33 @@ export class HTMLRewriter {
   private documentHandlers: DocumentHandlers[] = [];
 
   on(selector: string, handlers: ElementHandlers): this {
-    this.elementHandlers.push([
-      selector,
-      {
-        element: handlers.element?.bind(handlers),
-        comments: handlers.comments?.bind(handlers),
-        text: handlers.text?.bind(handlers),
-      },
-    ]);
+    this.elementHandlers.push([selector, handlers]);
     return this;
   }
 
   onDocument(handlers: DocumentHandlers): this {
-    this.documentHandlers.push({
-      doctype: handlers.doctype?.bind(handlers),
-      comments: handlers.comments?.bind(handlers),
-      text: handlers.text?.bind(handlers),
-      end: handlers.end?.bind(handlers),
-    });
+    this.documentHandlers.push(handlers);
     return this;
   }
 
   async transform(input: string): Promise<string> {
     let output = "";
     const rewriter = new RawHTMLRewriter((chunk) => {
       output += decoder.decode(chunk);
-      if (chunk.length === 0) queueMicrotask(() => rewriter.free());
     });
     for (const [selector, handlers] of this.elementHandlers) {
       rewriter.on(selector, handlers);
     }
     for (const handlers of this.documentHandlers) {
       rewriter.onDocument(handlers);
     }
-    await rewriter.write(encoder.encode(input));
-    await rewriter.end();
-    return output;
+    try {
+      await rewriter.write(encoder.encode(input));
+      await rewriter.end();
+      return output;
+    } finally {
+      rewriter.free();
+    }
   }
 }