Skip to content

Commit

Permalink
Add shortcut for lazy loading chunks
Browse files Browse the repository at this point in the history
  • Loading branch information
Nuckyz committed Jun 1, 2024
1 parent 78fd37a commit c4c92ed
Show file tree
Hide file tree
Showing 6 changed files with 189 additions and 160 deletions.
11 changes: 9 additions & 2 deletions scripts/generateReport.ts
Original file line number Diff line number Diff line change
Expand Up @@ -242,16 +242,23 @@ page.on("console", async e => {
});

break;
case "LazyChunkLoader:":
console.error(await getText());

switch (message) {
case "A fatal error occurred:":
process.exit(1);
}
case "Reporter:":
console.error(await getText());

switch (message) {
case "A fatal error occurred:":
process.exit(1);
case "Webpack Find Fail:":
process.exitCode = 1;
report.badWebpackFinds.push(otherMessage);
break;
case "A fatal error occurred:":
process.exit(1);
case "Finished test":
await browser.close();
await printReport();
Expand Down
167 changes: 167 additions & 0 deletions src/debug/loadLazyChunks.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,167 @@
/*
* Vencord, a Discord client mod
* Copyright (c) 2024 Vendicated and contributors
* SPDX-License-Identifier: GPL-3.0-or-later
*/

import { Logger } from "@utils/Logger";
import { canonicalizeMatch } from "@utils/patches";
import * as Webpack from "@webpack";
import { wreq } from "@webpack";

const LazyChunkLoaderLogger = new Logger("LazyChunkLoader");

export async function loadLazyChunks() {
try {
LazyChunkLoaderLogger.log("Loading all chunks...");

const validChunks = new Set<string>();
const invalidChunks = new Set<string>();
const deferredRequires = new Set<string>();

let chunksSearchingResolve: (value: void | PromiseLike<void>) => void;
const chunksSearchingDone = new Promise<void>(r => chunksSearchingResolve = r);

// True if resolved, false otherwise
const chunksSearchPromises = [] as Array<() => boolean>;

const LazyChunkRegex = canonicalizeMatch(/(?:(?:Promise\.all\(\[)?(\i\.e\("[^)]+?"\)[^\]]*?)(?:\]\))?)\.then\(\i\.bind\(\i,"([^)]+?)"\)\)/g);

async function searchAndLoadLazyChunks(factoryCode: string) {
const lazyChunks = factoryCode.matchAll(LazyChunkRegex);
const validChunkGroups = new Set<[chunkIds: string[], entryPoint: string]>();

// Workaround for a chunk that depends on the ChannelMessage component but may be be force loaded before
// the chunk containing the component
const shouldForceDefer = factoryCode.includes(".Messages.GUILD_FEED_UNFEATURE_BUTTON_TEXT");

await Promise.all(Array.from(lazyChunks).map(async ([, rawChunkIds, entryPoint]) => {
const chunkIds = rawChunkIds ? Array.from(rawChunkIds.matchAll(Webpack.ChunkIdsRegex)).map(m => m[1]) : [];

if (chunkIds.length === 0) {
return;
}

let invalidChunkGroup = false;

for (const id of chunkIds) {
if (wreq.u(id) == null || wreq.u(id) === "undefined.js") continue;

const isWasm = await fetch(wreq.p + wreq.u(id))
.then(r => r.text())
.then(t => (IS_WEB && t.includes(".module.wasm")) || !t.includes("(this.webpackChunkdiscord_app=this.webpackChunkdiscord_app||[]).push"));

if (isWasm && IS_WEB) {
invalidChunks.add(id);
invalidChunkGroup = true;
continue;
}

validChunks.add(id);
}

if (!invalidChunkGroup) {
validChunkGroups.add([chunkIds, entryPoint]);
}
}));

// Loads all found valid chunk groups
await Promise.all(
Array.from(validChunkGroups)
.map(([chunkIds]) =>
Promise.all(chunkIds.map(id => wreq.e(id as any).catch(() => { })))
)
);

// Requires the entry points for all valid chunk groups
for (const [, entryPoint] of validChunkGroups) {
try {
if (shouldForceDefer) {
deferredRequires.add(entryPoint);
continue;
}

if (wreq.m[entryPoint]) wreq(entryPoint as any);
} catch (err) {
console.error(err);
}
}

// setImmediate to only check if all chunks were loaded after this function resolves
// We check if all chunks were loaded every time a factory is loaded
// If we are still looking for chunks in the other factories, the array will have that factory's chunk search promise not resolved
// But, if all chunk search promises are resolved, this means we found every lazy chunk loaded by Discord code and manually loaded them
setTimeout(() => {
let allResolved = true;

for (let i = 0; i < chunksSearchPromises.length; i++) {
const isResolved = chunksSearchPromises[i]();

if (isResolved) {
// Remove finished promises to avoid having to iterate through a huge array everytime
chunksSearchPromises.splice(i--, 1);
} else {
allResolved = false;
}
}

if (allResolved) chunksSearchingResolve();
}, 0);
}

Webpack.factoryListeners.add(factory => {
let isResolved = false;
searchAndLoadLazyChunks(factory.toString()).then(() => isResolved = true);

chunksSearchPromises.push(() => isResolved);
});

for (const factoryId in wreq.m) {
let isResolved = false;
searchAndLoadLazyChunks(wreq.m[factoryId].toString()).then(() => isResolved = true);

chunksSearchPromises.push(() => isResolved);
}

await chunksSearchingDone;

// Require deferred entry points
for (const deferredRequire of deferredRequires) {
wreq!(deferredRequire as any);
}

// All chunks Discord has mapped to asset files, even if they are not used anymore
const allChunks = [] as string[];

// Matches "id" or id:
for (const currentMatch of wreq!.u.toString().matchAll(/(?:"(\d+?)")|(?:(\d+?):)/g)) {
const id = currentMatch[1] ?? currentMatch[2];
if (id == null) continue;

allChunks.push(id);
}

if (allChunks.length === 0) throw new Error("Failed to get all chunks");

// Chunks that are not loaded (not used) by Discord code anymore
const chunksLeft = allChunks.filter(id => {
return !(validChunks.has(id) || invalidChunks.has(id));
});

await Promise.all(chunksLeft.map(async id => {
const isWasm = await fetch(wreq.p + wreq.u(id))
.then(r => r.text())
.then(t => (IS_WEB && t.includes(".module.wasm")) || !t.includes("(this.webpackChunkdiscord_app=this.webpackChunkdiscord_app||[]).push"));

// Loads and requires a chunk
if (!isWasm) {
await wreq.e(id as any);
if (wreq.m[id]) wreq(id as any);
}
}));

LazyChunkLoaderLogger.log("Finished loading all chunks!");
} catch (e) {
LazyChunkLoaderLogger.log("A fatal error occurred:", e);
}
}
164 changes: 8 additions & 156 deletions src/debug/runReporter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,171 +5,23 @@
*/

