NETSCRIPT: Faster API wrapping on script launch. (#229)

* ns API is wrapped once
* when a new workerscript is created, each layer of ns is stamped with a private workerscript field that allows the functions to work.
* Test has been refactored to account for new method of wrapping
* BREAKING: ns functions need access to `this` value of their parent ns layer (or any ns layer)
* Enums are passed directly to player (no cloning) but are frozen.

src/Netscript/APIWrapper.ts

@@ -2,7 +2,6 @@ import { getRamCost } from "./RamCostGenerator";
 import type { WorkerScript } from "./WorkerScript";
 import { helpers } from "./NetscriptHelpers";
 import { ScriptArg } from "./ScriptArg";
-import { NSFull } from "src/NetscriptFunctions";
 import { cloneDeep } from "lodash";
 
 /** Generic type for an enums object */
@@ -37,42 +36,57 @@ export type InternalAPI<API> = {
 /** 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>, func: InternalFn<APIFn>, tree: string[], key: Key<API>) {
+    const arrayPath = [...tree, key];
+    const functionPath = arrayPath.join(".");
+    function wrappedFunction(this: StampedLayer, ...args: unknown[]): unknown {
+      const ctx = { workerScript: this.#workerScript, function: key, functionPath };
+      helpers.checkEnvFlags(ctx);
+      helpers.updateDynamicRam(ctx, getRamCost(...tree, key));
+      return func(ctx)(...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;
   function: string;
   functionPath: string;
 };
 
-export function wrapAPI(ws: WorkerScript, internalAPI: InternalAPI<NSFull>, args: ScriptArg[]): ExternalAPI<NSFull> {
+export function wrapAPILayer<API>(
-  function wrapAPILayer<API>(eLayer: ExternalAPI<API>, iLayer: InternalAPI<API>, tree: string[]): ExternalAPI<API> {
+  eLayer: ExternalAPI<API>,
-    for (const [key, value] of Object.entries(iLayer) as [Key<API>, InternalValues][]) {
+  iLayer: InternalAPI<API>,
-      if (key === "enums") {
+  tree: string[],
-        (eLayer[key] as Enums) = cloneDeep(value as Enums);
+): ExternalAPI<API> {
-      } else if (key === "args") continue;
+  for (const [key, value] of Object.entries(iLayer) as [Key<API>, InternalValues][]) {
-      // Args are added in wrapAPI function and should only exist at top level
+    if (key === "enums") {
-      else if (typeof value === "function") {
+      const enumObj = Object.freeze(cloneDeep(value as Enums));
-        wrapFunction(eLayer, value as InternalFn<APIFn>, tree, key);
+      for (const member of Object.values(enumObj)) Object.freeze(member);
-      } else if (typeof value === "object") {
+      (eLayer[key] as Enums) = enumObj;
-        wrapAPILayer((eLayer[key] = {} as ExternalAPI<API>[Key<API>]), value, [...tree, key as string]);
+    } else if (key === "args") continue;
-      } else {
+    // Args only added on individual instances.
-        console.warn(`Unexpected data while wrapping API.`, "tree:", tree, "key:", key, "value:", value);
+    else if (typeof value === "function") {
-        throw new Error("Error while wrapping netscript API. See console.");
+      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);
+      throw new Error("Error while wrapping netscript API. See console.");
     }
-    return eLayer;
   }
-  function wrapFunction<API>(eLayer: ExternalAPI<API>, func: InternalFn<APIFn>, tree: string[], key: Key<API>) {
+  return eLayer;
-    const arrayPath = [...tree, key];
-    const functionPath = arrayPath.join(".");
-    const ctx = { workerScript: ws, function: key, functionPath };
-    function wrappedFunction(...args: unknown[]): unknown {
-      helpers.checkEnvFlags(ctx);
-      helpers.updateDynamicRam(ctx, getRamCost(...tree, key));
-      return func(ctx)(...args);
-    }
-    (eLayer[key] as WrappedFn) = wrappedFunction;
-  }
-
-  const wrappedAPI = wrapAPILayer({ args } as ExternalAPI<NSFull>, internalAPI, []);
-  return wrappedAPI;
 }

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, wrapAPI } from "./Netscript/APIWrapper";
+import { ExternalAPI, InternalAPI, StampedLayer, wrapAPILayer } from "./Netscript/APIWrapper";
 import { INetscriptExtra } from "./NetscriptFunctions/Extra";
 import { ScriptDeath } from "./Netscript/ScriptDeath";
 import { getBitNodeMultipliers } from "./BitNode/BitNode";
@@ -93,10 +93,6 @@ export const enums: NSEnums = {
 
 export type NSFull = Readonly<NS & INetscriptExtra>;
 
-export function NetscriptFunctions(workerScript: WorkerScript): ExternalAPI<NSFull> {
-  return wrapAPI(workerScript, ns, workerScript.args.slice());
-}
-
 const base: InternalAPI<NS> = {
   args: [],
   enums,
@@ -1913,6 +1909,34 @@ export const ns = {
   ...NetscriptExtra(),
 };
 
+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: 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();
+  return instance;
+}
+
 const possibleLogs = Object.fromEntries([...getFunctionNames(ns, "")].map((a) => [a, true]));
 /** Provides an array of all function names on a nested object */
 function getFunctionNames(obj: object, prefix: string): string[] {

src/NetscriptWorker.ts

@@ -11,7 +11,7 @@ import { generateNextPid } from "./Netscript/Pid";
 
 import { CONSTANTS } from "./Constants";
 import { Interpreter } from "./ThirdParty/JSInterpreter";
-import { NetscriptFunctions } from "./NetscriptFunctions";
+import { NetscriptFunctions, wrappedNS } from "./NetscriptFunctions";
 import { compile, Node } from "./NetscriptJSEvaluator";
 import { IPort } from "./NetscriptPort";
 import { RunningScript } from "./Script/RunningScript";
@@ -81,14 +81,14 @@ async function startNetscript1Script(workerScript: WorkerScript): Promise<void>
 
   //TODO: Make NS1 wrapping type safe instead of using BasicObject
   type BasicObject = Record<string, any>;
-  function wrapNS1Layer(int: Interpreter, intLayer: unknown, nsLayer = workerScript.env.vars as BasicObject) {
+  function wrapNS1Layer(int: Interpreter, intLayer: unknown, nsLayer = wrappedNS as BasicObject) {
     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(...args.map((arg) => int.pseudoToNative(arg)));
+            const result = await entry.bind(workerScript.env.vars)(...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 } from "../../../src/NetscriptFunctions";
+import { NetscriptFunctions, wrappedNS } from "../../../src/NetscriptFunctions";
 import { RamCosts, getRamCost, RamCostConstants } from "../../../src/Netscript/RamCostGenerator";
 import { Environment } from "../../../src/Netscript/Environment";
 import { RunningScript } from "../../../src/Script/RunningScript";
@@ -99,11 +99,11 @@ describe("Netscript RAM Calculation/Generation Tests", function () {
   }
 
   describe("ns", () => {
-    Object.entries(ns as unknown as NSLayer).forEach(([key, val]) => {
+    Object.entries(wrappedNS as unknown as NSLayer).forEach(([key, val]) => {
       if (key === "args" || key === "enums") return;
       if (typeof val === "function") {
         const expectedRam = grabCost(RamCosts, [key]);
-        it(`${key}()`, () => combinedRamCheck(val, [key], expectedRam));
+        it(`${key}()`, () => combinedRamCheck(val.bind(ns), [key], expectedRam));
       }
       //The only other option should be an NSLayer
       const extraLayerCost = { corporation: 1022.4, hacknet: 4 }[key] ?? 0;
@@ -112,13 +112,15 @@ describe("Netscript RAM Calculation/Generation Tests", function () {
   });
 
   function testLayer(nsLayer: NSLayer, ramLayer: RamLayer, path: string[], extraLayerCost: number) {
+    // nsLayer is the layer on the main, unstamped wrappedNS object. The actualLayer is needed to check correct stamping.
+    const actualLayer = path.reduce((prev, curr) => prev[curr], ns as any); //todo: do this typesafely?
     describe(path[path.length - 1], () => {
       Object.entries(nsLayer).forEach(([key, val]) => {
         const newPath = [...path, key];
         if (typeof val === "function") {
           const fnName = newPath.join(".");
           const expectedRam = grabCost(ramLayer, newPath);
-          it(`${fnName}()`, () => combinedRamCheck(val, newPath, expectedRam, extraLayerCost));
+          it(`${fnName}()`, () => combinedRamCheck(val.bind(actualLayer), newPath, expectedRam, extraLayerCost));
         }
         //Skip enums layers
         else if (key === "enums") return;
@@ -130,11 +132,11 @@ describe("Netscript RAM Calculation/Generation Tests", function () {
 
   describe("Singularity multiplier checks", () => {
     sf4.lvl = 3;
-    const singFunctions = Object.entries(ns.singularity).filter(([__, val]) => typeof val === "function");
+    const singFunctions = Object.entries(wrappedNS.singularity).filter(([__, val]) => typeof val === "function");
     const singObjects = singFunctions.map(([key, val]) => {
       return {
         name: key,
-        fn: val,
+        fn: val.bind(ns.singularity),
         baseRam: grabCost(RamCosts.singularity, ["singularity", key]),
       };
     });