starboard-python
Version:
Python cells for Starboard Notebook
643 lines (595 loc) • 21.5 kB
text/typescript
import { nanoid } from "nanoid";
import { AsyncMemory } from "./async-memory";
const SERIALIZATION = {
UNDEFINED: 0,
NULL: 1,
FALSE: 2,
TRUE: 3,
NUMBER: 4,
DATE: 5,
KNOWN_SYMBOL: 6,
STRING: 10,
BIGINT: 11,
OBJECT: 255,
} as const;
// https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Symbol
const KNOWN_SYMBOLS = [
Symbol.asyncIterator,
Symbol.hasInstance,
Symbol.isConcatSpreadable,
Symbol.iterator,
Symbol.match,
Symbol.matchAll,
Symbol.replace,
Symbol.search,
Symbol.species,
Symbol.split,
Symbol.toPrimitive,
Symbol.toStringTag,
Symbol.unscopables,
];
const textEncoder = new TextEncoder();
const textDecoder = new TextDecoder("utf-8");
const encodeFloat = useFloatEncoder();
const decodeFloat = useFloatDecoder();
function useFloatEncoder() {
// https://stackoverflow.com/a/14379836/3492994
const temp = new ArrayBuffer(8);
const tempFloat64 = new Float64Array(temp);
const tempUint8 = new Uint8Array(temp);
return function (value: number): Uint8Array {
tempFloat64[0] = value;
return tempUint8;
};
}
function useFloatDecoder() {
const temp = new ArrayBuffer(8);
const tempFloat64 = new Float64Array(temp);
const tempUint8 = new Uint8Array(temp);
return function (value: Uint8Array): number {
tempUint8.set(value);
return tempFloat64[0];
};
}
/**
* Lets one other thread access the objects on this thread.
* Usually runs on the main thread.
*/
export class ObjectProxyHost {
readonly rootReferences = new Map<string, any>();
readonly temporaryReferences = new Map<string, any>();
readonly memory: AsyncMemory;
private writeMemoryContinuation?: () => void;
constructor(memory: AsyncMemory) {
this.memory = memory;
}
/** Creates a valid, random id for a given object */
private getId(value: any) {
return nanoid() + "-" + (typeof value === "function" ? "f" : "o");
}
registerRootObject(value: any) {
const id = this.getId(value);
this.rootReferences.set(id, value);
return id;
}
registerTempObject(value: any) {
const id = this.getId(value);
this.temporaryReferences.set(id, value);
return id;
}
clearTemporary() {
this.temporaryReferences.clear();
}
getObject(id: string) {
return this.rootReferences.get(id) ?? this.temporaryReferences.get(id);
}
// A serializePostMessage isn't needed here, because all we're ever going to pass to the worker are ids
serializeMemory(value: any, memory: AsyncMemory) {
// Format:
// [1 byte][n bytes ]
// [type ][data ]
memory.writeSize(8); // Anything that fits into 8 bytes is fine
// Simple primitives. Guaranteed to fit into the shared memory.
if (value === undefined) {
memory.memory[0] = SERIALIZATION.UNDEFINED;
memory.unlockSize();
} else if (value === null) {
memory.memory[0] = SERIALIZATION.NULL;
memory.unlockSize();
} else if (value === false) {
memory.memory[0] = SERIALIZATION.FALSE;
memory.unlockSize();
} else if (value === true) {
memory.memory[0] = SERIALIZATION.TRUE;
memory.unlockSize();
} else if (typeof value === "number") {
memory.memory[0] = SERIALIZATION.NUMBER;
memory.memory.set(encodeFloat(value), 1);
memory.unlockSize();
} else if (value instanceof Date) {
memory.memory[0] = SERIALIZATION.DATE;
const time = value.getTime();
memory.memory.set(encodeFloat(time), 1);
memory.unlockSize();
} else if (typeof value === "symbol" && KNOWN_SYMBOLS.includes(value)) {
memory.memory[0] = SERIALIZATION.KNOWN_SYMBOL;
memory.memory[1] = KNOWN_SYMBOLS.indexOf(value);
memory.unlockSize();
}
// Variable length primitives. Not guaranteed to fit into the shared memory, but we know their size.
else if (typeof value === "string") {
memory.memory[0] = SERIALIZATION.STRING;
// A string encoded in utf-8 uses at most 4 bytes per character
if (value.length * 4 <= memory.memory.byteLength) {
const data = textEncoder.encode(value);
memory.memory.set(data, 1);
memory.writeSize(data.byteLength);
memory.unlockSize();
} else {
// Longer strings need to be sent piece by piece
const bytes = textEncoder.encode(value);
const memorySize = memory.memory.byteLength;
let offset = 0;
let remainingBytes = bytes.byteLength;
memory.memory.set(bytes.subarray(offset, memorySize - 1), 1);
offset += memorySize - 1;
remainingBytes -= memorySize - 1;
this.writeMemoryContinuation = () => {
if (remainingBytes > 0) {
memory.memory.set(bytes.subarray(offset, memorySize), 0);
offset += memorySize;
remainingBytes -= memorySize;
} else {
this.writeMemoryContinuation = undefined;
}
memory.unlockSize();
};
memory.writeSize(bytes.byteLength);
memory.unlockSize();
}
} else if (typeof value === "bigint") {
memory.memory[0] = SERIALIZATION.BIGINT;
const digits = value.toString();
// TODO: Implement this (just like the text ^)
console.warn("Bigint support is not implemented");
memory.unlockSize();
}
// Object. Serialized as ID, guaranteed to fit into shared memory
else {
memory.memory[0] = SERIALIZATION.OBJECT;
const id = this.registerTempObject(value);
const data = textEncoder.encode(id);
memory.memory.set(data, 1);
memory.writeSize(data.byteLength);
memory.unlockSize();
}
}
/**
* Deserializes an object that was sent through postMessage
*/
deserializePostMessage(value: any): any {
if (typeof value === "object" && value !== null) {
// Special cases
if (value.id) return this.getObject(value.id);
if (value.value) return value.value;
if (value.symbol) return KNOWN_SYMBOLS[value.symbol];
}
// It's a primitive
return value;
}
handleProxyMessage(message: ProxyMessage, memory: AsyncMemory) {
if (message.type === "proxy_reflect") {
try {
if (message.method === "apply") {
const method = Reflect[message.method];
const args = (message.args ?? []).map((v) => this.deserializePostMessage(v));
const result = (method as any)(this.getObject(message.target), this.getObject(message.thisArg), args);
// Write result to shared memory
this.serializeMemory(result, memory);
} else {
const method = Reflect[message.method];
const args = (message.args ?? []).map((v) => this.deserializePostMessage(v));
const result = (method as any)(this.getObject(message.target), ...args);
// Write result to shared memory
this.serializeMemory(result, memory);
}
} catch (e) {
console.error(message);
throw e;
}
} else if (message.type === "proxy_shared_memory") {
// Write remaining data to shared memory
if (this.writeMemoryContinuation === undefined) {
console.warn("No more data to write to shared memory");
} else {
this.writeMemoryContinuation();
}
} else if (message.type === "proxy_print_object") {
console.log("Object with id", message.target, "is", this.getObject(message.target));
} else if (message.type === "proxy_promise") {
const promiseObject: Promise<any> = this.getObject(message.target);
if (message.method === "then") {
promiseObject[message.method](
(value) => {
const result = { value: value };
this.serializeMemory(result, memory);
},
(err) => {
const result = { error: err };
this.serializeMemory(result, memory);
}
);
} else {
console.error("Unknown proxy promise method", message);
}
} else {
console.warn("Unknown proxy message", message);
}
}
}
export const ObjectId = Symbol.for("id");
/**
* Allows this thread to access objects from another thread.
* Must run on a worker thread.
*/
export class ObjectProxyClient {
readonly memory: AsyncMemory;
readonly postMessage: (message: ProxyMessage) => void;
constructor(memory: AsyncMemory, postMessage: (message: ProxyMessage) => void) {
this.memory = memory;
this.postMessage = postMessage;
}
/**
* Serializes an object so that it can be sent using postMessage
*/
serializePostMessage(value: any): any {
if (isSimplePrimitive(value)) {
return value;
} else if (isSymbolPrimitive(value)) {
return { symbol: KNOWN_SYMBOLS.indexOf(value) };
} else if (isVariableLengthPrimitive(value)) {
return value;
} else if (value[ObjectId] !== undefined) {
return { id: value[ObjectId] };
} else {
// Maybe serialize simple functions https://stackoverflow.com/questions/1833588/javascript-clone-a-function
return { value: value }; // Might fail to get serialized
}
}
/**
* Deserializes an object from a shared array buffer. Can return a proxy.
*/
deserializeMemory(memory: AsyncMemory) {
const numberOfBytes = memory.readSize();
// Uint8Arrays have the convenient property of having 1 byte per element.
let resultBytes: Uint8Array;
if (numberOfBytes <= memory.sharedMemory.byteLength) {
resultBytes = memory.memory;
} else {
const memorySize = memory.sharedMemory.byteLength;
let offset = 0;
let remainingBytes = numberOfBytes;
resultBytes = new Uint8Array(numberOfBytes);
while (remainingBytes >= memorySize) {
resultBytes.set(memory.memory, offset);
offset += memorySize;
remainingBytes -= memorySize;
memory.lockSize();
this.postMessage({ type: "proxy_shared_memory" });
memory.waitForSize();
}
if (remainingBytes > 0) {
resultBytes.set(memory.memory.subarray(0, remainingBytes), offset);
}
}
// Simple primitives. Guaranteed to fit into the shared memory.
if (resultBytes[0] === SERIALIZATION.UNDEFINED) {
return undefined;
} else if (resultBytes[0] === SERIALIZATION.NULL) {
return null;
} else if (resultBytes[0] === SERIALIZATION.FALSE) {
return false;
} else if (resultBytes[0] === SERIALIZATION.TRUE) {
return true;
} else if (resultBytes[0] === SERIALIZATION.NUMBER) {
return decodeFloat(resultBytes.subarray(1, 9));
} else if (resultBytes[0] === SERIALIZATION.DATE) {
const date = new Date();
date.setTime(decodeFloat(resultBytes.subarray(1, 9)));
return date;
} else if (resultBytes[0] === SERIALIZATION.KNOWN_SYMBOL) {
const symbol = KNOWN_SYMBOLS[resultBytes[1]];
return symbol;
}
// Variable length primitives. We already read all of their data
else if (resultBytes[0] === SERIALIZATION.STRING) {
// Note: This *copies* the entire thing, because you aren't allowed to call decode on a SharedArrayBuffer
return textDecoder.decode(resultBytes.slice(1, numberOfBytes + 1));
} else if (resultBytes[0] === SERIALIZATION.BIGINT) {
return BigInt(textDecoder.decode(resultBytes.slice(1, numberOfBytes + 1)));
}
// Object. Serialized as ID, guaranteed to fit into shared memory
else if (resultBytes[0] === SERIALIZATION.OBJECT) {
const id = textDecoder.decode(resultBytes.slice(1, numberOfBytes + 1));
return this.getObjectProxy(id);
} else {
console.warn("Unknown type", resultBytes[0]);
return null;
}
}
/**
* Calls a Reflect function on an object from the other thread
* @returns The result of the operation, can be a primitive or a proxy
*/
private proxyReflect(method: keyof typeof Reflect, targetId: string, args: any[]) {
let value: any = undefined;
try {
this.memory.lockWorker();
this.memory.lockSize();
this.memory.writeSize(0);
if (method === "apply") {
// Special case for "apply"
this.postMessage({
type: "proxy_reflect",
method: method,
target: targetId,
thisArg: args[0],
args: args[1].map((v: any[]) => this.serializePostMessage(v)),
});
} else {
this.postMessage({
type: "proxy_reflect",
method: method,
target: targetId,
args: args.map((v) => this.serializePostMessage(v)),
});
}
this.memory.waitForSize();
value = this.deserializeMemory(this.memory);
} catch (e) {
console.error({ method, targetId, args });
console.error(e);
this.postMessage({
type: "proxy_print_object",
target: targetId,
});
} finally {
// Regardless of what happened, unlock the size and the worker
this.memory.forceUnlockSize();
this.memory.unlockWorker();
}
return value;
}
private proxyPromise(method: "then", targetId: string): { value?: any; error?: any } {
let value: any = undefined;
try {
this.memory.lockWorker();
this.memory.lockSize();
this.memory.writeSize(0);
this.postMessage({
type: "proxy_promise",
method: method,
target: targetId,
});
this.memory.waitForSize();
value = this.deserializeMemory(this.memory);
} catch (e) {
console.error({ method, targetId });
console.error(e);
this.postMessage({
type: "proxy_print_object",
target: targetId,
});
} finally {
// Regardless of what happened, unlock the size and the worker
this.memory.forceUnlockSize();
this.memory.unlockWorker();
}
return value;
}
/** Checks if an id encodes a function. Mostly a silly hack to ensure that proxies can work as expected */
private isFunction(id: string) {
return id.endsWith("-f");
}
/**
* Gets a proxy object for a given id
*/
getObjectProxy<T = any>(id: string): T {
const client = this;
return new Proxy(this.isFunction(id) ? function () {} : {}, {
get(target, prop, receiver) {
if (prop === ObjectId) {
return id;
}
// const value = Reflect.get(target, prop, receiver);
const value = client.proxyReflect("get", id, [prop, receiver]);
if (typeof value !== "function") return value;
// TODO: Special handling for .bind, .apply and more
/* Functions need special handling
* https://stackoverflow.com/questions/27983023/proxy-on-dom-element-gives-error-when-returning-functions-that-implement-interfa
* https://stackoverflow.com/questions/37092179/javascript-proxy-objects-dont-work
*/
return new Proxy(value, {
apply(_, thisArg, argumentsList) {
// thisArg: the object the function was called with. Can be the proxy or something else
// receiver: the object the propery was gotten from. Is always the proxy or something inheriting from the proxy
// target: the original object
const calledWithProxy = thisArg === receiver;
// return Reflect.apply(value, calledWithProxy ? target : thisArg, args);
const functionReturnValue = client.proxyReflect("apply", value[ObjectId], [
calledWithProxy ? id : thisArg[ObjectId],
argumentsList,
]);
return functionReturnValue;
},
});
},
set(target, prop, value, receiver) {
// return Reflect.set(target, prop, value, receiver);
return client.proxyReflect("set", id, [prop, value, receiver]);
},
ownKeys(target) {
// return Reflect.ownKeys(target);
return client.proxyReflect("ownKeys", id, []);
},
has(target, prop) {
// return Reflect.has(target, prop);
return client.proxyReflect("has", id, [prop]);
},
defineProperty(target, prop, attributes) {
// return Reflect.defineProperty(target, prop, attributes);
return client.proxyReflect("defineProperty", id, [prop, attributes]);
},
deleteProperty(target, prop) {
// return Reflect.deleteProperty(target, prop);
return client.proxyReflect("deleteProperty", id, [prop]);
},
// TODO: Those functions might be interesting as well
/*
getOwnPropertyDescriptor(target, prop) {
// return Reflect.getOwnPropertyDescriptor(target, prop);
return client.proxyReflect("getOwnPropertyDescriptor", id, [prop]);
},
isExtensible(target) {
// return Reflect.isExtensible(target);
return client.proxyReflect("isExtensible", id, []);
},
preventExtensions(target) {
// return Reflect.preventExtensions(target);
return client.proxyReflect("preventExtensions", id, []);
},
getPrototypeOf(target) {
// return Reflect.getPrototypeOf(target);
return client.proxyReflect("getPrototypeOf", id, []);
},
setPrototypeOf(target, proto) {
// return Reflect.setPrototypeOf(target, proto);
return client.proxyReflect("setPrototypeOf", id, [proto]);
},*/
// For function objects
apply(target, thisArg, argumentsList) {
// Note: It can happen that a function gets called with a thisArg that cannot be serialized (due to it being an object on the worker thread)
// One solution is to provide a `[ObjectId]: ""` property
// TODO: Can it also happen that a function gets called with an easily serializeable thisArg that doesn't have an ObjectId?
// return Reflect.apply(target, thisArg, argumentsList);
return client.proxyReflect("apply", id, [thisArg[ObjectId], argumentsList]);
},
construct(target, argumentsList, newTarget) {
// return Reflect.construct(target, argumentsList, newTarget)
return client.proxyReflect("construct", id, [argumentsList, newTarget]);
},
}) as T;
}
/**
* Wraps an object in a proxy that does not proxy certain properties
*/
wrapExcluderProxy<T extends object>(obj: T, underlyingObject: T, exclude: Set<string | symbol>): T {
return new Proxy<T>(obj, {
get(target, prop, receiver) {
if (exclude.has(prop)) {
target = underlyingObject;
}
const value = Reflect.get(target, prop, receiver);
if (typeof value !== "function") return value;
return new Proxy(value, {
apply(_, thisArg, args) {
const calledWithProxy = thisArg === receiver;
return Reflect.apply(value, calledWithProxy ? target : thisArg, args);
},
});
},
has(target, prop) {
if (exclude.has(prop)) {
target = underlyingObject;
}
return Reflect.has(target, prop);
},
});
}
/**
* Blocks until an proxy object promise has returned a result or an error
*/
thenSync<T>(obj: Promise<T>): T {
const objectId = (obj as any)[ObjectId];
if (!objectId) {
throw new Error("Not a proxy object");
}
const result = this.proxyPromise("then", objectId);
if (result.error) {
throw result.error;
}
return result.value;
}
}
function isSimplePrimitive(value: any) {
if (value === undefined) {
return true;
} else if (value === null) {
return true;
} else if (value === false) {
return true;
} else if (value === true) {
return true;
} else if (typeof value === "number") {
return true;
} else if (value instanceof Date) {
return true;
} else {
return false;
}
}
function isSymbolPrimitive(value: any) {
if (typeof value === "symbol" && KNOWN_SYMBOLS.includes(value)) {
return true;
}
return false;
}
function isVariableLengthPrimitive(value: any) {
if (typeof value === "string") {
return true;
} else if (typeof value === "bigint") {
return true;
}
}
export type ProxyMessage =
| {
type: "proxy_reflect";
method: Exclude<keyof typeof Reflect, "apply">;
/**
* An object id
*/
target: string;
/**
* Further parameters. Have to be serialized
*/
args?: any[];
}
| {
type: "proxy_reflect";
method: "apply";
/**
* An object id
*/
target: string;
/**
* An object id
*/
thisArg: string;
/**
* Further parameters. Have to be serialized
*/
args?: any[];
}
| {
/** For requesting more bytes from the shared memory*/
type: "proxy_shared_memory";
}
| {
type: "proxy_print_object";
target: string;
}
| {
type: "proxy_promise";
method: "then";
target: string;
};