squeeze the last ms of performance out of the rewriter

rewriter/src/lib.rs

@@ -1,7 +1,8 @@
 pub mod rewrite;
 
-use std::str::FromStr;
+use std::str::{from_utf8, FromStr};
 
+use js_sys::Uint8Array;
 use rewrite::rewrite;
 use url::Url;
 use wasm_bindgen::prelude::*;
@@ -13,12 +14,12 @@ extern "C" {
 }
 
 #[wasm_bindgen]
-pub fn rewrite_js(js: &str, url: &str) -> String {
+pub fn rewrite_js(js: &str, url: &str) -> Vec<u8> {
     rewrite(js, Url::from_str(url).unwrap())
 }
 
 #[wasm_bindgen]
-pub fn rewrite_js_from_arraybuffer(js: &[u8], url: &str) -> String {
+pub fn rewrite_js_from_arraybuffer(js: &[u8], url: &str) -> Vec<u8> {
     // technically slower than the c++ string conversion but it will create *less copies*
 
     let js = unsafe { std::str::from_utf8_unchecked(js) };

rewriter/src/main.rs

@@ -1,5 +1,9 @@
 #![allow(clippy::print_stdout)]
-use std::{env, path::Path, str::FromStr};
+use std::{
+    env,
+    path::Path,
+    str::{from_utf8, FromStr},
+};
 
 use oxc_allocator::Allocator;
 use oxc_ast::{
@@ -28,10 +32,15 @@ fn main() -> std::io::Result<()> {
 
     println!(
         "{}",
+        from_utf8(
             rewrite(
                 &source_text,
                 Url::from_str("https://google.com/glorngle/si.js").unwrap()
             )
+            .as_slice()
+        )
+        .unwrap()
+        .to_string()
     );
 
     Ok(())

rewriter/src/rewrite.rs

@@ -184,7 +184,7 @@ impl Rewriter {
     }
 }
 
-pub fn rewrite(js: &str, url: Url) -> String {
+pub fn rewrite(js: &str, url: Url) -> Vec<u8> {
     let allocator = Allocator::default();
     let source_type = SourceType::default();
     let ret = Parser::new(&allocator, &js, source_type).parse();
@@ -241,10 +241,8 @@ pub fn rewrite(js: &str, url: Url) -> String {
         }
     }
 
-    let mut buffer = String::new();
-    // pre-allocate the space we need. should make copies faster
     let size_estimate = (original_len as i32 + difference) as usize;
-    buffer.reserve(size_estimate);
+    let mut buffer: Vec<u8> = Vec::with_capacity(size_estimate);
 
     let mut offset = 0;
     for change in ast_pass.jschanges {
@@ -254,8 +252,9 @@ pub fn rewrite(js: &str, url: Url) -> String {
                 let start = span.start as usize;
                 let end = span.end as usize;
 
-                buffer.push_str(&js[offset..start]);
+                buffer.extend_from_slice(unsafe { js.slice_unchecked(offset, start) }.as_bytes());
-                buffer.push_str(&text);
+
+                buffer.extend_from_slice(text.as_bytes());
                 offset = end;
 
                 // offset = (offset as i64 + (text.len() as i64 - len as i64)) as usize;
@@ -281,7 +280,7 @@ pub fn rewrite(js: &str, url: Url) -> String {
             _ => {}
         }
     }
-    buffer.push_str(&js[offset..]);
+    buffer.extend_from_slice(js[offset..].as_bytes());
 
     return buffer;
 }

src/shared/rewriters/js.ts

@@ -1,22 +1,7 @@
-// import { parseModule } from "meriyah";
+import { decodeUrl } from "./url";
-// import { generate } from "astring";
-// import { makeTraveler } from "astravel";
-import { decodeUrl, encodeUrl } from "./url";
-// import * as ESTree from "estree";
 
 // i am a cat. i like to be petted. i like to be fed. i like to be
-
+import { initSync, rewrite_js, rewrite_js_from_arraybuffer } from "../../../rewriter/out/rewriter.js";
-// js rewiter is NOT finished
-
-// location
-// window
-// self
-// globalThis
-// this
-// top
-// parent
-
-import { initSync, rewrite_js } from "../../../rewriter/out/rewriter.js";
 import "../../../static/wasm.js";
 
 initSync(
@@ -30,130 +15,15 @@ export function rewriteJs(js: string | ArrayBuffer, origin?: URL) {
 	if ("window" in globalThis)
 		origin ??= new URL(decodeUrl(location.href));
 
-	let before = performance.now();
+	const before = performance.now();
 	if (typeof js === "string") {
-		js = rewrite_js(js, origin.toString());
+		js = new TextDecoder().decode(rewrite_js(js, origin.toString()));
 	} else {
-		js = new TextDecoder().decode(js);
+		js = rewrite_js_from_arraybuffer(new Uint8Array(js), origin.toString());
-		js = rewrite_js(js, origin.toString());
 	}
-	let after = performance.now();
+	const after = performance.now();
 
 	console.log("Rewrite took", Math.floor((after - before) * 10) / 10, "ms");
-	//
-	// let offset = 0;
-	//
-	// for (const rewrite of rewrites) {
-	// 	if (rewrite.genericchange) {
-	// 		let change = rewrite.genericchange;
-	// 		let start = change.span.start + offset;
-	// 		let end = change.span.end + offset;
-	// 		let len = end - start;
-	//
-	// 		js = js.slice(0, start) + change.text + js.slice(end);
-	//
-	// 		offset += change.text.length - len;
-	// 	}
-	// }
-	// console.log(js)
-	//
-	return js;
 
-	// console.log(f)
+	return js;
-	//
-	// return f
-	// try {
-	// 	const ast = parseModule(js, {
-	// 		module: true,
-	// 		webcompat: true,
-	// 	});
-	//
-	// 	const identifierList = [
-	// 		"window",
-	// 		"self",
-	// 		"globalThis",
-	// 		"this",
-	// 		"parent",
-	// 		"top",
-	// 		"location",
-	// 	];
-	//
-	// 	const customTraveler = makeTraveler({
-	// 		ImportDeclaration: (node: ESTree.ImportDeclaration) => {
-	// 			node.source.value = encodeUrl(node.source.value as string, origin);
-	// 		},
-	//
-	// 		ImportExpression: (node: ESTree.ImportExpression) => {
-	// 			if (node.source.type === "Literal") {
-	// 				node.source.value = encodeUrl(node.source.value as string, origin);
-	// 			} else if (node.source.type === "Identifier") {
-	// 				// this is for things that import something like
-	// 				// const moduleName = "name";
-	// 				// await import(moduleName);
-	// 				node.source.name = `__wrapImport(${node.source.name})`;
-	// 			}
-	// 		},
-	//
-	// 		ExportAllDeclaration: (node: ESTree.ExportAllDeclaration) => {
-	// 			node.source.value = encodeUrl(node.source.value as string, origin);
-	// 		},
-	//
-	// 		ExportNamedDeclaration: (node: ESTree.ExportNamedDeclaration) => {
-	// 			// strings are Literals in ESTree syntax but these will always be strings
-	// 			if (node.source)
-	// 				node.source.value = encodeUrl(node.source.value as string, origin);
-	// 		},
-	//
-	// 		MemberExpression: (node: ESTree.MemberExpression) => {
-	// 			if (
-	// 				node.object.type === "Identifier" &&
-	// 				identifierList.includes(node.object.name)
-	// 			) {
-	// 				node.object.name = `globalThis.$s(${node.object.name})`;
-	// 			}
-	// 		},
-	//
-	// 		AssignmentExpression: (node: ESTree.AssignmentExpression, more) => {
-	// 			if (
-	// 				node.left.type === "Identifier" &&
-	// 				identifierList.includes(node.left.name)
-	// 			) {
-	// 				node.left.name = `globalThis.$s(${node.left.name})`;
-	// 			}
-	//
-	// 			if (
-	// 				node.right.type === "Identifier" &&
-	// 				identifierList.includes(node.right.name)
-	// 			) {
-	// 				node.right.name = `globalThis.$s(${node.right.name})`;
-	// 			}
-	// 		},
-	// 		ArrayExpression: (node: ESTree.ArrayExpression) => {
-	// 			node.elements.forEach((element) => {
-	// 				if (element.type === "Identifier" && identifierList.includes(element.name)) {
-	// 					element.name = `globalThis.$s(${element.name})`;
-	// 				}
-	// 			});
-	// 		},
-	//
-	// 		VariableDeclarator: (node: ESTree.VariableDeclarator) => {
-	// 			if (
-	// 				node.init &&
-	// 				node.init.type === "Identifier" &&
-	// 				identifierList.includes(node.init.name)
-	// 			) {
-	// 				node.init.name = `globalThis.$s(${node.init.name})`;
-	// 			}
-	// 		},
-	// 	});
-	//
-	// 	customTraveler.go(ast);
-	//
-	// 	return generate(ast);
-	// } catch (e) {
-	// 	console.error(e);
-	// 	console.log(js);
-	//
-	// 	return js;
-	// }
 }

src/worker/fetch.ts

@@ -110,7 +110,7 @@ export async function swfetch(
 					}
 					break;
 				case "script":
-					responseBody = rewriteJs(await response.text(), url);
+					responseBody = rewriteJs(await response.arrayBuffer(), url);
 					// Disable threading for now, it's causing issues.
 					// responseBody = await this.threadpool.rewriteJs(await responseBody.arrayBuffer(), url.toString());
 					break;