feat(lib/modules/metro): module segments support

lib/modules/src/metro/patches.ts

@@ -21,6 +21,60 @@ export const mDeps = new Map<Metro.ModuleID, Metro.DependencyMap>()
 
 export const mList: RevengeMetro.ModuleList = new Map()
 
+/**
+ * Metro allows a bundle to be split into segments.
+ *
+ * A segment registers itself with `globalThis.__registerSegment(segmentId, moduleDefiner, moduleIds?)`:
+ * - `moduleDefiner(id)` is a callback that triggers the (deferred) `__d(...)`
+ *   call for module `id` belonging to this segment.
+ * - `moduleIds` (optional, non-zero segments) is the list of module IDs that
+ *   live in this segment, so the runtime can route requires to the right
+ *   definer before the segment has actually defined them.
+ *
+ * When `metroRequire(id)` is called and `id` is not yet in `mList`, we look up
+ * its segment, invoke that segment's definer (which is expected to call
+ * `__d(...)` synchronously), and retry the lookup.
+ *
+ * @see https://github.com/facebook/metro/blob/6d63660/packages/metro-runtime/src/polyfills/require.js#L387
+ */
+export const mSegmentDefiners: Array<
+    ((moduleId: Metro.ModuleID) => void) | undefined
+> = []
+export const mModuleIdToSegmentId = new Map<Metro.ModuleID, number>()
+
+const registerSegment: Metro.RegisterSegmentFn = (
+    segmentId,
+    moduleDefiner,
+    moduleIds,
+) => {
+    mSegmentDefiners[segmentId] = moduleDefiner
+    if (moduleIds) {
+        for (const id of moduleIds) {
+            if (!mList.has(id) && !mModuleIdToSegmentId.has(id)) {
+                mModuleIdToSegmentId.set(id, segmentId)
+            }
+        }
+    }
+}
+
+/**
+ * Resolve a module that is not yet present in `mList` by invoking its segment definer (if available)
+ *
+ * @returns The freshly-registered definition or `undefined`.
+ */
+export function loadModuleFromSegment(
+    moduleId: Metro.ModuleID,
+): RevengeMetro.ModuleDefinition | undefined {
+    const segmentId = mModuleIdToSegmentId.get(moduleId) ?? 0
+    const definer = mSegmentDefiners[segmentId]
+    if (!definer) return undefined
+
+    definer(moduleId)
+    mModuleIdToSegmentId.delete(moduleId)
+
+    return mList.get(moduleId)
+}
+
 const metroDefine = (
     factory: Metro.FactoryFn,
     id: Metro.ModuleID,
@@ -53,17 +107,22 @@ const metroDefine = (
  *    - Why don't we patch it here? Because we need to know the value of `__METRO_GLOBAL_PREFIX__` first
  *    - And since Revenge runs before everything else, we need to patch in the next steps:
  *      - We chose to patch it in 4. because we can access all the global functions at that point.
- * 3. `__r = function metroRequire(...) {}`
- * 4. `__c = function clear() {}`
+ * 3. `__r`, `__d`, `__c` is set
+ * 4. `__registerSegment` is set
  *    - **PATCH**: Set `__d` to our own implementation.
+ *    - **PATCH**: Set `__registerSegment` to our own implementation.
  *
  * 5. `clear()`
  *    - Metro clears the module list directly with `clear()` and not `__c()`.
  * 6. `metroRequire.importDefault = ...`, `metroRequire.importAll = ...`
- * 7. `__d(..., 0, [...])`
+ * 7. `global.__registerSegment = registerSegment`
+ *    - **PATCH**: Replace `__registerSegment` with our own implementation so that
+ *      segment-loaded modules become resolvable through our runtime.
+ * 8. `__d(..., 0, [...])`
  *    - The first module is defined with ID 0, which is the index module.
  *    - **PATCH**: Override the `importDefault` and `importAll` functions in `__r`.
  * #. `__d` is called with subsequent module definitions
+ * #. Segments are loaded and registered via `__registerSegment(...)`
  */
 const defineKey = `${__METRO_GLOBAL_PREFIX__}__d` as const
 
@@ -79,6 +138,8 @@ globalThis[defineKey] = function define(origFactory, id, deps) {
     metroDefine(origFactory, id, deps)
 }
 
+globalThis.__registerSegment = registerSegment
+
 // Why don't we use all the arguments from Metro.FactoryFn?
 // Because there's too many for Hermes to be able to use its dedicated CallN operation
 // which only supports up to 4 arguments. (Call, Call1, Call2, Call3, Call4)
@@ -127,13 +188,17 @@ function handleFactoryCall(
             }
         }
     } catch (e) {
+        const msg = `Module ${mInitializingId} failed to initialize:\n\n${getErrorStack(e)}`
+
         if (__DEV__) {
-            callBridgeMethodSync('revenge.alertError', [
-                `Module ${mInitializingId} failed to initialize:\n\n${getErrorStack(e)}`,
-                FullVersion,
-            ])
+            callBridgeMethodSync('revenge.alertError', [msg, FullVersion])
         } else {
-            throw e
+            // So... it wasn't a great idea to throw, Discord has pushed a broken build that has some failing modules
+            // Vanilla Metro would swallow the error and just return an empty object as the exports..., insanity
+            // throw e
+            moduleObject.exports = {}
+            cacheBlacklistedModule(mInitializingId)
+            nativeLoggingHook(msg, 2)
         }
     } finally {
         mUninitialized.delete(mInitializingId)

lib/modules/src/metro/runtime.ts

@@ -5,7 +5,7 @@
  * Also avoids cloning exports, allowing for patches to be applied directly without checking for clones.
  */
 
-import { mList } from './patches'
+import { loadModuleFromSegment, mList } from './patches'
 import { executeInitializeSubscriptions } from './subscriptions/_internal'
 import type { Metro } from '../types'
 
@@ -21,7 +21,12 @@ const NotInitializedOrInitializingMask = ~InitializedOrInitializing
 export const global = globalThis
 
 export const metroRequire = (moduleId => {
-    const mod = mList.get(moduleId)!
+    let mod = mList.get(moduleId)
+
+    // Module isn't registered yet, maybe it's in another segment that hasn't been loaded yet.
+    if (!mod) mod = loadModuleFromSegment(moduleId)
+    if (!mod) throw new Error(`Requiring unknown module: ${moduleId}`)
+
     let { flags, module: moduleObject } = mod
 
     if (flags & InitializedOrInitializing) return moduleObject!.exports
@@ -59,11 +64,13 @@ export const metroRequire = (moduleId => {
 global.__r = metroRequire
 
 export const metroImportDefault: Metro.RequireFn = moduleId => {
+    const exports = metroRequire(moduleId)
+
+    // metroRequire may have lazily registered the module via a segment, so we
+    // look up the definition after the require call.
     const mod = mList.get(moduleId)!
     if (mod.flags & HasImportedDefault) return mod.importedDefault
 
-    const exports = metroRequire(moduleId)
-
     mod.flags |= HasImportedDefault
 
     return (mod.importedDefault = exports?.__esModule
@@ -72,10 +79,11 @@ export const metroImportDefault: Metro.RequireFn = moduleId => {
 }
 
 export const metroImportAll: Metro.RequireFn = moduleId => {
+    const exports = metroRequire(moduleId)
+
     const mod = mList.get(moduleId)!
     if (mod.flags & HasImportedAll) return mod.importedAll
 
-    const exports = metroRequire(moduleId)
     // Our implementation doesn't match Metro's because we modify the exports directly instead of cloning
     // But this is why ours is superior, it allows patching the exports without needing to do it more than a single time
     if (!exports?.__esModule) exports.default = exports

lib/modules/src/types.ts

@@ -18,6 +18,12 @@ export namespace Metro {
         dependencyMap: DependencyMap,
     ) => void
 
+    export type RegisterSegmentFn = (
+        segmentId: number,
+        moduleDefiner: (moduleId: ModuleID) => void,
+        moduleIds?: ReadonlyArray<ModuleID> | null,
+    ) => void
+
     export type ModuleID = number
 
     export interface ModuleDefinition<Initialized = boolean> {

src/index.ts

@@ -11,13 +11,13 @@ globalThis.ErrorUtils = {
     reportFatalError: onError,
 }
 
-Object.defineProperty(globalThis, '__c', {
+Object.defineProperty(globalThis, '__registerSegment', {
     configurable: true,
-    set(clear: Metro.ClearFn) {
+    set(registerSegment: Metro.RegisterSegmentFn) {
         // @ts-expect-error
         // biome-ignore lint/performance/noDelete: Prevent infinite set loop
-        delete globalThis.__c
-        globalThis.__c = clear
+        delete globalThis.__registerSegment
+        globalThis.__registerSegment = registerSegment
 
         // @as-require
         import('./preinit')

types/globals.d.ts

@@ -51,7 +51,6 @@ declare global {
 /// HERMES PROMISES
 
 declare global {
-    // biome-ignore lint/correctness/noUnusedVariables: Type parameter names must match
     interface Promise<T> {
         /// PROMISE POLYFILLS FROM: https://github.com/then/promise
         /// AND: https://github.com/facebook/hermes/blob/main/lib/InternalBytecode/01-Promise.js
@@ -101,6 +100,7 @@ declare global {
         importAll: Metro.ImportAllFn
     }
     var __c: Metro.ClearFn
+    var __registerSegment: Metro.RegisterSegmentFn
 }
 
 /// REACT NATIVE