New ns implementation using Proxy (#292)

src/Netscript/APIWrapper.ts

@@ -12,7 +12,6 @@ type APIFn = (...args: any[]) => void;
 type WrappedFn = (...args: unknown[]) => unknown;
 /** Type for internal, unwrapped ctx function that produces an APIFunction */
 type InternalFn<F extends APIFn> = (ctx: NetscriptContext) => ((...args: unknown[]) => ReturnType<F>) & F;
-type Key<API> = keyof API & string;
 
 export type ExternalAPI<API> = {
   [key in keyof API]: API[key] extends Enums
@@ -33,42 +32,6 @@ export type InternalAPI<API> = {
     ? InternalFn<API[key]>
     : InternalAPI<API[key]>;
 };
-/** Any of the possible values on a internal API layer */
-type InternalValues = Enums | ScriptArg[] | InternalFn<APIFn> | InternalAPI<unknown>;
-
-export class StampedLayer {
-  #workerScript: WorkerScript;
-  constructor(ws: WorkerScript, obj: ExternalAPI<unknown>) {
-    this.#workerScript = ws;
-    Object.setPrototypeOf(this, obj);
-  }
-  static wrapFunction<API>(eLayer: ExternalAPI<API>, internalFunc: InternalFn<APIFn>, tree: string[], key: Key<API>) {
-    const arrayPath = [...tree, key];
-    const functionPath = arrayPath.join(".");
-    function wrappedFunction(this: StampedLayer, ...args: unknown[]): unknown {
-      if (!this)
-        throw new Error(`
-ns.${functionPath} called with no this value.
-ns functions must be bound to ns if placed in a new
-variable. e.g.
-
-const ${key} = ns.${functionPath}.bind(ns);
-${key}(${JSON.stringify(args).replace(/^\[|\]$/g, "")});\n\n`);
-      const ctx = { workerScript: this.#workerScript, function: key, functionPath };
-      const func = internalFunc(ctx); //Allows throwing before ram chack
-      helpers.checkEnvFlags(ctx);
-      helpers.updateDynamicRam(ctx, getRamCost(...tree, key));
-      return func(...args);
-    }
-    Object.defineProperty(eLayer, key, { value: wrappedFunction, enumerable: true, writable: false });
-  }
-}
-Object.defineProperty(StampedLayer.prototype, "constructor", {
-  value: Object,
-  enumerable: false,
-  writable: false,
-  configurable: false,
-});
 
 export type NetscriptContext = {
   workerScript: WorkerScript;
@@ -76,28 +39,64 @@ export type NetscriptContext = {
   functionPath: string;
 };
 
-export function wrapAPILayer<API>(
-  eLayer: ExternalAPI<API>,
-  iLayer: InternalAPI<API>,
+export function NSProxy<API>(
+  ws: WorkerScript,
+  ns: InternalAPI<API>,
   tree: string[],
+  additionalData?: Record<string, unknown>,
 ): ExternalAPI<API> {
-  for (const [key, value] of Object.entries(iLayer) as [Key<API>, InternalValues][]) {
-    if (key === "enums") {
-      const enumObj = Object.freeze(cloneDeep(value as Enums));
-      for (const member of Object.values(enumObj)) Object.freeze(member);
-      (eLayer[key] as Enums) = enumObj;
-    } else if (key === "args") continue;
-    // Args only added on individual instances.
-    else if (typeof value === "function") {
-      StampedLayer.wrapFunction(eLayer, value as InternalFn<APIFn>, tree, key);
-    } else if (typeof value === "object") {
-      wrapAPILayer((eLayer[key] = {} as ExternalAPI<API>[Key<API>]), value, [...tree, key as string]);
-    } else {
-      console.warn(`Unexpected data while wrapping API.`, "tree:", tree, "key:", key, "value:", value);
+  const memoed: ExternalAPI<API> = Object.assign({} as ExternalAPI<API>, additionalData ?? {});
+
+  const handler = {
+    has(__target: unknown, key: string) {
+      return Reflect.has(ns, key);
+    },
+    ownKeys(__target: unknown) {
+      return Reflect.ownKeys(ns);
+    },
+    getOwnPropertyDescriptor(__target: unknown, key: string) {
+      if (!Reflect.has(ns, key)) return undefined;
+      return { value: this.get(__target, key, this), configurable: true, enumerable: true, writable: false };
+    },
+    defineProperty(__target: unknown, __key: unknown, __attrs: unknown) {
+      throw new TypeError("ns instances are not modifiable!");
+    },
+    get(__target: unknown, key: string, __receiver: any) {
+      const ours = memoed[key as keyof API];
+      if (ours) return ours;
+
+      const field = ns[key as keyof API];
+      if (!field) return field;
+
+      if (key === "enums") {
+        const enumObj = Object.freeze(cloneDeep(field as Enums));
+        for (const member of Object.values(enumObj)) Object.freeze(member);
+        return ((memoed[key as keyof API] as Enums) = enumObj);
+      }
+      if (typeof field === "function") {
+        const arrayPath = [...tree, key];
+        const functionPath = arrayPath.join(".");
+        const wrappedFunction = function (...args: unknown[]): unknown {
+          const ctx = { workerScript: ws, function: key, functionPath };
+          const func = field(ctx); //Allows throwing before ram chack
+          helpers.checkEnvFlags(ctx);
+          helpers.updateDynamicRam(ctx, getRamCost(...tree, key));
+          return func(...args);
+        };
+        return ((memoed[key as keyof API] as WrappedFn) = wrappedFunction);
+      }
+      if (typeof field === "object") {
+        // TODO unplanned: Make this work generically
+        return ((memoed[key as keyof API] as unknown) = NSProxy(ws, field as InternalAPI<unknown>, [...tree, key]));
+      }
+      console.warn(`Unexpected data while wrapping API.`, "tree:", tree, "key:", key, "field:", field);
       throw new Error("Error while wrapping netscript API. See console.");
-    }
-  }
-  return eLayer;
+    },
+  };
+
+  // We target an empty Object, so that unproxied methods don't do anything.
+  // We *can't* freeze the target, because it would break invariants on ownKeys.
+  return new Proxy({}, handler) as ExternalAPI<API>;
 }
 
 /** Specify when a function was removed from the game, and its replacement function. */

src/NetscriptFunctions.ts

@@ -72,7 +72,7 @@ import { Flags } from "./NetscriptFunctions/Flags";
 import { calculateIntelligenceBonus } from "./PersonObjects/formulas/intelligence";
 import { CalculateShareMult, StartSharing } from "./NetworkShare/Share";
 import { recentScripts } from "./Netscript/RecentScripts";
-import { ExternalAPI, InternalAPI, removedFunction, StampedLayer, wrapAPILayer } from "./Netscript/APIWrapper";
+import { ExternalAPI, InternalAPI, removedFunction, NSProxy } from "./Netscript/APIWrapper";
 import { INetscriptExtra } from "./NetscriptFunctions/Extra";
 import { ScriptDeath } from "./Netscript/ScriptDeath";
 import { getBitNodeMultipliers } from "./BitNode/BitNode";
@@ -1893,34 +1893,8 @@ Object.assign(ns, {
   getServerRam: removedFunction("v2.2.0", "getServerMaxRam and getServerUsedRam"),
 });
 
-// add undocumented functions
-export const wrappedNS = wrapAPILayer({} as ExternalAPI<NSFull>, ns, []);
-
-// Figure out once which layers of ns have functions on them and will need to be stamped with a private workerscript field for API access
-const layerLocations: string[][] = [];
-function populateLayers(nsLayer: ExternalAPI<unknown>, currentLayers: string[] = []) {
-  for (const [k, v] of Object.entries(nsLayer)) {
-    if (typeof v === "object" && k !== "enums") {
-      if (Object.values(v as object).some((member) => typeof member === "function"))
-        layerLocations.push([...currentLayers, k]);
-      populateLayers(v as ExternalAPI<unknown>, [...currentLayers, k]);
-    }
-  }
-}
-populateLayers(wrappedNS);
-
 export function NetscriptFunctions(ws: WorkerScript): ExternalAPI<NSFull> {
-  //TODO unplanned: better typing instead of relying on an any
-  const instance = new StampedLayer(ws, wrappedNS) as any;
-  for (const layerLocation of layerLocations) {
-    const key = layerLocation.pop() as string;
-    const obj = layerLocation.reduce((prev, curr) => prev[curr], instance);
-    layerLocation.push(key);
-    obj[key] = new StampedLayer(ws, obj[key]);
-  }
-  instance.args = ws.args.slice();
-  instance.pid = ws.pid;
-  return instance;
+  return NSProxy(ws, ns, [], { args: ws.args.slice(), pid: ws.pid });
 }
 
 const possibleLogs = Object.fromEntries([...getFunctionNames(ns, "")].map((a) => [a, true]));

src/NetscriptWorker.ts

@@ -11,7 +11,7 @@ import { generateNextPid } from "./Netscript/Pid";
 
 import { CONSTANTS } from "./Constants";
 import { Interpreter } from "./ThirdParty/JSInterpreter";
-import { NetscriptFunctions, wrappedNS } from "./NetscriptFunctions";
+import { NetscriptFunctions } from "./NetscriptFunctions";
 import { compile, Node } from "./NetscriptJSEvaluator";
 import { IPort } from "./NetscriptPort";
 import { RunningScript } from "./Script/RunningScript";
@@ -81,15 +81,15 @@ async function startNetscript1Script(workerScript: WorkerScript): Promise<void>
 
   //TODO unplanned: Make NS1 wrapping type safe instead of using BasicObject.
   type BasicObject = Record<string, any>;
+  const wrappedNS = NetscriptFunctions(workerScript);
   function wrapNS1Layer(int: Interpreter, intLayer: unknown, nsLayer = wrappedNS as BasicObject) {
-    if (nsLayer === wrappedNS) int.setProperty(intLayer, "args", int.nativeToPseudo(workerScript.args));
     for (const [name, entry] of Object.entries(nsLayer)) {
       if (typeof entry === "function") {
         const wrapper = async (...args: unknown[]) => {
           try {
             // Sent a resolver function as an extra arg. See createAsyncFunction JSInterpreter.js:3209
             const callback = args.pop() as (value: unknown) => void;
-            const result = await entry.bind(workerScript.env.vars)(...args.map((arg) => int.pseudoToNative(arg)));
+            const result = await entry(...args.map((arg) => int.pseudoToNative(arg)));
             return callback(int.nativeToPseudo(result));
           } catch (e: unknown) {
             errorToThrow = e;

test/jest/Netscript/RamCalculation.test.ts

@@ -1,5 +1,5 @@
 import { Player } from "../../../src/Player";
-import { NetscriptFunctions, wrappedNS } from "../../../src/NetscriptFunctions";
+import { NetscriptFunctions } from "../../../src/NetscriptFunctions";
 import { RamCosts, getRamCost, RamCostConstants } from "../../../src/Netscript/RamCostGenerator";
 import { Environment } from "../../../src/Netscript/Environment";
 import { RunningScript } from "../../../src/Script/RunningScript";
@@ -108,7 +108,7 @@ describe("Netscript RAM Calculation/Generation Tests", function () {
   }
 
   describe("ns", () => {
-    Object.entries(wrappedNS as unknown as NSLayer).forEach(([key, val]) => {
+    Object.entries(NetscriptFunctions(workerScript) as unknown as NSLayer).forEach(([key, val]) => {
       if (key === "args" || key === "enums") return;
       if (typeof val === "function") {
         // Removed functions have no ram cost and should be skipped.
@@ -145,7 +145,9 @@ describe("Netscript RAM Calculation/Generation Tests", function () {
 
   describe("Singularity multiplier checks", () => {
     sf4.lvl = 3;
-    const singFunctions = Object.entries(wrappedNS.singularity).filter(([__, val]) => typeof val === "function");
+    const singFunctions = Object.entries(NetscriptFunctions(workerScript).singularity).filter(
+      ([__, val]) => typeof val === "function",
+    );
     const singObjects = singFunctions
       .filter((fn) => !isRemovedFunction(fn))
       .map(([key, val]) => {