massive speedup in sourcemaps.ts

rewriter/Cargo.lock

@@ -346,8 +346,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7"
 dependencies = [
  "cfg-if",
+ "js-sys",
  "libc",
  "wasi",
+ "wasm-bindgen",
 ]
 
 [[package]]
@@ -980,6 +982,7 @@ version = "0.1.0"
 dependencies = [
  "boa_engine",
  "console_error_panic_hook",
+ "getrandom",
  "js-sys",
  "obfstr",
  "oxc_allocator",
@@ -987,6 +990,7 @@ dependencies = [
  "oxc_parser",
  "oxc_span",
  "oxc_syntax",
+ "rand",
  "serde",
  "serde-wasm-bindgen",
  "url",

rewriter/Cargo.toml

@@ -24,6 +24,7 @@ panic = "abort"
 
 [dependencies]
 console_error_panic_hook = "0.1.7"
+getrandom = { version = "0.2.15", features = ["js"] }
 js-sys = "0.3.69"
 obfstr = "0.4.3"
 oxc_allocator = "0.20.0"
@@ -31,6 +32,7 @@ oxc_ast = { version = "0.20.0", features = ["serialize"]}
 oxc_parser = { version = "0.20.0" }
 oxc_span = "0.20.0"
 oxc_syntax = "0.20.0"
+rand = "0.8.5"
 serde = "1.0.204"
 serde-wasm-bindgen = "0.6.5"
 url = "2.5.2"

rewriter/src/rewrite.rs

@@ -1,4 +1,5 @@
 use core::str;
+use std::str::from_utf8;
 
 use oxc_allocator::Allocator;
 use oxc_ast::{
@@ -19,8 +20,8 @@ enum JsChange {
 		span: Span,
 		text: String,
 	},
-	DebugInject {
-		span: Span,
+	SourceTag {
+		tagstart: u32,
 	},
 	Assignment {
 		name: String,
@@ -254,6 +255,16 @@ impl<'a> Visit<'a> for Rewriter {
 		walk::walk_object_expression(self, it);
 	}
 
+	fn visit_function_body(&mut self, it: &oxc_ast::ast::FunctionBody<'a>) {
+		// tag function for use in sourcemaps
+		if self.config.do_sourcemaps {
+			self.jschanges.push(JsChange::SourceTag {
+				tagstart: it.span.start,
+			});
+		}
+		walk::walk_function_body(self, it);
+	}
+
 	fn visit_return_statement(&mut self, it: &oxc_ast::ast::ReturnStatement<'a>) {
 		// if let Some(arg) = &it.argument {
 		// 	self.jschanges.push(JsChange::GenericChange {
@@ -391,6 +402,19 @@ const UNSAFE_GLOBALS: [&str; 9] = [
 	"eval",
 ];
 
+fn random_string() -> String {
+	use rand::{distributions::Alphanumeric, thread_rng, Rng};
+
+	from_utf8(
+		&thread_rng()
+			.sample_iter(&Alphanumeric)
+			.take(10)
+			.collect::<Vec<u8>>(),
+	)
+	.unwrap()
+	.to_string()
+}
+
 pub fn rewrite(js: &str, url: Url, config: Config) -> Vec<u8> {
 	let allocator = Allocator::default();
 	let source_type = SourceType::default();
@@ -406,6 +430,8 @@ pub fn rewrite(js: &str, url: Url, config: Config) -> Vec<u8> {
 
 	// dbg!(&program);
 
+	let sourcetag = random_string();
+
 	let mut ast_pass = Rewriter {
 		jschanges: Vec::new(),
 		base: url,
@@ -424,7 +450,7 @@ pub fn rewrite(js: &str, url: Url, config: Config) -> Vec<u8> {
 				rhsspan: _,
 				op: _,
 			} => entirespan.start,
-			JsChange::DebugInject { span } => span.start,
+			JsChange::SourceTag { tagstart } => *tagstart,
 		};
 		let b = match b {
 			JsChange::GenericChange { span, text: _ } => span.start,
@@ -434,7 +460,7 @@ pub fn rewrite(js: &str, url: Url, config: Config) -> Vec<u8> {
 				rhsspan: _,
 				op: _,
 			} => entirespan.start,
-			JsChange::DebugInject { span } => span.start,
+			JsChange::SourceTag { tagstart } => *tagstart,
 		};
 		a.cmp(&b)
 	});
@@ -514,11 +540,14 @@ pub fn rewrite(js: &str, url: Url, config: Config) -> Vec<u8> {
 
 				offset = entirespan.end as usize;
 			}
-			JsChange::DebugInject { span } => {
-				let start = span.start as usize;
+			JsChange::SourceTag { tagstart } => {
+				let start = *tagstart as usize;
 				buffer.extend_from_slice(unsafe { js.get_unchecked(offset..start) }.as_bytes());
 
-				offset = span.end as usize;
+				let inject = format!("/*scramtag {} {}*/", start, sourcetag);
+				buffer.extend_from_slice(inject.as_bytes());
+
+				offset = start;
 			}
 		}
 	}
@@ -527,8 +556,7 @@ pub fn rewrite(js: &str, url: Url, config: Config) -> Vec<u8> {
 	if ast_pass.config.do_sourcemaps {
 		sourcemap.extend_from_slice(b"],");
 		sourcemap.extend_from_slice(b"\"");
-		sourcemap
-			.extend_from_slice(json_escape_string(str::from_utf8(&buffer).unwrap()).as_bytes());
+		sourcemap.extend_from_slice(&sourcetag.as_bytes());
 		sourcemap.extend_from_slice(b"\");\n");
 
 		sourcemap.extend_from_slice(&buffer);

src/client/shared/sourcemaps.ts

@@ -1,17 +1,17 @@
 import { ScramjetClient } from "../client";
 
-const sourcemaps: {
-	source: string;
-	map: [string, number, number][];
-}[] = [];
+type Mapping = [string, number, number];
+
+const sourcemaps: Record<string, Mapping[]> = {};
 
 export const enabled = () => self.$scramjet.config.flags.sourcemaps;
 
+let t = 0;
 export default function (client: ScramjetClient, self: Self) {
 	// every script will push a sourcemap
 	Object.defineProperty(self, "$scramjet$pushsourcemap", {
-		value: (map, source) => {
-			sourcemaps.push({ map, source });
+		value: (maps: Mapping[], tag: string) => {
+			sourcemaps[tag] = maps;
 		},
 		enumerable: false,
 		writable: false,
@@ -22,44 +22,52 @@ export default function (client: ScramjetClient, self: Self) {
 	// this can lead to double rewrites which is bad
 	client.Proxy("Function.prototype.toString", {
 		apply(ctx) {
-			const stringified = ctx.fn.call(ctx.this);
+			let stringified: string = ctx.fn.call(ctx.this);
 			let newString = "";
 
-			// find the sourcemap, brute force, just check every file until the body of the function shows up
-			// it doesnt matter if there's multiple with the same content because it will be the same function
-			const sourcemap = sourcemaps.find(({ source }) =>
-				source.includes(stringified)
+			// every function rewritten will have a scramtag comment
+			// it will look like this:
+			// function name() /*scramtag [index] [tag] */ { ... }
+			const scramtag_ident = "/*scramtag ";
+			const scramtagstart = stringified.indexOf(scramtag_ident);
+
+			if (scramtagstart === -1) return ctx.return(stringified); // it's either a native function or something stolen from scramjet itself
+
+			// [index] holds the index of the first character in the scramtag (/)
+			const abstagindex = parseInt(
+				stringified
+					.substring(scramtagstart + scramtag_ident.length)
+					.split(" ")[0]
 			);
 
-			// i don't know what cases this would happen under, but it does
-			if (!sourcemap) return ctx.return(stringified);
-			const { source, map } = sourcemap;
+			// subtracting that from the index of the scramtag gives us the starting index of the function relative to the entire file
+			let absindex = abstagindex - scramtagstart;
 
-			// first we need to find the character # where the function starts relative to the *transformed* source
-			let starting = source.indexOf(stringified);
+			const scramtagend = stringified.indexOf("*/", scramtagstart);
+			const tag = stringified
+				.substring(scramtagstart + scramtag_ident.length, scramtagend)
+				.split(" ")[1];
 
-			const beforeFunctionRewrites = map.filter(
-				([str, start, end]) => start < starting
-			);
+			// delete the scramtag now that we're done with it
+			stringified =
+				stringified.slice(0, scramtagstart) +
+				stringified.slice(scramtagend + 2);
 
-			// map the offsets of the original source to the transformed source
-			for (const [str, start, end] of beforeFunctionRewrites) {
-				starting -= end - start - str.length;
-			}
+			const maps = sourcemaps[tag];
 
-			const relevantRewrites = map.filter(
+			const relevantRewrites = maps.filter(
 				([str, start, end]) =>
-					start >= starting && end <= starting + stringified.length
+					start >= absindex && end <= absindex + stringified.length
 			);
 
 			let i = 0;
 			let offset = 0;
 			for (const [str, start, end] of relevantRewrites) {
 				// ooh i should really document this before i forget how it works
-				newString += stringified.slice(i, start - starting + offset);
+				newString += stringified.slice(i, start - absindex + offset);
 				newString += str;
 				offset += end - start - str.length;
-				i = start - starting + offset + str.length;
+				i = start - absindex + offset + str.length;
 			}
 
 			return ctx.return(newString + stringified.slice(i));