starboard-python
Version:
Python cells for Starboard Notebook
2 lines (1 loc) • 11.2 kB
JavaScript
class AsyncMemory{constructor(sharedLock,sharedMemory){this.sharedLock=sharedLock??new SharedArrayBuffer(8*Int32Array.BYTES_PER_ELEMENT);this.lockAndSize=new Int32Array(this.sharedLock);if(this.lockAndSize.length<8)throw new Error("Expected an sharedLock with at least 8x32 bytes");this.sharedMemory=sharedMemory??new SharedArrayBuffer(1024);this.memory=new Uint8Array(this.sharedMemory);if(this.sharedMemory.byteLength<1024)throw new Error("Expected an sharedMemory with at least 1024 bytes")}lockWorker(){const oldValue=Atomics.compareExchange(this.lockAndSize,AsyncMemory.LOCK_WORKER_INDEX,AsyncMemory.UNLOCKED,AsyncMemory.LOCKED);if(oldValue!==AsyncMemory.UNLOCKED)throw new Error(`Cannot lock worker, the worker has to be unlocked ${AsyncMemory.UNLOCKED} !== ${oldValue}`)}lockSize(){const oldValue=Atomics.compareExchange(this.lockAndSize,AsyncMemory.LOCK_SIZE_INDEX,AsyncMemory.UNLOCKED,AsyncMemory.LOCKED);if(oldValue!==AsyncMemory.UNLOCKED)throw new Error(`Cannot set size flag, the size has to be unlocked ${AsyncMemory.UNLOCKED} !== ${oldValue}`)}waitForSize(){Atomics.wait(this.lockAndSize,AsyncMemory.LOCK_SIZE_INDEX,AsyncMemory.LOCKED)}writeSize(value){return Atomics.store(this.lockAndSize,AsyncMemory.SIZE_INDEX,value)}readSize(){return Atomics.load(this.lockAndSize,AsyncMemory.SIZE_INDEX)}unlockSize(){const oldValue=Atomics.compareExchange(this.lockAndSize,AsyncMemory.LOCK_SIZE_INDEX,AsyncMemory.LOCKED,AsyncMemory.UNLOCKED);if(oldValue!=AsyncMemory.LOCKED)throw new Error("Tried to unlock, but was already unlocked");Atomics.notify(this.lockAndSize,AsyncMemory.LOCK_SIZE_INDEX)}forceUnlockSize(){const oldValue=Atomics.compareExchange(this.lockAndSize,AsyncMemory.LOCK_SIZE_INDEX,AsyncMemory.LOCKED,AsyncMemory.UNLOCKED);if(oldValue!=AsyncMemory.LOCKED)Atomics.store(this.lockAndSize,AsyncMemory.LOCK_SIZE_INDEX,AsyncMemory.UNLOCKED);Atomics.notify(this.lockAndSize,AsyncMemory.LOCK_SIZE_INDEX)}unlockWorker(){const oldValue=Atomics.compareExchange(this.lockAndSize,AsyncMemory.LOCK_WORKER_INDEX,AsyncMemory.LOCKED,AsyncMemory.UNLOCKED);if(oldValue!=AsyncMemory.LOCKED)throw new Error("Tried to unlock, but was already unlocked");Atomics.notify(this.lockAndSize,AsyncMemory.LOCK_WORKER_INDEX)}}AsyncMemory.LOCK_WORKER_INDEX=0;AsyncMemory.LOCK_SIZE_INDEX=2;AsyncMemory.SIZE_INDEX=4;AsyncMemory.UNLOCKED=0;AsyncMemory.LOCKED=1;const SERIALIZATION={UNDEFINED:0,NULL:1,FALSE:2,TRUE:3,NUMBER:4,DATE:5,KNOWN_SYMBOL:6,STRING:10,BIGINT:11,OBJECT:255};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];new TextEncoder;const textDecoder=new TextDecoder("utf-8");const decodeFloat=useFloatDecoder();function useFloatDecoder(){const temp=new ArrayBuffer(8);const tempFloat64=new Float64Array(temp);const tempUint8=new Uint8Array(temp);return function(value){tempUint8.set(value);return tempFloat64[0]}}const ObjectId=Symbol.for("id");class ObjectProxyClient{constructor(memory,postMessage){this.memory=memory;this.postMessage=postMessage}serializePostMessage(value){if(isSimplePrimitive(value))return value;else if(isSymbolPrimitive(value))return{symbol:KNOWN_SYMBOLS.indexOf(value)};else if(isVariableLengthPrimitive(value))return value;else if(void 0!==value[ObjectId])return{id:value[ObjectId]};else return{value}}deserializeMemory(memory){const numberOfBytes=memory.readSize();let resultBytes;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)}if(resultBytes[0]===SERIALIZATION.UNDEFINED)return;else if(resultBytes[0]===SERIALIZATION.NULL)return null;else if(resultBytes[0]===SERIALIZATION.FALSE)return!1;else if(resultBytes[0]===SERIALIZATION.TRUE)return!0;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}else if(resultBytes[0]===SERIALIZATION.STRING)return textDecoder.decode(resultBytes.slice(1,numberOfBytes+1));else if(resultBytes[0]===SERIALIZATION.BIGINT)return BigInt(textDecoder.decode(resultBytes.slice(1,numberOfBytes+1)));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}}proxyReflect(method,targetId,args){let value;try{this.memory.lockWorker();this.memory.lockSize();this.memory.writeSize(0);if("apply"===method)this.postMessage({type:"proxy_reflect",method,target:targetId,thisArg:args[0],args:args[1].map((v=>this.serializePostMessage(v)))});else this.postMessage({type:"proxy_reflect",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{this.memory.forceUnlockSize();this.memory.unlockWorker()}return value}proxyPromise(method,targetId){let value;try{this.memory.lockWorker();this.memory.lockSize();this.memory.writeSize(0);this.postMessage({type:"proxy_promise",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{this.memory.forceUnlockSize();this.memory.unlockWorker()}return value}isFunction(id){return id.endsWith("-f")}getObjectProxy(id){const client=this;return new Proxy(this.isFunction(id)?function(){}:{},{get(target,prop,receiver){if(prop===ObjectId)return id;const value=client.proxyReflect("get",id,[prop,receiver]);if("function"!==typeof value)return value;return new Proxy(value,{apply(_,thisArg,argumentsList){const calledWithProxy=thisArg===receiver;const functionReturnValue=client.proxyReflect("apply",value[ObjectId],[calledWithProxy?id:thisArg[ObjectId],argumentsList]);return functionReturnValue}})},set(target,prop,value,receiver){return client.proxyReflect("set",id,[prop,value,receiver])},ownKeys(target){return client.proxyReflect("ownKeys",id,[])},has(target,prop){return client.proxyReflect("has",id,[prop])},defineProperty(target,prop,attributes){return client.proxyReflect("defineProperty",id,[prop,attributes])},deleteProperty(target,prop){return client.proxyReflect("deleteProperty",id,[prop])},apply(target,thisArg,argumentsList){return client.proxyReflect("apply",id,[thisArg[ObjectId],argumentsList])},construct(target,argumentsList,newTarget){return client.proxyReflect("construct",id,[argumentsList,newTarget])}})}wrapExcluderProxy(obj,underlyingObject,exclude){return new Proxy(obj,{get(target,prop,receiver){if(exclude.has(prop))target=underlyingObject;const value=Reflect.get(target,prop,receiver);if("function"!==typeof value)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)}})}thenSync(obj){const objectId=obj[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){if(void 0===value)return!0;else if(null===value)return!0;else if(!1===value)return!0;else if(!0===value)return!0;else if("number"===typeof value)return!0;else if(value instanceof Date)return!0;else return!1}function isSymbolPrimitive(value){if("symbol"===typeof value&&KNOWN_SYMBOLS.includes(value))return!0;return!1}function isVariableLengthPrimitive(value){if("string"===typeof value)return!0;else if("bigint"===typeof value)return!0}var _a;function assertUnreachable(_x){throw new Error("This case should have never been reached")}class KernelManager{constructor(){this.kernels=new Map;this.input=()=>"\n";this[_a]="";self.addEventListener("message",(async e=>{if(!e.data){console.warn("Kernel worker received unexpected message:",e);return}const data=e.data;switch(data.type){case"initialize":if(data.asyncMemory){const asyncMemory=new AsyncMemory(data.asyncMemory.lockBuffer,data.asyncMemory.dataBuffer);this.proxy=new ObjectProxyClient(asyncMemory,(message=>{this.postMessage(message)}));if(data.getInputId)this.input=this.proxy.getObjectProxy(data.getInputId);if(data.filesystemId){const proxy=this.proxy;const asyncFs=this.proxy.getObjectProxy(data.filesystemId);this.syncFs={get(opts){return proxy.thenSync(asyncFs.get(opts))},put(opts){return proxy.thenSync(asyncFs.put(opts))},delete(opts){return proxy.thenSync(asyncFs.delete(opts))},move(opts){return proxy.thenSync(asyncFs.move(opts))},listDirectory(opts){return proxy.thenSync(asyncFs.listDirectory(opts))}}}}else console.warn("Missing async memory, accessing objects from the main thread will not work. Please make sure that COOP/COEP is enabled.");break;case"import_kernel":try{if("url"===data.source.type)importScripts(data.source.url);else{const blob=new Blob([data.source.code],{type:"text/javascript"});importScripts(URL.createObjectURL(blob))}const KernelClass=globalThis[data.className];if(!data.options.id)data.options.id=data.kernelId;const kernel=new KernelClass(data.options);this.kernels.set(kernel.kernelId,kernel);kernel.init().then((()=>{this.postMessage({type:"kernel_initialized",kernelId:kernel.kernelId})}))}catch(e){this.postMessage({type:"error",kernelId:data.kernelId,id:"",error:e+""})}break;case"run":try{const kernel=this.kernels.get(data.kernelId);if(!kernel)throw new Error("Failed to find kernel with id "+data.kernelId);const result=await kernel.runCode(data.code);this.postMessage({type:"result",kernelId:kernel.kernelId,id:data.id,value:result})}catch(e){this.postMessage({type:"error",kernelId:data.kernelId,id:data.id,error:e+""})}break;case"custom":{const kernel=this.kernels.get(data.kernelId);if(kernel)kernel.customMessage(data.message);else console.warn("Custom message was sent to an nonexistent kernel",data);break}default:assertUnreachable();break}}))}postMessage(message){self.postMessage(message)}log(kernel,...args){this.postMessage({kernelId:kernel.kernelId,type:"console",method:"log",data:args})}logWarning(kernel,...args){this.postMessage({kernelId:kernel.kernelId,type:"console",method:"warn",data:args})}logError(kernel,...args){this.postMessage({kernelId:kernel.kernelId,type:"console",method:"error",data:args})}}_a=ObjectId;globalThis.manager=new KernelManager;