poe-code 13.0.2 → 13.0.4

package/packages/safe-js/dist/snapshot/restore.js

@@ -1,10 +1,11 @@
-import { Budget } from "../interp/budget.js";
+import { Budget, SandboxError } from "../interp/budget.js";
+import { CompileScope } from "../interp/regex/compile-guard.js";
 import { decodeFloat32Storage } from "./float32array.js";
 import { sandboxErrorTypes } from "../error/shape.js";
 import { SnapshotMismatchError } from "../restore.js";
 import { Scope } from "../interp/interpreter.js";
 import { wrapCallerInjectedBindings } from "../interp/host-bridge.js";
-import { createSandboxArguments, createSandboxClosure, createSandboxGenerator, createSandboxMap, createSandboxPromise, createSandboxRegex, createSandboxSet, measureSandboxData } from "../interp/values.js";
+import { createSandboxArguments, createSandboxClosure, createSandboxGenerator, createSandboxMap, createSandboxPromise, createSandboxRegex, createSandboxSet, reconcileCompiledValues } from "../interp/values.js";
 import { createGeneratorChannel, restoreGeneratorChannel } from "../interp/generator.js";
 import { hashSource } from "../parse/hash.js";
 import { getFunctionLength } from "../parse/bindings.js";
@@ -14,70 +15,84 @@ import { interpret } from "../interp/interpreter.js";
 import { bindPattern } from "../interp/patterns.js";
 import { resolvePendingHostCallResumePolicy } from "./policy.js";
 import { validateInterpreterSnapshot, validateSnapshotSourceHash } from "./validation.js";
