Add support for Element#onEndTag

mrbbot · mrbbot · commit 1bd5f694d121 · 2022-01-18T17:43:03.000Z
diff --git a/html_rewriter.js.patch b/html_rewriter.js.patch
@@ -1,5 +1,5 @@
---- pkg/html_rewriter.js	2021-07-24 10:12:13.000000000 +0100
-+++ pkg2/html_rewriter.js	2021-07-24 10:11:57.000000000 +0100
+--- pkg/html_rewriter.js	2022-01-18 17:37:39.000000000 +0000
++++ pkg2/html_rewriter.js	2022-01-18 17:37:19.000000000 +0000
 @@ -1,7 +1,7 @@
  let imports = {};
  imports['__wbindgen_placeholder__'] = module.exports;
@@ -118,7 +118,7 @@
      }
      /**
      * @param {string} content
-@@ -530,11 +543,13 @@
+@@ -530,17 +543,19 @@
          var ptr0 = passStringToWasm0(content, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
          var len0 = WASM_VECTOR_LEN;
          wasm.element_setInnerContent(this.ptr, ptr0, len0, isLikeNone(content_type) ? 0 : addHeapObject(content_type));
@@ -130,9 +130,38 @@
          wasm.element_removeAndKeepContent(this.ptr);
 +        return this;
      }
+     /**
+     * @param {any} handler
+     */
+     onEndTag(handler) {
+-        wasm.element_onEndTag(this.ptr, addHeapObject(handler));
++        wasm.element_onEndTag(this.ptr, addHeapObject(handler.bind(this)));
+     }
  }
  module.exports.Element = Element;
