feat(lib/modules/metro)!: more performance wins

lib/modules/package.json

@@ -24,9 +24,6 @@
     "./native": {
       "default": "./src/native.ts"
     },
-    "./preinit": {
-      "default": "./src/preinit.ts"
-    },
     "./types": {
       "types": "./src/types.ts"
     },
@@ -34,4 +31,4 @@
       "default": "./src/metro/_internal.ts"
     }
   }
-}
+}

lib/modules/src/metro/_internal.ts

@@ -1,4 +1,13 @@
 import {
+    global,
+    define as metroDefine,
+    metroImportAll,
+    metroImportDefault,
+    metroRequire,
+} from './custom'
+import { onModuleInitialized } from './subscriptions'
+import {
+    executeImportedPathSubscriptions,
     executeInitializeSubscriptions,
     executeRequireSubscriptions,
 } from './subscriptions/_internal'
@@ -38,79 +47,32 @@ export let mList: Metro.ModuleList
  */
 export function patchMetroDefine() {
     const defineKey = `${__METRO_GLOBAL_PREFIX__}__d` as const
-    const metroDefine = globalThis[defineKey]
 
-    // Actual patched define function
-    const defineRest: Metro.DefineFn = function define(factory, id, deps) {
-        mDeps[id] = deps!
-        mUninitialized.add(id)
-
-        metroDefine(
-            (_, __, ___, ____, moduleObject) => {
-                handleFactoryCall(factory, moduleObject)
-            },
-            id,
-            deps,
-        )
-    }
-
-    // Patched function for first __d call
+    // First __d call
     globalThis[defineKey] = function define(origFactory, id, deps) {
         // Clear the module list so we can keep in sync with Metro's
         mList = __c()
 
-        metroImportDefault = metroRequire.importDefault
-        metroImportAll = metroRequire.importAll
+        // Set own implementation of metroImportDefault and metroImportAll
+        metroRequire.importDefault = metroImportDefault
+        metroRequire.importAll = metroImportAll
 
-        // Call the actual patched function
-        defineRest(origFactory, id, deps)
-        // Set to the actual patched function so we can use it later
-        globalThis[defineKey] = defineRest
+        // Set to the our own implementation function
+        // And then call the function to define the first module
+        ;(globalThis[defineKey] = metroDefine)(origFactory, id, deps)
     }
 }
 
