feat(metro): assets caching

src/lib/api/assets.ts

@@ -1,24 +1,47 @@
+// TODO: Rewrite the whole thing to fit new metro :3
+
 import { after } from "@lib/api/patcher";
-import { assets } from "@metro/common";
+import { getImportingModuleId, requireModule } from "@lib/metro";
+import { getMetroCache, registerAssetCacheId } from "@lib/metro/caches";
 
 export const all: Record<string, Asset> = {};
 
 export interface Asset {
     name: string;
     id: number;
+    moduleId: number;
 }
 
+interface AssetModule {
+    registerAsset(assetDefinition: any): number;
+    getAssetByID(id: number): any;
+}
+
+let assetsModule: AssetModule;
+
 /**
  * @internal
  */
-export function patchAssets() {
-    const unpatch = after("registerAsset", assets, (args: Asset[], id: number) => {
+export function patchAssets(module: AssetModule) {
+    if (assetsModule) return;
+    assetsModule = module;
+
+    const unpatch = after("registerAsset", assetsModule, (args: Asset[], id: number) => {
         const asset = args[0];
-        all[asset.name] = { ...asset, id: id };
+
+        const moduleId = Number(getImportingModuleId());
+
+        registerAssetCacheId(asset.name, moduleId);
+
+        all[asset.name] = {
+            ...asset,
+            id,
+            moduleId
+        };
     });
 
     for (let id = 1; ; id++) {
-        const asset = assets.getAssetByID(id);
+        const asset = assetsModule.getAssetByID(id);
         if (!asset) break;
         if (all[asset.name]) continue;
         all[asset.name] = { ...asset, id: id };
@@ -29,5 +52,5 @@ export function patchAssets() {
 
 export const find = (filter: (a: any) => void): Asset | null | undefined => Object.values(all).find(filter);
 export const getAssetByName = (name: string): Asset => all[name];
-export const getAssetByID = (id: number): Asset => assets.getAssetByID(id);
-export const getAssetIDByName = (name: string) => all[name]?.id;
+export const getAssetByID = (id: number): Asset => assetsModule.getAssetByID(id);
+export const getAssetIDByName = (name: string) => all[name]?.id ?? requireModule(getMetroCache().assetsCache[name]);

src/lib/metro/caches.ts

@@ -1,17 +1,17 @@
 import { ClientInfoManager, MMKVManager } from "@lib/api/native/modules";
 import { throttle } from "@lib/utils/throttle";
 import SparkMD5 from "spark-md5";
-import { before } from "spitroast";
 
 // fucked types right here
 interface MetroCacheStore {
-    v: 1;
+    v: 4;
     buildNumber: number;
     findCache: Record<string, Record<number | "_", 1 | undefined> | undefined>;
     polyfillCache: Record<string, Record<number, 1 | undefined> | undefined>;
+    assetsCache: Record<string, number>;
 }
 
-const PYON_METRO_CACHE_KEY = "__bunny_metro_cache_key_v1__";
+const PYON_METRO_CACHE_KEY = "__bunny_metro_cache_key_v4__";
 
 let metroCache = null as unknown as MetroCacheStore;
 
@@ -22,21 +22,25 @@ window.getMetroCache = getMetroCache;
 
 export function getFuncUniqCall() {
     const { stack } = new Error();
-    return SparkMD5.hash(stack!!);
+    return SparkMD5.hash(stack!);
 }
 
 function buildInitCache() {
-    return metroCache = {
-        v: 1,
+    metroCache = {
+        v: 4,
         buildNumber: ClientInfoManager.Build,
         findCache: {},
-        polyfillCache: {}
+        polyfillCache: {},
+        assetsCache: {}
     } as const;
+
+    // Make sure all assets are cached
+    for (const id in window.modules) {
+        require("@metro/filters").requireModule(id);
+    }
 }
 
 export async function initMetroCache() {
-    before("refresh", MMKVManager, ([a]) => void (a.push(PYON_METRO_CACHE_KEY)), true);
-
     const rawCache = await MMKVManager.getItem(PYON_METRO_CACHE_KEY);
     if (rawCache == null) return buildInitCache();
 
@@ -65,3 +69,9 @@ export function registerPolyfillCacheId(name: string, moduleId: number, remove =
     else (metroCache.polyfillCache[name] ??= {})[moduleId] = 1;
     saveCache();
 }
+
+export function registerAssetCacheId(name: string, moduleId: number) {
+    if (isNaN(moduleId)) return;
+    metroCache.assetsCache[name] = moduleId;
+    saveCache();
+}

src/lib/metro/filters.ts

@@ -11,6 +11,11 @@ export type PropsFinderAll = <T extends PropertyKey>(...props: T[]) => PropIntel
 metroRequire(0);
 
 const blacklistedIds = new Set<String>();
+const noopHandler = () => undefined;
+const functionToString = Function.prototype.toString;
+
+let patchedInspectSource = false;
+let _importingModuleId: string | null = null;
 
 /** Makes the module associated with the specified ID non-enumberable. */
 function blacklistModule(id: string) {
@@ -22,13 +27,12 @@ function isBadExports(exports: any) {
     return !exports || exports === window || exports["<!@ pylix was here :fuyusquish: !@>"] === null;
 }
 
-const functionToString = Function.prototype.toString;
-
 for (const id in modules) {
     const metroModule = modules[id];
 
     if (metroModule!.factory) {
         instead("factory", metroModule, ((args: Parameters<Metro.FactoryFn>, origFunc: Metro.FactoryFn) => {
+            _importingModuleId = id;
             const { 1: metroRequire, 4: moduleObject } = args;
 
             args[2 /* metroImportDefault */] = id => {
@@ -52,19 +56,21 @@ for (const id in modules) {
             } else {
                 blacklistModule(id);
             }
+            _importingModuleId = null;
         }) as any); // If only spitroast had better types
     }
 }
 
-const noopHandler = () => undefined;
-let patchedInspectSource = false;
-
 function onModuleRequire(moduleExports: any, id: Metro.ModuleID) {
     // Temporary
     moduleExports.initSentry &&= () => undefined;
     if (moduleExports.default?.track && moduleExports.default.trackMaker)
         moduleExports.default.track = () => Promise.resolve();
 
+    if (moduleExports.registerAsset) {
+        require("@lib/api/assets").patchAssets(moduleExports);
+    }
+
     // There are modules registering the same native component
     if (moduleExports?.default?.name === "requireNativeComponent") {
         instead("default", moduleExports, (args, origFunc) => {
@@ -118,6 +124,10 @@ function onModuleRequire(moduleExports: any, id: Metro.ModuleID) {
     }
 }
 
+export function getImportingModuleId() {
+    return _importingModuleId;
+}
+
 export function requireModule(id: Metro.ModuleID) {
     if (blacklistedIds.has(String(id))) return undefined;