-@@ -579,25 +594,27 @@
+@@ -597,6 +612,7 @@
+         var ptr0 = passStringToWasm0(content, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
+         var len0 = WASM_VECTOR_LEN;
+         wasm.endtag_before(this.ptr, ptr0, len0, isLikeNone(content_type) ? 0 : addHeapObject(content_type));
++        return this;
+     }
+     /**
+     * @param {string} content
+@@ -606,11 +622,13 @@
+         var ptr0 = passStringToWasm0(content, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
+         var len0 = WASM_VECTOR_LEN;
+         wasm.endtag_after(this.ptr, ptr0, len0, isLikeNone(content_type) ? 0 : addHeapObject(content_type));
++        return this;
+     }
+     /**
+     */
+     remove() {
+         wasm.endtag_remove(this.ptr);
++        return this;
+     }
+ }
+ module.exports.EndTag = EndTag;
+@@ -656,25 +674,27 @@
          var ptr0 = passStringToWasm0(selector, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
          var len0 = WASM_VECTOR_LEN;
          wasm.htmlrewriter_on(this.ptr, ptr0, len0, addHeapObject(handlers));
@@ -164,23 +193,23 @@
      }
      /**
      * @returns {number}
-@@ -638,6 +655,7 @@
+@@ -715,6 +735,7 @@
          var ptr0 = passStringToWasm0(content, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
          var len0 = WASM_VECTOR_LEN;
          wasm.textchunk_before(this.ptr, ptr0, len0, isLikeNone(content_type) ? 0 : addHeapObject(content_type));
 +        return this;
      }
      /**
      * @param {string} content
-@@ -647,6 +665,7 @@
+@@ -724,6 +745,7 @@
          var ptr0 = passStringToWasm0(content, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
          var len0 = WASM_VECTOR_LEN;
          wasm.textchunk_after(this.ptr, ptr0, len0, isLikeNone(content_type) ? 0 : addHeapObject(content_type));
 +        return this;
      }
      /**
      * @param {string} content
-@@ -656,11 +675,13 @@
+@@ -733,11 +755,13 @@
          var ptr0 = passStringToWasm0(content, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
          var len0 = WASM_VECTOR_LEN;
          wasm.textchunk_replace(this.ptr, ptr0, len0, isLikeNone(content_type) ? 0 : addHeapObject(content_type));
@@ -194,8 +223,8 @@
      }
      /**
      * @returns {boolean}
-@@ -801,7 +822,8 @@
- }, arguments) };
+@@ -893,7 +917,8 @@
+ };
  
  module.exports.__wbg_instanceof_Promise_c6535fc791fcc4d2 = function(arg0) {
 -    var ret = getObject(arg0) instanceof Promise;
@@ -204,7 +233,7 @@
      return ret;
  };
  
-@@ -847,5 +869,6 @@
+@@ -939,5 +964,6 @@
  const wasmModule = new WebAssembly.Module(bytes);
  const wasmInstance = new WebAssembly.Instance(wasmModule, imports);
  wasm = wasmInstance.exports;
diff --git a/src/element.rs b/src/element.rs
@@ -1,6 +1,10 @@
+use super::end_tag::EndTag;
+use super::handlers::{await_promise, make_handler, HandlerJsErrorWrap};
 use super::*;
+use js_sys::{Function as JsFunction, Promise as JsPromise};
 use lol_html::html_content::Element as NativeElement;
 use serde_wasm_bindgen::to_value as to_js_value;
+use wasm_bindgen::JsCast;
 
 #[wasm_bindgen]
 pub struct Element(NativeRefWrap<NativeElement<'static, 'static>>);
@@ -100,4 +104,14 @@ impl Element {
     pub fn remove_and_keep_content(&mut self) -> Result<(), JsValue> {
         self.0.get_mut().map(|e| e.remove_and_keep_content())
     }
+
+    #[wasm_bindgen(method, js_name=onEndTag)]
+    pub fn on_end_tag(&mut self, handler: JsFunction) -> Result<(), JsValue> {
+        let this = JsValue::NULL;
+        let stack_ptr = self.0.stack_ptr;
+        self.0
+            .get_mut()?
+            .on_end_tag(make_handler!(handler, EndTag, this, stack_ptr))
+            .into_js_result()
+    }
 }
diff --git a/src/end_tag.rs b/src/end_tag.rs
@@ -0,0 +1,44 @@
+use super::*;
+use lol_html::html_content::EndTag as NativeEndTag;
+
+#[wasm_bindgen]
+pub struct EndTag(NativeRefWrap<NativeEndTag<'static>>);
+
+impl_from_native!(NativeEndTag --> EndTag);
+
+#[wasm_bindgen]
+impl EndTag {
+    #[wasm_bindgen(method, getter=name)]
+    pub fn name(&self) -> JsResult<String> {
+        self.0.get().map(|e| e.name())
+    }
+
+    #[wasm_bindgen(method, setter=name)]
+    pub fn set_name(&mut self, name: &str) -> JsResult<()> {
+        self.0.get_mut().map(|e| e.set_name_str(String::from(name)))
+    }
+
+    pub fn before(
+        &mut self,
+        content: &str,
+        content_type: Option<ContentTypeOptions>,
+    ) -> JsResult<()> {
+        self.0
+            .get_mut()
+            .map(|e| e.before(content, content_type.into_native()))
+    }
+
+    pub fn after(
+        &mut self,
+        content: &str,
+        content_type: Option<ContentTypeOptions>,
+    ) -> JsResult<()> {
+        self.0
+            .get_mut()
+            .map(|e| e.after(content, content_type.into_native()))
+    }
+
+    pub fn remove(&mut self) -> JsResult<()> {
+        self.0.get_mut().map(|e| e.remove())
+    }
+}
diff --git a/src/handlers.rs b/src/handlers.rs
@@ -26,13 +26,13 @@ unsafe impl Sync for HandlerJsErrorWrap {}
 #[wasm_bindgen(raw_module = "./asyncify.js")]
 extern "C" {
     #[wasm_bindgen(js_name = awaitPromise)]
-    fn await_promise(stack_ptr: *mut u8, promise: &JsPromise);
+    pub(crate) fn await_promise(stack_ptr: *mut u8, promise: &JsPromise);
 }
 
 macro_rules! make_handler {
     ($handler:ident, $JsArgType:ident, $this:ident, $stack_ptr:ident) => {
         move |arg: &mut _| {
-            let (js_arg, anchor) = $JsArgType::from_native(arg);
+            let (js_arg, anchor) = $JsArgType::from_native(arg, $stack_ptr);
             let js_arg = JsValue::from(js_arg);
 
             let res = match $handler.call1(&$this, &js_arg) {
@@ -51,6 +51,7 @@ macro_rules! make_handler {
         }
     };
 }
+pub(crate) use make_handler;
 
 pub trait IntoNativeHandlers<T> {
     fn into_native(self, stack_ptr: *mut u8) -> T;
diff --git a/src/html_rewriter.d.ts b/src/html_rewriter.d.ts
@@ -19,6 +19,14 @@ export class Element {
   readonly namespaceURI: string;
   readonly removed: boolean;
   tagName: string;
+  onEndTag(handler: (this: this, endTag: EndTag) => void | Promise<void>): void;
+}
+
+export class EndTag {
+  before(content: string, options?: ContentTypeOptions): this;
+  after(content: string, options?: ContentTypeOptions): this;
+  remove(): this;
+  name: string;
 }
 
 export class Comment {
@@ -68,7 +76,10 @@ export interface HTMLRewriterOptions {
 }
 
 export class HTMLRewriter {
-  constructor(outputSink: (chunk: Uint8Array) => void, options?: HTMLRewriterOptions);
+  constructor(
+    outputSink: (chunk: Uint8Array) => void,
+    options?: HTMLRewriterOptions
+  );
   on(selector: string, handlers: ElementHandlers): this;
   onDocument(handlers: DocumentHandlers): this;
   write(chunk: Uint8Array): Promise<void>;
diff --git a/src/lib.rs b/src/lib.rs
@@ -35,16 +35,19 @@ impl Drop for Anchor<'_> {
 //
 // When anchor goes out of scope, wrapper becomes poisoned and any attempt to get inner
 // object results in exception.
+#[derive(Clone)]
 struct NativeRefWrap<R> {
     inner_ptr: *mut R,
     poisoned: Rc<Cell<bool>>,
+    stack_ptr: *mut u8,
 }
 
 impl<R> NativeRefWrap<R> {
-    pub fn wrap<I>(inner: &mut I) -> (Self, Anchor) {
+    pub fn wrap<I>(inner: &mut I, stack_ptr: *mut u8) -> (Self, Anchor) {
         let wrap = NativeRefWrap {
             inner_ptr: unsafe { mem::transmute(inner) },
             poisoned: Rc::new(Cell::new(false)),
+            stack_ptr,
         };
 
         let anchor = Anchor::new(Rc::clone(&wrap.poisoned));
@@ -158,8 +161,11 @@ macro_rules! impl_mutations {
 macro_rules! impl_from_native {
     ($Ty:ident --> $JsTy:ident) => {
         impl $JsTy {
-            pub(crate) fn from_native<'r>(inner: &'r mut $Ty) -> (Self, Anchor<'r>) {
-                let (ref_wrap, anchor) = NativeRefWrap::wrap(inner);
+            pub(crate) fn from_native<'r>(
+                inner: &'r mut $Ty,
+                stack_ptr: *mut u8,
+            ) -> (Self, Anchor<'r>) {
+                let (ref_wrap, anchor) = NativeRefWrap::wrap(inner, stack_ptr);
 
                 ($JsTy(ref_wrap), anchor)
             }
@@ -171,6 +177,7 @@ mod comment;
 mod doctype;
 mod document_end;
 mod element;
+mod end_tag;
 mod handlers;
 mod html_rewriter;
 mod text_chunk;
diff --git a/test/element.spec.ts b/test/element.spec.ts
@@ -134,7 +134,7 @@ test("element allows chaining", async (t) => {
     })
     .transform("<p>test</p>");
 });
-test("handles element async handler", async (t) => {
+test.serial("handles element async handler", async (t) => {
   const res = await new HTMLRewriter()
     .on("p", {
       async element(element) {
@@ -158,3 +158,102 @@ test("handles element class handler", async (t) => {
     .transform("<p>test</p>");
   t.is(res, "<p>new</p>");
 });
+
+test("handles end tag properties", async (t) => {
+  const res = await new HTMLRewriter()
+    .on("p", {
+      element(element) {
+        element.onEndTag(function (end) {
+          t.is(this, element);
+          t.is(end.name, "p");
+          end.name = "h1";
+        });
+      },
+    })
+    .transform("<p>test</p>");
+  t.is(res, "<p>test</h1>");
+});
+test("handles end tag mutations", async (t) => {
+  const input = "<p>test</p>";
+  const beforeAfterExpected = [
+    "<p>",
+    "test",
+    "&lt;span&gt;before&lt;/span&gt;",
+    "<span>before html</span>",
+    "</p>",
+    "<span>after html</span>",
+    "&lt;span&gt;after&lt;/span&gt;",
+  ].join("");
+  const removeExpected = "<p>test";
+
+  // before/after
+  let res = await new HTMLRewriter()
+    .on("p", {
+      element(element) {
+        const that = this;
+        element.onEndTag((end) => {
+          t.is(this, that);
+          end.before("<span>before</span>");
+          end.before("<span>before html</span>", { html: true });
+          end.after("<span>after</span>");
+          end.after("<span>after html</span>", { html: true });
+        });
+      },
+    })
+    .transform(input);
+  t.is(res, beforeAfterExpected);
+
+  // remove
+  res = await new HTMLRewriter()
+    .on("p", {
+      element(element) {
+        element.onEndTag((end) => {
+          end.remove();
+        });
+      },
+    })
+    .transform(input);
+  t.is(res, removeExpected);
+});
+test("end tag allows chaining", async (t) => {
+  t.plan(3);
+  await new HTMLRewriter()
+    .on("p", {
+      element(element) {
+        element.onEndTag((end) => {
+          t.is(end.before(""), end);
+          t.is(end.after(""), end);
+          t.is(end.remove(), end);
+        });
+      },
+    })
+    .transform("<p>test</p>");
+});
+test.serial("handles end tag async handler", async (t) => {
+  const res = await new HTMLRewriter()
+    .on("p", {
+      element(element) {
+        element.onEndTag(async (end) => {
+          await wait(50);
+          end.before("!");
+        });
+      },
+    })
+    .transform("<p>test</p>");
+  t.is(res, "<p>test!</p>");
+});
+test("uses last end tag handler", async (t) => {
+  const res = await new HTMLRewriter()
+    .on("p", {
+      element(element) {
+        element.onEndTag((end) => {
+          end.before("1");
+        });
+        element.onEndTag((end) => {
+          end.before("2");
+        });
+      },
+    })
+    .transform("<p>test</p>");
+  t.is(res, "<p>test2</p>");
+});