import { Logger } from "@utils/Logger";
import { canonicalizeMatch } from "@utils/patches";
import * as Webpack from "@webpack";
import { wreq } from "@webpack";
import { patches } from "plugins";

import { loadLazyChunks } from "./loadLazyChunks";


const ReporterLogger = new Logger("Reporter");

async function runReporter() {
ReporterLogger.log("Starting test...");

try {
const validChunks = new Set<string>();
const invalidChunks = new Set<string>();
const deferredRequires = new Set<string>();

let chunksSearchingResolve: (value: void | PromiseLike<void>) => void;
const chunksSearchingDone = new Promise<void>(r => chunksSearchingResolve = r);

// True if resolved, false otherwise
const chunksSearchPromises = [] as Array<() => boolean>;

const LazyChunkRegex = canonicalizeMatch(/(?:(?:Promise\.all\(\[)?(\i\.e\("[^)]+?"\)[^\]]*?)(?:\]\))?)\.then\(\i\.bind\(\i,"([^)]+?)"\)\)/g);

async function searchAndLoadLazyChunks(factoryCode: string) {
const lazyChunks = factoryCode.matchAll(LazyChunkRegex);
const validChunkGroups = new Set<[chunkIds: string[], entryPoint: string]>();

// Workaround for a chunk that depends on the ChannelMessage component but may be be force loaded before
// the chunk containing the component
const shouldForceDefer = factoryCode.includes(".Messages.GUILD_FEED_UNFEATURE_BUTTON_TEXT");

await Promise.all(Array.from(lazyChunks).map(async ([, rawChunkIds, entryPoint]) => {
const chunkIds = rawChunkIds ? Array.from(rawChunkIds.matchAll(Webpack.ChunkIdsRegex)).map(m => m[1]) : [];

if (chunkIds.length === 0) {
return;
}

let invalidChunkGroup = false;

for (const id of chunkIds) {
if (wreq.u(id) == null || wreq.u(id) === "undefined.js") continue;

const isWasm = await fetch(wreq.p + wreq.u(id))
.then(r => r.text())
.then(t => (IS_WEB && t.includes(".module.wasm")) || !t.includes("(this.webpackChunkdiscord_app=this.webpackChunkdiscord_app||[]).push"));

if (isWasm && IS_WEB) {
invalidChunks.add(id);
invalidChunkGroup = true;
continue;
}

validChunks.add(id);
}

if (!invalidChunkGroup) {
validChunkGroups.add([chunkIds, entryPoint]);
}
}));

// Loads all found valid chunk groups
await Promise.all(
Array.from(validChunkGroups)
.map(([chunkIds]) =>
Promise.all(chunkIds.map(id => wreq.e(id as any).catch(() => { })))
)
);

// Requires the entry points for all valid chunk groups
for (const [, entryPoint] of validChunkGroups) {
try {
if (shouldForceDefer) {
deferredRequires.add(entryPoint);
continue;
}

if (wreq.m[entryPoint]) wreq(entryPoint as any);
} catch (err) {
console.error(err);
}
}
ReporterLogger.log("Starting test...");

// setImmediate to only check if all chunks were loaded after this function resolves
// We check if all chunks were loaded every time a factory is loaded
// If we are still looking for chunks in the other factories, the array will have that factory's chunk search promise not resolved
// But, if all chunk search promises are resolved, this means we found every lazy chunk loaded by Discord code and manually loaded them
setTimeout(() => {
let allResolved = true;

for (let i = 0; i < chunksSearchPromises.length; i++) {
const isResolved = chunksSearchPromises[i]();

if (isResolved) {
// Remove finished promises to avoid having to iterate through a huge array everytime
chunksSearchPromises.splice(i--, 1);
} else {
allResolved = false;
}
}

if (allResolved) chunksSearchingResolve();
}, 0);
}

Webpack.beforeInitListeners.add(async () => {
ReporterLogger.log("Loading all chunks...");

Webpack.factoryListeners.add(factory => {
let isResolved = false;
searchAndLoadLazyChunks(factory.toString()).then(() => isResolved = true);

chunksSearchPromises.push(() => isResolved);
});

// setImmediate to only search the initial factories after Discord initialized the app
// our beforeInitListeners are called before Discord initializes the app
setTimeout(() => {
for (const factoryId in wreq.m) {
let isResolved = false;
searchAndLoadLazyChunks(wreq.m[factoryId].toString()).then(() => isResolved = true);

chunksSearchPromises.push(() => isResolved);
}
}, 0);
});

await chunksSearchingDone;

// Require deferred entry points
for (const deferredRequire of deferredRequires) {
wreq!(deferredRequire as any);
}

// All chunks Discord has mapped to asset files, even if they are not used anymore
const allChunks = [] as string[];

// Matches "id" or id:
for (const currentMatch of wreq!.u.toString().matchAll(/(?:"(\d+?)")|(?:(\d+?):)/g)) {
const id = currentMatch[1] ?? currentMatch[2];
if (id == null) continue;

allChunks.push(id);
}

if (allChunks.length === 0) throw new Error("Failed to get all chunks");

// Chunks that are not loaded (not used) by Discord code anymore
const chunksLeft = allChunks.filter(id => {
return !(validChunks.has(id) || invalidChunks.has(id));
});

await Promise.all(chunksLeft.map(async id => {
const isWasm = await fetch(wreq.p + wreq.u(id))
.then(r => r.text())
.then(t => (IS_WEB && t.includes(".module.wasm")) || !t.includes("(this.webpackChunkdiscord_app=this.webpackChunkdiscord_app||[]).push"));

// Loads and requires a chunk
if (!isWasm) {
await wreq.e(id as any);
if (wreq.m[id]) wreq(id as any);
}
}));
let loadLazyChunksResolve: (value: void | PromiseLike<void>) => void;
const loadLazyChunksDone = new Promise<void>(r => loadLazyChunksResolve = r);

ReporterLogger.log("Finished loading all chunks!");
Webpack.beforeInitListeners.add(() => loadLazyChunks().then(() => loadLazyChunksResolve()));
await loadLazyChunksDone;

for (const patch of patches) {
if (!patch.all) {
Expand Down
2 changes: 2 additions & 0 deletions src/plugins/consoleShortcuts/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import definePlugin, { PluginNative, StartAt } from "@utils/types";
import * as Webpack from "@webpack";
import { extract, filters, findAll, findModuleId, search } from "@webpack";
import * as Common from "@webpack/common";
import { loadLazyChunks } from "debug/loadLazyChunks";
import type { ComponentType } from "react";

const DESKTOP_ONLY = (f: string) => () => {
Expand Down Expand Up @@ -82,6 +83,7 @@ function makeShortcuts() {
wpsearch: search,
wpex: extract,
wpexs: (code: string) => extract(findModuleId(code)!),
loadLazyChunks: IS_DEV ? loadLazyChunks : () => { throw new Error("loadLazyChunks is dev only."); },

This comment has been minimized.

Copy link
@Vendicated

Vendicated Jun 1, 2024

Owner

it's always better to use inline requires for stuff like this as otherwise any top level code that could have side effects (such as const LazyChunkLoaderLogger = new Logger("LazyChunkLoader");) will be bundled anyway

loadLazyChunks: IS_DEV
	? require("debug/loadLazyChunks").loadLazyChunks
	: () => { throw new Error("loadLazyChunks is dev only."); }
find,
findAll: findAll,
findByProps,
Expand Down
Loading

0 comments on commit c4c92ed

Please sign in to comment.