Detect language the WASM is written in instead of asking user to prov…

…ide it (#301)
FluffyLabs · Jan 21, 2025 · 4f4b3f6 · 4f4b3f6
1 parent 907f379
commit 4f4b3f6
Show file tree

Hide file tree

Showing 6 changed files with 38 additions and 61 deletions.
diff --git a/src/components/PvmSelect/index.tsx b/src/components/PvmSelect/index.tsx
@@ -4,8 +4,6 @@ import { useCallback, useEffect, useState } from "react";
 import path from "path-browserify";
 import { MultiSelect } from "@/components/ui/multi-select.tsx";
 import { AvailablePvms } from "@/types/pvm.ts";
-import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
-import { SupportedLangs } from "@/packages/web-worker/utils";
 import { ExternalLink } from "lucide-react";
 import { PvmTypes } from "@/packages/web-worker/types.ts";
 import { useDebuggerActions } from "@/hooks/useDebuggerActions";
@@ -39,7 +37,6 @@ export interface SelectedPvmWithPayload {
   label: string;
   params?: {
     file?: SerializedFile;
-    lang?: SupportedLangs;
     url?: string;
   };
   removable?: boolean;
@@ -72,7 +69,6 @@ export const PvmSelect = () => {
   const selectedPvms = useAppSelector(selectSelectedPvms);
   const pvmsWithPayload = useAppSelector(selectAllAvailablePvms);
   const dispatch = useAppDispatch();
-  const [selectedLang, setSelectedLang] = useState<SupportedLangs>(SupportedLangs.Rust);
 
   const mapValuesToPvmWithPayload = useCallback(
     (values: string[]) => {
@@ -100,7 +96,6 @@ export const PvmSelect = () => {
         id,
         type: PvmTypes.WASM_FILE,
         params: {
-          lang: selectedLang,
           file: await serializeFile(file),
         },
         label: `${id} - last modified: ${new Date(file.lastModified).toUTCString()}`,
@@ -229,18 +224,6 @@ export const PvmSelect = () => {
           <DialogTitle>Upload WASM file</DialogTitle>
           <DialogDescription>
             <div className="flex justify-between">
-              <div>
-                <Select onValueChange={(value: SupportedLangs) => setSelectedLang(value)} value={selectedLang}>
-                  <SelectTrigger>
-                    <SelectValue placeholder="Select Language" />
-                  </SelectTrigger>
-                  <SelectContent>
-                    <SelectItem value={SupportedLangs.Go}>Go</SelectItem>
-                    <SelectItem value={SupportedLangs.Rust}>Rust</SelectItem>
-                    <SelectItem value={SupportedLangs.AssemblyScript}>AssemblyScript</SelectItem>
-                  </SelectContent>
-                </Select>
-              </div>
               <div>
                 <Input
                   type="file"

diff --git a/src/packages/web-worker/command-handlers/load.ts b/src/packages/web-worker/command-handlers/load.ts
@@ -1,10 +1,10 @@
 import { logger } from "@/utils/loggerService";
-import { getMemorySize, loadArrayBufferAsWasm, SupportedLangs } from "../utils";
+import { getMemorySize, loadArrayBufferAsWasm } from "../utils";
 import { CommandStatus, PvmApiInterface, PvmTypes } from "../types";
 import { Pvm as InternalPvmInstance } from "@typeberry/pvm-debugger-adapter";
 import { deserializeFile, SerializedFile } from "@/lib/utils.ts";
 
-export type LoadParams = { type: PvmTypes; params: { url?: string; file?: SerializedFile; lang?: SupportedLangs } };
+export type LoadParams = { type: PvmTypes; params: { url?: string; file?: SerializedFile } };
 export type LoadResponse = {
   pvm: PvmApiInterface | null;
   memorySize: number | null;
@@ -24,7 +24,7 @@ const load = async (args: LoadParams): Promise<PvmApiInterface | null> => {
 
     logger.info("Load WASM from file", file);
     const bytes = await file.arrayBuffer();
-    return await loadArrayBufferAsWasm(bytes, args.params.lang);
+    return await loadArrayBufferAsWasm(bytes);
   } else if (args.type === PvmTypes.WASM_URL) {
     const url = args.params.url ?? "";
     const isValidUrl = Boolean(new URL(url));
@@ -37,7 +37,7 @@ const load = async (args: LoadParams): Promise<PvmApiInterface | null> => {
     const response = await fetch(url);
     const bytes = await response.arrayBuffer();
 
-    return await loadArrayBufferAsWasm(bytes, args.params.lang);
+    return await loadArrayBufferAsWasm(bytes);
   }
 
   return null;

diff --git a/src/packages/web-worker/types.ts b/src/packages/web-worker/types.ts
@@ -1,5 +1,4 @@
 import { CurrentInstruction, ExpectedState, HostCallIdentifiers, InitialState } from "@/types/pvm";
-import { SupportedLangs } from "./utils";
 import { WasmPvmShellInterface } from "./wasmBindgenShell";
 import { Pvm as InternalPvm } from "@/types/pvm";
 import { bytes } from "@typeberry/jam-host-calls";
@@ -49,7 +48,7 @@ type CommonWorkerRequestParams = { messageId: string };
 export type CommandWorkerRequestParams =
   | {
       command: Commands.LOAD;
-      payload: { type: PvmTypes; params: { url?: string; file?: SerializedFile; lang?: SupportedLangs } };
+      payload: { type: PvmTypes; params: { url?: string; file?: SerializedFile } };
     }
   | { command: Commands.INIT; payload: { program: Uint8Array; initialState: InitialState } }
   | { command: Commands.STEP; payload: { program: Uint8Array; stepsToPerform: number } }

diff --git a/src/packages/web-worker/utils.ts b/src/packages/web-worker/utils.ts
@@ -8,12 +8,6 @@ import { logger } from "@/utils/loggerService.tsx";
 import { PvmApiInterface } from "./types.ts";
 import { createAssemblyScriptWasmPvmShell } from "./wasmAsShell.ts";
 
-export enum SupportedLangs {
-  Go = "Go",
-  Rust = "Rust",
-  AssemblyScript = "AssemblyScript",
-}
-
 export function getState(pvm: PvmApiInterface): ExpectedState {
   const regs = isInternalPvm(pvm)
     ? (Array.from(pvm.getRegisters()) as RegistersArray)
@@ -80,34 +74,38 @@ export function isInternalPvm(pvm: PvmApiInterface): pvm is InternalPvm {
   return pvm instanceof InternalPvm;
 }
 
-export async function loadArrayBufferAsWasm(bytes: ArrayBuffer, lang?: SupportedLangs): Promise<PvmApiInterface> {
-  if (lang === SupportedLangs.Go) {
-    const go = new Go();
-    const wasmModule = await WebAssembly.instantiate(bytes, go.importObject);
-    go.run(wasmModule.instance);
-    logger.info("Go WASM module loaded", wasmModule.instance.exports);
-    const wasmPvmShell = createGoWasmPvmShell();
-    wasmPvmShell.__wbg_set_wasm(wasmModule.instance.exports);
-    return wasmPvmShell;
-  }
-
-  if (lang === SupportedLangs.Rust) {
-    const wasmModule = await WebAssembly.instantiate(bytes, {});
-    logger.info("Rust WASM module loaded", wasmModule.instance.exports);
-    const wasmPvmShell = createWasmPvmShell();
-    wasmPvmShell.__wbg_set_wasm(wasmModule.instance.exports);
-    return wasmPvmShell;
-  }
-
-  if (lang === SupportedLangs.AssemblyScript) {
-    const compiled = await WebAssembly.compile(bytes);
-    logger.info("AssemblyScript WASM module compiled", compiled);
-    const wasmPvmShell = await createAssemblyScriptWasmPvmShell(compiled);
-    logger.info("AssemblyScript WASM module loaded", wasmPvmShell);
-    return wasmPvmShell;
+export async function loadArrayBufferAsWasm(bytes: ArrayBuffer): Promise<PvmApiInterface> {
+  const go = new Go();
+  const wasmModule = await WebAssembly.compile(bytes);
+
+  try {
+    // Even though we're not always using Go, we instantiate it here as any Rust WASM has a subset of Go WASM import object (it's empty for Rust)
+    const wasmInstance = await WebAssembly.instantiate(wasmModule, go.importObject);
+    const wasmMethods = wasmInstance.exports;
+
+    if (wasmMethods.__wbindgen_add_to_stack_pointer) {
+      logger.info("Rust WASM module loaded", wasmInstance.exports);
+      const wasmPvmShell = createWasmPvmShell();
+      wasmPvmShell.__wbg_set_wasm(wasmInstance.exports);
+      return wasmPvmShell;
+    } else {
+      go.run(wasmInstance);
+      logger.info("Go WASM module loaded", wasmInstance.exports);
+      const wasmPvmShell = createGoWasmPvmShell();
+      wasmPvmShell.__wbg_set_wasm(wasmInstance.exports);
+      return wasmPvmShell;
+    }
+  } catch (_) {
+    try {
+      logger.info("AssemblyScript WASM module compiled", wasmModule);
+      const wasmPvmShell = await createAssemblyScriptWasmPvmShell(wasmModule);
+      logger.info("AssemblyScript WASM module loaded", wasmPvmShell);
+      return wasmPvmShell;
+    } catch (e) {
+      logger.warn(e);
+      throw new Error(`Error instantiating WASM. Possibly unsupported language.`);
+    }
   }
-
-  throw new Error(`Unsupported lang: ${lang}`);
 }
 
 export function getMemorySize(pvm: PvmApiInterface | null) {

diff --git a/src/store/debugger/debuggerSlice.ts b/src/store/debugger/debuggerSlice.ts
@@ -4,7 +4,6 @@ import { InstructionMode } from "@/components/Instructions/types.ts";
 import { RootState } from "@/store";
 import { isEqual } from "lodash";
 import { SelectedPvmWithPayload } from "@/components/PvmSelect";
-import { SupportedLangs } from "@/packages/web-worker/utils.ts";
 import { PvmTypes } from "@/packages/web-worker/types.ts";
 
 export interface DebuggerState {
@@ -42,15 +41,14 @@ const initialState: DebuggerState = {
       {
         id: AvailablePvms.POLKAVM,
         type: PvmTypes.WASM_URL,
-        params: { url: "https://todr.me/polkavm/pvm-metadata.json", lang: SupportedLangs.Rust },
+        params: { url: "https://todr.me/polkavm/pvm-metadata.json" },
         label: "PolkaVM",
       },
       {
         id: AvailablePvms.ANANAS,
         type: PvmTypes.WASM_URL,
         params: {
           url: "https://todr.me/anan-as/pvm-metadata.json",
-          lang: SupportedLangs.AssemblyScript,
         },
         label: "Anan-AS",
       },

diff --git a/src/store/workers/workersSlice.ts b/src/store/workers/workersSlice.ts
@@ -10,7 +10,6 @@ import {
   setStorage,
 } from "@/store/debugger/debuggerSlice.ts";
 import PvmWorker from "@/packages/web-worker/worker?worker&inline";
-import { SupportedLangs } from "@/packages/web-worker/utils.ts";
 import { virtualTrapInstruction } from "@/utils/virtualTrapInstruction.ts";
 import { logger } from "@/utils/loggerService";
 import { Commands, PvmTypes } from "@/packages/web-worker/types";
@@ -96,7 +95,7 @@ export const loadWorker = createAsyncThunk(
       payload,
     }: {
       id: string;
-      payload: { type: PvmTypes; params?: { url?: string; file?: SerializedFile; lang?: SupportedLangs } };
+      payload: { type: PvmTypes; params?: { url?: string; file?: SerializedFile } };
     },
     { getState, dispatch },
   ) => {