-export function restore(snapshot, options) {
+export function restore(snapshot, options, owner) {
     const budget = options.budget ?? new Budget();
-    validateSnapshotSourceHash(snapshot);
-    let currentSourceHash;
+    const operation = budget.acquireCompileOwner(false, owner);
+    const compilation = new CompileScope(operation.owner);
     try {
-        currentSourceHash = hashSource(options.source);
-    }
-    catch {
-        throw new Error(`source changed since snapshot was taken (hash ${snapshot.sourceHash} expected, but current source could not be hashed); pass --reset to discard`);
-    }
-    if (snapshot.sourceHash !== currentSourceHash) {
-        throw new SnapshotMismatchError(snapshot.sourceHash, currentSourceHash);
-    }
-    const ast = parseModule(options.source);
-    const nodeById = indexAstNodes(ast);
-    validateInterpreterSnapshot(snapshot, nodeById, budget);
-    const currentNode = nodeById.get(snapshot.currentAstNodeId);
-    if (currentNode === undefined) {
-        throw new Error(`Snapshot references unknown AST node ${snapshot.currentAstNodeId}.`);
-    }
-    const state = {
-        budget,
-        heap: snapshot.heap ?? {},
-        heapValueById: new Map(),
-        moduleBindings: restoreModuleBindings(snapshot.moduleBindings, options.modules, {
+        validateSnapshotSourceHash(snapshot);
+        let currentSourceHash;
+        try {
+            currentSourceHash = hashSource(options.source, operation.owner);
+        }
+        catch (error) {
+            if (error instanceof SandboxError)
+                throw error;
+            throw new Error(`source changed since snapshot was taken (hash ${snapshot.sourceHash} expected, but current source could not be hashed); pass --reset to discard`);
+        }
+        if (snapshot.sourceHash !== currentSourceHash) {
+            throw new SnapshotMismatchError(snapshot.sourceHash, currentSourceHash);
+        }
+        const ast = parseModule(options.source, "<input>", operation.owner);
+        const nodeById = indexAstNodes(ast);
+        validateInterpreterSnapshot(snapshot, nodeById, budget);
+        const currentNode = nodeById.get(snapshot.currentAstNodeId);
+        if (currentNode === undefined) {
+            throw new Error(`Snapshot references unknown AST node ${snapshot.currentAstNodeId}.`);
+        }
+        const state = {
             budget,
-            signal: options.signal
-        }),
-        nodeById,
-        pendingPromiseById: new Map(),
-        serializedScopeById: new Map(snapshot.scopeChain.map((frame) => [frame.id, frame])),
-        scopeById: new Map(),
-        scopeFrameById: new Map()
-    };
-    const pendingPromises = snapshot.pendingPromises.map((entry) => restorePendingPromise(entry, state));
-    const scopeChain = snapshot.scopeChain.map((frame) => restoreScopeFrame(frame, state));
-    const currentScopeId = snapshot.callStack.at(-1)?.scopeId ?? scopeChain.at(-1)?.id;
-    if (currentScopeId === undefined) {
-        throw new Error("Snapshot does not contain a scope to resume.");
-    }
-    const currentScope = state.scopeById.get(currentScopeId);
-    if (currentScope === undefined) {
-        throw new Error(`Snapshot references unknown scope ${String(currentScopeId)}.`);
-    }
-    budget.reconcileDataUsage(measureSandboxData([
-        ...currentScope.retainedValues(),
-        ...pendingPromises.flatMap((pending) => Object.values(pending).filter(isSandboxSnapshotValue))
-    ]));
-    return {
-        ast,
-        budget,
-        callStack: snapshot.callStack.map((frame) => restoreCallFrame(frame, state)),
-        currentAstNodeId: snapshot.currentAstNodeId,
-        currentNode,
-        currentScope,
-        moduleBindings: state.moduleBindings,
-        pendingPromises,
-        scopeChain: scopeChain.map((frame) => ({
-            ...frame,
-            scope: state.scopeById.get(frame.id) ?? currentScope
-        })),
-        signal: options.signal,
-        sourceHash: snapshot.sourceHash
-    };
+            compilation,
+            heap: snapshot.heap ?? {},
+            heapValueById: new Map(),
+            moduleBindings: restoreModuleBindings(snapshot.moduleBindings, options.modules, {
+                budget,
+                compileOwner: operation.owner,
+                signal: options.signal
+            }),
+            nodeById,
+            pendingPromiseById: new Map(),
+            serializedScopeById: new Map(snapshot.scopeChain.map((frame) => [frame.id, frame])),
+            scopeById: new Map(),
+            scopeFrameById: new Map()
+        };
+        const pendingPromises = snapshot.pendingPromises.map((entry) => restorePendingPromise(entry, state));
+        const scopeChain = snapshot.scopeChain.map((frame) => restoreScopeFrame(frame, state));
+        const currentScopeId = snapshot.callStack.at(-1)?.scopeId ?? scopeChain.at(-1)?.id;
+        if (currentScopeId === undefined) {
+            throw new Error("Snapshot does not contain a scope to resume.");
+        }
+        const currentScope = state.scopeById.get(currentScopeId);
+        if (currentScope === undefined) {
+            throw new Error(`Snapshot references unknown scope ${String(currentScopeId)}.`);
+        }
+        const callStack = snapshot.callStack.map((frame) => restoreCallFrame(frame, state));
+        reconcileCompiledValues(budget, [
+            ...currentScope.retainedValues(),
+            ...budget.retainedValues(),
+            ...pendingPromises.flatMap((pending) => Object.values(pending).filter(isSandboxSnapshotValue))
+        ], compilation);
+        return {
+            ast,
+            budget,
+            callStack,
+            currentAstNodeId: snapshot.currentAstNodeId,
+            currentNode,
+            currentScope,
+            moduleBindings: state.moduleBindings,
+            pendingPromises,
+            scopeChain: scopeChain.map((frame) => ({
+                ...frame,
+                scope: state.scopeById.get(frame.id) ?? currentScope
+            })),
+            signal: options.signal,
+            sourceHash: snapshot.sourceHash
+        };
+    }
+    finally {
+        compilation.dispose();
+        operation.release();
+    }
 }
 function isSandboxSnapshotValue(value) {
     return value !== undefined && !(value instanceof Promise);
@@ -229,7 +244,7 @@ function deserializeValue(value, state) {
         return restoreGeneratorValue(value, state);
     }
     if (isSerializedRegexValue(value)) {
-        return createSandboxRegex(value.source, value.flags, value.lastIndex);
+        return createSandboxRegex(value.source, value.flags, value.lastIndex, state.compilation);
     }
     const object = Object.create(null);
     for (const [key, entry] of Object.entries(value)) {
@@ -255,6 +270,7 @@ function restoreGeneratorValue(value, state) {
             restoreParentScope(value.capturedScopeId, state);
         const result = await interpret(node.body, {
             budget: state.budget,
+            compileOwner: state.compilation.owner,
             ...(value.state === "suspended"
                 ? {
                     generatorResume: {
@@ -438,28 +454,60 @@ function restoreClosureValue(astNodeId, capturedScopeId, state) {
     return restoredClosure;
 }
 async function executeRestoredClosure(node, capturedScopeId, closure, args, thisValue, state) {
-    const capturedScope = state.scopeById.get(capturedScopeId) ?? restoreParentScope(capturedScopeId, state);
-    const wrapperScope = node.type === "FunctionExpression" && node.id !== undefined
-        ? capturedScope.child()
-        : capturedScope;
-    const scope = wrapperScope.child();
-    if (node.type === "FunctionExpression" && node.id !== undefined) {
-        wrapperScope.declare(node.id.name, "const", closure);
-    }
-    if (node.type !== "ArrowFunctionExpression") {
-        scope.declare("this", "const", thisValue);
-        state.budget.allocateArrayLength(args.length);
-        scope.declare("arguments", "let", createSandboxArguments(args));
-    }
-    for (let index = 0; index < node.params.length; index += 1) {
-        const param = node.params[index];
-        if (param.type === "RestElement") {
-            const rest = args.slice(index);
-            state.budget.allocateArrayLength(rest.length);
-            const binding = await bindPattern(param, rest, { kind: "let" }, scope, {
+    const operation = state.budget.acquireCompileOwner(false, state.compilation.owner);
+    const compilation = new CompileScope(operation.owner);
+    state = { ...state, compilation };
+    try {
+        const capturedScope = state.scopeById.get(capturedScopeId) ?? restoreParentScope(capturedScopeId, state);
+        const wrapperScope = node.type === "FunctionExpression" && node.id !== undefined
+            ? capturedScope.child()
+            : capturedScope;
+        const scope = wrapperScope.child();
+        if (node.type === "FunctionExpression" && node.id !== undefined) {
+            wrapperScope.declare(node.id.name, "const", closure);
+        }
+        if (node.type !== "ArrowFunctionExpression") {
+            scope.declare("this", "const", thisValue);
+            state.budget.allocateArrayLength(args.length);
+            scope.declare("arguments", "let", createSandboxArguments(args));
+        }
+        for (let index = 0; index < node.params.length; index += 1) {
+            const param = node.params[index];
+            if (param.type === "RestElement") {
+                const rest = args.slice(index);
+                state.budget.allocateArrayLength(rest.length);
+                const binding = await bindPattern(param, rest, { kind: "let" }, scope, {
+                    evaluate: async (defaultNode) => {
+                        const result = await interpret(defaultNode, {
+                            budget: state.budget,
+                            compilation,
+                            scope,
+                            useScopeDirectly: true
+                        });
+                        return result.ok
+                            ? {
+                                kind: "normal",
+                                hasValue: "returnValue" in result,
+                                value: result.returnValue
+                            }
+                            : { kind: "error", error: result.error };
+                    }
+                });
+                if (!binding.ok) {
+                    if (binding.result.kind === "error") {
+                        throw binding.result.error;
+                    }
+                    if (binding.result.kind === "throw") {
+                        throw binding.result.value;
+                    }
+                }
+                break;
+            }
+            const binding = await bindPattern(param, args[index], { kind: "let" }, scope, {
                 evaluate: async (defaultNode) => {
                     const result = await interpret(defaultNode, {
                         budget: state.budget,
+                        compilation,
                         scope,
                         useScopeDirectly: true
                     });
@@ -480,41 +528,22 @@ async function executeRestoredClosure(node, capturedScopeId, closure, args, this
                     throw binding.result.value;
                 }
             }
-            break;
         }
-        const binding = await bindPattern(param, args[index], { kind: "let" }, scope, {
-            evaluate: async (defaultNode) => {
-                const result = await interpret(defaultNode, {
-                    budget: state.budget,
-                    scope,
-                    useScopeDirectly: true
-                });
-                return result.ok
-                    ? {
-                        kind: "normal",
-                        hasValue: "returnValue" in result,
-                        value: result.returnValue
-                    }
-                    : { kind: "error", error: result.error };
-            }
+        const result = await interpret(node.body, {
+            budget: state.budget,
+            compilation,
+            scope
         });
-        if (!binding.ok) {
-            if (binding.result.kind === "error") {
-                throw binding.result.error;
-            }
-            if (binding.result.kind === "throw") {
-                throw binding.result.value;
-            }
+        if (!result.ok) {
+            throw new Error(result.error.message);
         }
+        reconcileCompiledValues(state.budget, [...scope.retainedValues(), ...state.budget.retainedValues(), result.returnValue], compilation);
+        return result.returnValue;
     }
-    const result = await interpret(node.body, {
-        budget: state.budget,
-        scope
-    });
-    if (!result.ok) {
-        throw new Error(result.error.message);
+    finally {
+        compilation.dispose();
+        operation.release();
     }
-    return result.returnValue;
 }
 function restorePromiseValue(id, state) {
     return (state.pendingPromiseById.get(id) ??