-/**
- * Returns whether the module has bad exports. If it does, it will be "blacklisted" to avoid filtering issues.
- *
- * **Which module exports are considered bad?** Anything not an object or function, or an empty object.
- *
- * @param exports The exports of the module.
- */
-export function isModuleExportsBad(exports: Metro.ModuleExports): boolean {
-    return (
-        // Nullish?
-        exports == null ||
-        // Isn't an object or function?
-        // - Number exports are not useful, usually just an asset ID
-        // - String, Boolean, Symbol, BigInt exports are not useful (who would do `module.exports = ...`?)
-        !(
-            (exports.__proto__ === Object.prototype &&
-                Reflect.ownKeys(exports).length) ||
-            exports.__proto__ === Function.prototype
-        ) ||
-        // Can't run isProxy() on this because this isn't your typical proxy:
-        // https://github.com/facebook/react-native/blob/master/packages/react-native/ReactCommon/react/nativemodule/core/ReactCommon/TurboModuleBinding.cpp
-        exports === nativeModuleProxy
-    )
-}
-
-const global = globalThis
-const metroRequire = __r
-// We can't destructure here because metroRequire.* is set a little later
-let metroImportDefault: Metro.RequireFn
-let metroImportAll: Metro.RequireFn
-
-function handleFactoryCall(
+// Why don't we use all the arguments from Metro.FactoryFn?
+// Because there's too many for Hermes to be able to its dedicated CallN function which only supports up to 4 arguments. (Call, Call1, Call2, Call3, Call4)
+export function handleFactoryCall(
     factory: Metro.FactoryFn,
     moduleObject: Metro.Module,
 ) {
-    const exports = moduleObject.exports
-    const id = moduleObject.id!
-
     const prevIId = mInitializingId
-    mInitializingId = id
+    mInitializingId = moduleObject.id!
 
-    executeRequireSubscriptions(id)
+    executeRequireSubscriptions(mInitializingId)
 
     try {
         factory(
@@ -119,16 +81,40 @@ function handleFactoryCall(
             metroImportDefault,
             metroImportAll,
             moduleObject,
-            exports,
-            mDeps[id],
+            moduleObject.exports,
+            mDeps[mInitializingId],
         )
 
+        const { exports: actualExports } = moduleObject
+
         // Add the module to the initialized set only if the factory doesn't error or the exports aren't bad
         // Don't use exports here, as modules can set module.exports to a different object
-        if (!isModuleExportsBad(moduleObject.exports)) mInitialized.add(id)
-        executeInitializeSubscriptions(id, moduleObject.exports)
+        if (actualExports != null) mInitialized.add(mInitializingId)
+
+        executeInitializeSubscriptions(mInitializingId, actualExports)
     } finally {
         mUninitialized.delete(mInitializingId)
         mInitializingId = prevIId
     }
 }
+
+/// MODULE PATCHES AND BLACKLISTS
+
+const ImportTrackerModuleId = 2
+
+onModuleInitialized(ImportTrackerModuleId, (_, exports) => {
+    const orig = exports.fileFinishedImporting
+    exports.fileFinishedImporting = (path: string) => {
+        orig(path)
+        const id = mInitializingId!
+        mImportedPaths.set(path, id)
+        executeImportedPathSubscriptions(id, path)
+    }
+})
+
+const NativeModuleProxyModuleId = 45
+
+// Exports nativeModuleProxy, which we want to blacklist
+onModuleInitialized(NativeModuleProxyModuleId, () => {
+    mInitialized.delete(NativeModuleProxyModuleId)
+})

lib/modules/src/metro/custom.ts

@@ -0,0 +1,65 @@
+/**
+ * Custom Metro-compatible methods.
+ */
+
+import { handleFactoryCall, mDeps, mList, mUninitialized } from './_internal'
+import type { Nullish } from '@revenge-mod/utils/types'
+import type { Metro } from '../types'
+
+export const global = globalThis
+export const metroRequire = global.__r
+
+// Roll our own define function implementation that replaces Metro's
+// This is much more performant than Metro's own define function, which makes initialization faster
+export function define(
+    factory: Metro.FactoryFn,
+    id: Metro.ModuleID,
+    dependencyMap: Metro.DependencyMap | Nullish,
+) {
+    mDeps[id] = dependencyMap!
+    mUninitialized.add(id)
+
+    const moduleObject = { exports: {} }
+
+    mList.set(id, {
+        dependencyMap: dependencyMap!,
+        factory: () => {
+            handleFactoryCall(factory, moduleObject)
+        },
+        publicModule: moduleObject,
+    } satisfies Omit<
+        Metro.ModuleDefinition<false>,
+        'importedDefault' | 'importedAll' | 'isInitialized' | 'hasError'
+    > as unknown as Metro.ModuleDefinition<false>)
+}
+
+export function metroImportDefault(
+    moduleId: Metro.ModuleID,
+): Metro.ModuleExports {
+    const mod = mList.get(moduleId)!
+    if (mod.importedDefault) return mod.importedDefault
+
+    const exports: Metro.ModuleExports = metroRequire(moduleId)
+    const importedDefault: Metro.ModuleExports = exports?.__esModule
+        ? exports.default
+        : exports
+
+    return (mod.importedDefault = importedDefault)
+}
+
+export function metroImportAll(moduleId: Metro.ModuleID): Metro.ModuleExports {
+    const mod = mList.get(moduleId)!
+    if (mod.importedAll) return mod.importedAll
+
+    const exports: Metro.ModuleExports = metroRequire(moduleId)
+    let importedAll: Metro.ModuleExports
+
+    if (exports?.__esModule) importedAll = exports
+    else {
+        importedAll = {}
+        if (exports) for (const key in exports) importedAll[key] = exports[key]
+        importedAll.default = exports
+    }
+
+    return (mod.importedAll = importedAll)
+}

lib/modules/src/metro/patches.ts

@@ -1,16 +0,0 @@
-import { byProps } from '../finders/filters'
-import { waitForModules } from '../finders/wait'
-import { mImportedPaths, mInitializingId } from './_internal'
-import { executeImportedPathSubscriptions } from './subscriptions/_internal'
-
-const unsubFFI = waitForModules(byProps('fileFinishedImporting'), exports => {
-    unsubFFI()
-
-    const orig = exports.fileFinishedImporting
-    exports.fileFinishedImporting = (path: string) => {
-        orig(path)
-        const id = mInitializingId!
-        mImportedPaths.set(path, id)
-        executeImportedPathSubscriptions(id, path)
-    }
-})

lib/modules/src/metro/utils.ts

@@ -52,5 +52,3 @@ export function isModuleExportBad(
         (!exp.__proto__ && '\u0001' in exp)
     )
 }
-
-export { isModuleExportsBad } from './_internal'

lib/modules/src/preinit.ts

@@ -1 +0,0 @@
-import './metro/patches'

src/preinit.ts

@@ -18,7 +18,6 @@ onModuleFirstRequired(IndexModuleId, function onIndexRequired() {
             nativeLoggingHook(`\u001b[31m--- PREINIT STAGE ---\u001b[0m`, 1)
 
         // Initialize preinit libraries
-        require('@revenge-mod/modules/preinit')
         require('@revenge-mod/react/preinit')
         require('@revenge-mod/assets/preinit')
         require('@revenge-mod/storage/preinit')