Mangle all TS interface fields with a directive

mangle.ts

// TODO: make sure this doesn't run if `git diff --stat HEAD -- src/` actually returns something

import * as path from 'node:path';
import { fileURLToPath } from 'node:url';

import { type PropertySignature, InterfaceDeclaration, Project, SyntaxKind } from 'ts-morph';

import { ShortIdent } from '../lib/mangler/shortident.js';

const root = path.join(fileURLToPath(import.meta.url), '../..');
const projectPath = path.join(root, 'tsconfig.json');

const project = new Project({
	tsConfigFilePath: projectPath,
	compilerOptions: {
		outDir: 'src2',
	},
});

const interfacesToMangle = new Map<
	InterfaceDeclaration,
	{ id: number; candidates: Map<string, PropertySignature>; reserved: string[] }
>();
const idToInterface = new Map<number, InterfaceDeclaration>();

let uid = 0;

// Go through all the source files
for (const file of project.getSourceFiles()) {
	// Go through the interface declarations within it
	for (const declaration of file.getInterfaces()) {
		// Find if the interface contains a directive to mangle one of its properties
		const candidates = new Map<string, PropertySignature>();
		const reserved: string[] = [];

		let valid = false;
		let shouldMangleAllMember = false;

		// Check if this declaration has a directive set, that means all properties should be mangled
		for (const comment of declaration.getLeadingCommentRanges()) {
			const text = comment.getText();

			if (text.includes('@mangle')) {
				shouldMangleAllMember = true;
				break;
			}
		}

		// Go through each member of the declaration, and find the right candidates
		for (const member of declaration.getMembers()) {
			if (!member.isKind(SyntaxKind.PropertySignature)) {
				continue;
			}

			const name = member.getName();
			let shouldMangleMember = shouldMangleAllMember;

			if (!shouldMangleMember) {
				for (const comment of member.getLeadingCommentRanges()) {
					const text = comment.getText();

					if (text.includes('@mangle')) {
						shouldMangleMember = true;
						break;
					}
				}
			}

			if (shouldMangleMember) {
				candidates.set(name, member);
				valid = true;
			} else {
				reserved.push(name);
			}
		}

		// Put it on a map for later processing.
		if (valid) {
			idToInterface.set(uid, declaration);
			interfacesToMangle.set(declaration, { id: uid++, candidates, reserved });
		}
	}
}

// Next, we'll go through the references of each interfaces to see if they're in a union,
// - We'll have to bail out if the union contains another interface but the same field name aren't
//   set to mangle.
// - We'll use this opportunity to connect all relating interfaces into a single ShortIdent instance
//   for use in renaming.
const pairs: [from: number, to: number][] = [];
const standalones: number[] = [];

loop: for (const [declaration, { id, candidates }] of interfacesToMangle) {
	let standalone = true;

	for (const reference of declaration.findReferencesAsNodes()) {
		const typeReference = reference.getParentIfKind(SyntaxKind.TypeReference);

		if (!typeReference) {
			continue;
		}

		const referenceParent = typeReference.getParent();

		if (referenceParent.isKind(SyntaxKind.UnionType)) {
			for (const typeNode of referenceParent.getTypeNodes()) {
				if (typeNode === typeReference) {
					continue;
				}

				const type = typeNode.getType();
				const symbol = type.getSymbol();

				if (!symbol) {
					console.log(
						`[warn] bailing out ${declaration.getName()}\n       unknown error, why is the symbol missing?`,
					);

					interfacesToMangle.delete(declaration);
					standalone = false;
					continue loop;
				}

				for (const decl of symbol.getDeclarations()) {
					if (!decl.isKind(SyntaxKind.InterfaceDeclaration)) {
						console.log(
							`[warn] bailing out ${declaration.getName()}\n       is in a union with unknown type ${decl.getKindName()}`,
						);

						interfacesToMangle.delete(declaration);
						standalone = false;
						continue loop;
					}

					const declInfo = interfacesToMangle.get(decl);
					let valid = true;

					for (const member of decl.getMembers()) {
						if (!member.isKind(SyntaxKind.PropertySignature)) {
							continue;
						}

						const name = member.getName();
						if (candidates.has(name)) {
							if (!declInfo || !declInfo.candidates.has(name)) {
								console.log(
									`[warn] bailing out ${declaration.getName()}.${name}\n       is in a union with ${decl.getName()} where "${name}" is not being mangled`,
								);

								candidates.delete(name);
								continue;
							}

							valid = true;
						}
					}

					if (valid && declInfo) {
						pairs.push([id, declInfo.id]);
						standalone = false;
					}
				}
			}
		}
	}

	if (standalone) {
		standalones.push(id);
	}
}

// Connect all relations together with breadth-first search,
// instantiate a ShortIdent instance for each group.
const bfs = <T>(v: T, pairs: [T, T][], visited: Set<T>) => {
	const queue: T[] = [];
	const group: T[] = [];

	queue.push(v);

	while (queue.length > 0) {
		v = queue.shift();

		if (!visited.has(v)) {
			visited.add(v);

			group.push(v);

			for (let idx = 0, len = pairs.length; idx < len; idx++) {
				const pair = pairs[idx];

				if (pair[0] === v && !visited.has(pair[1])) {
					queue.push(pair[1]);
				} else if (pair[1] === v && !visited.has(pair[0])) {
					queue.push(pair[0]);
				}
			}
		}
	}

	return group;
};

const visited = new Set<number>();

for (let i = 0, len = pairs.length; i < len; i += 1) {
	const pair = pairs[i];
	const u = pair[0];
	const v = pair[1];

	const src = !visited.has(u) ? u : !visited.has(v) ? v : undefined;

	if (src !== undefined) {
		const group = bfs(src, pairs, visited);

		const shortident = new ShortIdent('');
		const sortmap = new Map<string, number>();

		const groupedCandidates: [name: string, signature: PropertySignature][] = [];

		for (let j = 0, jl = group.length; j < jl; j++) {
			const id = group[j];
			const decl = idToInterface.get(id)!;

			const info = interfacesToMangle.get(decl)!;
			const reserved = info.reserved;
			const candidates = info.candidates;

			for (let k = 0, kl = reserved.length; k < kl; k++) {
				const name = reserved[k];
				shortident.reserve(name);
			}

			for (const [name, signature] of candidates) {
				if (sortmap.has(name)) {
					sortmap.set(name, sortmap.get(name)! + 1);
				} else {
					sortmap.set(name, 1);
				}

				groupedCandidates.push([name, signature]);
			}
		}

		groupedCandidates.sort((a, b) => sortmap.get(a[0])! - sortmap.get(b[0])!);

		for (const [name, signature] of groupedCandidates) {
			const short = shortident.next(name);
			signature.rename(short);
		}
	}
}

for (let i = 0, len = standalones.length; i < len; i++) {
	const id = standalones[i];
	const decl = idToInterface.get(id)!;

	const info = interfacesToMangle.get(decl)!;
	const reserved = info.reserved;
	const candidates = info.candidates;

	const shortident = new ShortIdent('');

	for (let k = 0, kl = reserved.length; k < kl; k++) {
		const name = reserved[k];
		shortident.reserve(name);
	}

	for (const [name, signature] of candidates) {
		const short = shortident.next(name);
		signature.rename(short);
	}
}

// Save to disk.
await project.save();