prototype search

src/renderer/modules/webpack/filters.ts

@@ -1,4 +1,4 @@
-import { getExportsForProps } from "./get-modules";
+import { getExportsForProps, getFunctionForPrototypes } from "./get-modules";
 import { sourceStrings } from "./patch-load";
 import type { RawModule } from "../../../types";
 
@@ -10,6 +10,14 @@ export const byProps = <P extends PropertyKey = PropertyKey>(...props: P[]) => {
   return (m: RawModule) => typeof getExportsForProps(m.exports, props) !== "undefined";
 };
 
+/**
+ * Get a module that has all the given prototypes on one of its functions
+ * @param prototypes List of prototype names
+ */
+export const byPrototypes = <P extends PropertyKey = PropertyKey>(...prototypes: P[]) => {
+  return (m: RawModule) => typeof getFunctionForPrototypes(m.exports, prototypes) !== "undefined";
+};
+
 /**
  * Get a module whose source code matches the given string or RegExp
  * @param match String or RegExp to match in the module's source code

src/renderer/modules/webpack/get-modules.ts

@@ -1,4 +1,4 @@
-import type { Filter, GetModuleOptions, RawModule } from "src/types";
+import type { AnyFunction, Filter, GetModuleOptions, RawModule, WithPrototype } from "src/types";
 import { wpRequire } from "./patch-load";
 import { logError } from "./util";
 
@@ -27,7 +27,7 @@ export function getExports<T>(m: RawModule): T | undefined {
  */
 function* iterateModuleExports(m: unknown): IterableIterator<Record<PropertyKey, unknown>> {
   // if m is null or not an object/function, then it will obviously not have the props
-  // if if has no props, then it definitely has no children either
+  // if it has no props, then it definitely has no children either
   if (m && (typeof m === "object" || typeof m === "function")) {
     yield m as Record<PropertyKey, unknown>;
     for (const key in m) {
@@ -65,6 +65,62 @@ export function getExportsForProps<T, P extends PropertyKey = keyof T>(
   }
 }
 
+/**
+ * Iterates over an object and its top-level/second-level children that could have functions
+ * @param m Object (module exports) to iterate over
+ */
+function* iterateModuleFunctions(m: unknown): IterableIterator<AnyFunction> {
+  // if m is null or not an object/function, then it will obviously not be a function/object
+  // if it have no function then how would it have prototypes XD
+  if (m && (typeof m === "object" || typeof m === "function")) {
+    if (typeof m === "function") yield m as AnyFunction;
+    try {
+      // This could throw an error ("illegal invocation") if val === DOMTokenList.prototype
+      // and key === "length"
+      // There could be other cases too, hence this try-catch instead of a specific exclusion
+      if (typeof m === "object") {
+        for (const key in m) {
+          const val = (m as Record<PropertyKey, unknown>)[key];
+          if (val && typeof val === "function") {
+            yield val as AnyFunction;
+            continue;
+          }
+          if (typeof val === "object") {
+            for (const subkey in val) {
+              const subVal = (val as Record<PropertyKey, unknown>)[subkey];
+              if (subVal && typeof subVal === "function") {
+                yield subVal as AnyFunction;
+                continue;
+              }
+            }
+          }
+        }
+      }
+    } catch {
+      // ignore this export
+    }
+  }
+}
+
+/**
+ * Find an object in a module that has all the given properties. You will usually not need this function.
+ * @param m Module to search
+ * @param props Array of prototype names
+ * @returns Function that contains all the given prototypes (and any others), or undefined if not found
+ */
+export function getFunctionForPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(m: unknown, prototypes: P[]): T | undefined {
+  // Loop over the module and its exports at the top level
+  // Return the first thing that has all the indicated props
+  for (const exported of iterateModuleFunctions(m)) {
+    if (prototypes.every((p) => p in (exported.prototype as Record<P, AnyFunction>))) {
+      return exported as T;
+    }
+  }
+}
+
 // This doesn't have anywhere else to go
 
 export function getById<T>(id: number, raw?: false): T | undefined;

src/renderer/modules/webpack/helpers.ts

@@ -1,5 +1,11 @@
-import type { GetModuleOptions, RawModule, WaitForOptions } from "src/types";
-import { getExportsForProps, getModule } from "./get-modules";
+import type {
+  AnyFunction,
+  GetModuleOptions,
+  RawModule,
+  WaitForOptions,
+  WithPrototype,
+} from "src/types";
+import { getExportsForProps, getFunctionForPrototypes, getModule } from "./get-modules";
 import * as filters from "./filters";
 import Flux, { Store } from "../common/flux";
 import { waitForModule } from "./lazy";
@@ -165,6 +171,131 @@ export async function waitForProps<T, P extends PropertyKey = keyof T>(
   return getExportsForProps<T, P>(result as T, props)!;
 }
 
+// Get by prototypes
+
+export function getByPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(prototypes: P[], options?: { all?: false; raw?: false }): T | undefined;
+export function getByPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(prototypes: P[], options: { all: true; raw?: false }): T[];
+export function getByPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(prototypes: P[], options: { all?: false; raw: true }): RawModule<T> | undefined;
+export function getByPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(prototypes: P[], options: { all: true; raw: true }): Array<RawModule<T>>;
+export function getByPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(prototypes: P[], options?: { all: true; raw?: boolean }): T[] | Array<RawModule<T>>;
+export function getByPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(prototypes: P[], options: { all?: false; raw?: boolean }): T | RawModule<T> | undefined;
+export function getByPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(
+  prototypes: P[],
+  options: { all?: boolean; raw: true },
+): RawModule<T> | Array<RawModule<T>> | undefined;
+export function getByPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(prototypes: P, options: { all?: boolean; raw?: false }): T | T[] | undefined;
+export function getByPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(
+  prototypes: P[],
+  options?: { all?: boolean; raw?: boolean },
+): T | T[] | RawModule<T> | Array<RawModule<T>> | undefined;
+export function getByPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey[] = Array<keyof WithPrototype<T>>,
+>(...prototypes: P): T | undefined;
+
+/**
+ * Equivalent to `getModule(filters.byPrototypes(...props), options)`
+ *
+ * @see {@link filters.byPrototypes}
+ * @see {@link getModule}
+ */
+export function getByPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(
+  ...args: [P[], GetModuleOptions] | P[]
+): T | T[] | RawModule<T> | Array<RawModule<T>> | undefined {
+  const prototypes = (typeof args[0] === "string" ? args : args[0]) as P[];
+  const raw = typeof args[0] === "string" ? false : (args[1] as GetModuleOptions)?.raw;
+
+  const result =
+    typeof args.at(-1) === "object"
+      ? getModule<T>(filters.byPrototypes(...prototypes), args.at(-1) as GetModuleOptions)
+      : getModule<T>(filters.byPrototypes(...prototypes));
+
+  if (raw || typeof result === "undefined") {
+    return result as RawModule<T> | undefined;
+  }
+
+  if (result instanceof Array) {
+    // @ts-expect-error TypeScript isn't going to infer types based on the raw variable, so this is fine
+    return result.map((m) => getFunctionForPrototypes(m, prototypes));
+  }
+
+  return getFunctionForPrototypes<T, P>(result, prototypes);
+}
+
+// Wait for prototypes
+
+export function waitForPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(prototypes: P[], options: WaitForOptions & { raw?: false }): Promise<T>;
+export function waitForPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(prototypes: P[], options: WaitForOptions & { raw: true }): Promise<RawModule<T>>;
+export function waitForPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(prototypes: P[], options?: WaitForOptions): Promise<T | RawModule<T>>;
+export function waitForPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(...prototypes: P[]): Promise<T>;
+
+/**
+ * Like {@link getByPrototypes} but waits for the module to be loaded.
+ *
+ * @see {@link getByPrototypes}
+ * @see {@link waitForModule}
+ */
+export async function waitForPrototypes<
+  T extends AnyFunction,
+  P extends PropertyKey = keyof WithPrototype<T>,
+>(...args: [P[], WaitForOptions] | P[]): Promise<T | RawModule<T>> {
+  const prototypes = (typeof args[0] === "string" ? args : args[0]) as P[];
+  const raw = typeof args[0] === "string" ? false : (args[1] as WaitForOptions)?.raw;
+
+  const result = await (typeof args.at(-1) === "object"
+    ? waitForModule<T>(filters.byPrototypes(...prototypes), args.at(-1) as WaitForOptions)
+    : waitForModule<T>(filters.byPrototypes(...prototypes)));
+
+  if (raw) {
+    return result as RawModule<T>;
+  }
+
+  // We know this will always exist since filters.byPrototypes will always return a module that has the function with prototypes
+  return getFunctionForPrototypes<T, P>(result as T, prototypes)!;
+}
+
 // Get by value
 
 export function getByValue<T>(

src/renderer/modules/webpack/index.ts

@@ -2,7 +2,7 @@ export { waitForModule } from "./lazy";
 
 export { getFunctionBySource, getFunctionKeyBySource } from "./inner-search";
 
-export { getById, getExportsForProps, getModule } from "./get-modules";
+export { getById, getExportsForProps, getFunctionForPrototypes, getModule } from "./get-modules";
 
 /**
  * Filter functions to use with {@link getModule}