Properly type rust functions/worker

rezasm-app/rezasm-wasm/src/lib.rs

@@ -13,7 +13,7 @@ use rezasm_core::util::as_any::AsAny;
 use rezasm_web_core::{
     get_exit_status, get_memory_bounds, get_memory_slice, get_register_names, get_register_value,
     get_register_values, get_simulator_mut, get_word_size, initialize_simulator, is_completed,
-    load, reset, step, stop,
+    load, reset, step, step_back, stop,
 };
 use wasm_bindgen::prelude::*;
 
@@ -37,6 +37,11 @@ pub fn wasm_step() -> Result<(), String> {
     step()
 }
 
+#[wasm_bindgen]
+pub fn wasm_step_back() -> Result<(), String> {
+    step_back()
+}
+
 #[wasm_bindgen]
 pub fn wasm_is_completed() -> bool {
     is_completed()

src/components/Code.jsx

@@ -1,6 +1,6 @@
 import React, {useEffect, useState} from "react";
 import RegistryView from "./RegistryView.jsx";
-import {loadWasm} from "../rust_functions.js";
+import {loadWasm} from "../rust_functions.ts";
 import {Tabs, Tab} from "./Tabs.jsx";
 
 import MemoryView from "./MemoryView.jsx";

src/components/Console.jsx

@@ -1,7 +1,7 @@
 import React, {useCallback, useEffect, useReducer, useRef, useState} from "react";
 import {listen} from "@tauri-apps/api/event";
 import {CALLBACKS_TRIGGERS, CALLBACK_TYPES, STATE} from "./simulator.ts";
-import {RUST} from "../rust_functions.js";
+import {RUST} from "../rust_functions.ts";
 import { debounce } from "lodash";
 
 const ENTER = 13;

src/components/MemoryView.jsx

@@ -1,5 +1,5 @@
 import React, {useCallback, useEffect, useRef, useState} from "react";
-import {RUST} from "../rust_functions.js";
+import {RUST} from "../rust_functions.ts";
 import {CALLBACK_TYPES, CALLBACKS_TRIGGERS} from "./simulator.ts";
 
 const WIDTH = 4;

src/components/RegistryView.jsx

@@ -1,5 +1,5 @@
 import React, {useEffect, useState} from "react";
-import {RUST} from "../rust_functions.js";
+import {RUST} from "../rust_functions.ts";
 import {CALLBACK_TYPES, CALLBACKS_TRIGGERS} from "./simulator.ts";
 
 function RegistryView({loaded, registerCallback}) {

src/window.d.ts → src/globals.d.ts

@@ -3,6 +3,9 @@ import WorkerPromise from "webworker-promise";
 declare global {
     interface Window {
         __WASM_DEFINED__?: boolean;
+        __WASM_LOADED__?: boolean;
         worker: WorkerPromise;
+        // eslint-disable-next-line no-unused-vars
+        emitPrintString?: (string: string) => void;
     }
 }

src/rust_functions.js → src/rust_functions.ts

@@ -1,7 +1,10 @@
+/* eslint-disable no-unused-vars */
 import {invoke} from "@tauri-apps/api/tauri";
 import WorkerPromise from "webworker-promise";
+import workerURL from "./worker.ts?worker&url";
+import {type Message, type ValidWasmCommandStrings} from "./worker.ts";
 
-const setsEqual = (xs, ys) => xs.size === ys.size && [...xs].every((x) => ys.has(x));
+const setsEqual = <T>(xs: Set<T>, ys: Set<T>) => xs.size === ys.size && [...xs].every((x) => ys.has(x));
 
 const isWasmLoaded = () => {
     return callWorkerFunction({command: "status"});
@@ -11,7 +14,7 @@ const loadWasm = async () => {
     return import("../wasm/rezasm_wasm.js").then(() => {
         if (!window.__WASM_DEFINED__) {
             window.__WASM_DEFINED__ = true;
-            window.worker = new WorkerPromise(new Worker("/src/worker.js", { type: "module" }));
+            window.worker = new WorkerPromise(new Worker(workerURL, { type: "module" }));
             window.worker.postMessage({command: "ping"}).then((e) => {
                 return e === "pong";
             });
@@ -21,7 +24,7 @@ const loadWasm = async () => {
     });
 };
 
-const callWorkerFunction = (message) => {
+const callWorkerFunction = (message: Message) => {
     return new Promise((resolve, reject) => {
         window.worker.postMessage(message)
             .then(result => resolve(result))
@@ -34,14 +37,15 @@ const callWorkerFunction = (message) => {
 
 // name is the name of the function in rust (without "tauri_" or "wasm_")
 // shape is an array describing the keys that are expected to be defined in props
-const get_rust_function = (name, shape) => {
+export const get_rust_function = (name: ValidWasmCommandStrings, shape?: string[]) => {
     shape = shape ?? [];
     const shapeSet = new Set(shape);
-    return async (props) => {
+    return async (props: Record<string, unknown>) => {
         props = props ?? {};
         if (!setsEqual(shapeSet, new Set(Object.keys(props)))) {
             throw new Error(`Function '${name} passed with unexpected shape'`);
         }
+        // @ts-expect-error -- This is not always going to exist, but the compiler doesn't know that
         if (window.__TAURI_IPC__) {
             return invoke(`tauri_${name}`, props);
         } else {
@@ -53,6 +57,51 @@ const get_rust_function = (name, shape) => {
     };
 };
 
+export interface RustFunctions {
+    LOAD: (
+        props: {lines: string}
+    ) => Promise<void>;
+    STEP: (
+        props: Record<string, never>
+    ) => Promise<void>;
+    STEP_BACK: (
+        props: Record<string, never>
+    ) => Promise<void>;
+    RESET: (
+        props: Record<string, never>
+    ) => Promise<void>;
+    STOP: (
+        props: Record<string, never>
+    ) => Promise<void>;
+    IS_COMPLETED: (
+        props: Record<string, never>
+    ) => Promise<boolean>;
+    GET_EXIT_STATUS: (
+        props: Record<string, never>
+    ) => Promise<bigint>;
+    GET_REGISTER_VALUE: (
+        props: {register: string}
+    ) => Promise<bigint | undefined>;
+    GET_REGISTER_NAMES: (
+        props: Record<string, never>
+    ) => Promise<string[]>;
+    GET_REGISTER_VALUES: (
+        props: Record<string, never>
+    ) => Promise<BigInt64Array>;
+    GET_MEMORY_BOUNDS: (
+        props: Record<string, never>
+    ) => Promise<BigInt64Array>;
+    GET_MEMORY_SLICE: (
+        props: {address: number, length: number}
+    ) => Promise<BigInt64Array>;
+    GET_WORD_SIZE: (
+        props: Record<string, never>
+    ) => Promise<number>;
+    RECEIVE_INPUT: (
+        props: {data: string}
+    ) => Promise<void>;
+}
+
 const RUST = {
     LOAD: get_rust_function("load", ["lines"]),
     STEP: get_rust_function("step"),
@@ -68,7 +117,7 @@ const RUST = {
     GET_MEMORY_SLICE: get_rust_function("get_memory_slice", ["address", "length"]),
     GET_WORD_SIZE: get_rust_function("get_word_size"),
     RECEIVE_INPUT: get_rust_function("receive_input", ["data"]),
-};
+} as RustFunctions;
 
 export {
     RUST,

src/vite-env.d.ts

@@ -0,0 +1 @@
+/// <reference types="vite/client" />

src/worker.ts

@@ -0,0 +1,79 @@
+import registerWebworker from "webworker-promise/lib/register";
+
+type WasmExports = Omit<typeof import("../wasm/rezasm_wasm.js"), "default">;
+
+// https://stackoverflow.com/questions/53501721/typescript-exclude-property-key-when-starts-with-target-string
+// eslint-disable-next-line no-unused-vars,@typescript-eslint/no-unused-vars
+type FilterStartsWith<Set, Needle extends string> = Set extends `${Needle}${infer _}` ? Set : never
+// https://stackoverflow.com/questions/75418099/shorter-keys-with-typescript-key-remapping-removing-prefix
+type RemovePrefix<K extends PropertyKey, Prefix extends string> =
+    K extends `${Prefix}${infer P}` ? P : K
+
+type NonWasmCommandStrings = "status" | "ping"
+export type ValidWasmCommandStrings = RemovePrefix<FilterStartsWith<keyof WasmExports, "wasm_">, "wasm_">;
+type ValidCommandStrings = ValidWasmCommandStrings | NonWasmCommandStrings;
+type WasmFunctions = Pick<WasmExports, `wasm_${ValidWasmCommandStrings}`>
+
+let wasmFunctions: WasmFunctions;
+
+export interface Message {
+    command: ValidCommandStrings
+    argument?: Record<string, unknown>
+    shape?: string[]
+}
+
+if (self.__WASM_LOADED__ === undefined && self.document === undefined) {
+    self.__WASM_LOADED__ = false;
+    import("../wasm/rezasm_wasm.js").then(async (module) => {
+        await module.default();
+        wasmFunctions = Object.entries(module)
+            .filter(([x,]) => x.startsWith("wasm_"))
+            .reduce((left, [x, y]) => {
+                left[x] = y;
+                return left;
+            }, {} as Record<string, unknown>) as WasmFunctions;
+        self.__WASM_LOADED__ = true;
+        console.log("WebAssembly code loaded");
+    }).catch((error) => {
+        console.log("WebAssembly could not load", error);
+    });
+
+    const worker = registerWebworker(async (message: Message) => {
+        const command = message.command;
+        if (command === "status") {
+            return self.__WASM_LOADED__;
+        }
+        const data = message.argument ?? {};
+        const shape = message.shape ?? [];
+        const functionArguments: unknown[] = shape.map(arg => {
+            const argument = data[arg];
+            if (argument !== undefined) {
+                return argument;
+            } else {
+                throw new Error(`Function '${command}' passed without required argument ${arg}`);
+            }
+        });
+
+        if (command === "ping") {
+            return "pong";
+        }
+
+        const wasmFunction = wasmFunctions[`wasm_${command}`];
+
+        if (wasmFunction) {
+            try {
+                // @ts-expect-error -- There is no easy way to ensure this is correctly typed here.
+                // We have this typed correctly in rust_functions.ts.
+                return wasmFunction(...functionArguments);
+            } catch (error) {
+                throw new Error(String(error));
+            }
+        } else {
+            throw new Error(`Invalid command: '${command}'`);
+        }
+    });
+
+    self.emitPrintString = (string: string) => {
+        worker.emit("wasm_print", string);
+    };
+}

tsconfig.json

@@ -20,5 +20,5 @@
     "noUnusedParameters": true,
     "noFallthroughCasesInSwitch": true
   },
-  "include": ["src"]
+  "include": ["src", "wasm"]
 }

vite.config.js

@@ -8,22 +8,15 @@ export default defineConfig({
     plugins: [react(), wasm(), topLevelAwait()],
 
     build: {
-        rollupOptions: {
-            input: {
-                app: "./index.html",
-                worker: "./src/worker.js"
-            },
-            output: {
-                entryFileNames: assetInfo => {
-                    return assetInfo.name === "worker" ? "src/[name].js" : "assets/js/[name]-[hash].js";
-                }
-            },
-        },
         root: "src",
         outDir: "dist",
         emptyOutDir: true,
     },
 
+    worker: {
+        format: "es"
+    },
+
     // Vite options tailored for Tauri development and only applied in `tauri dev` or `tauri build`
     //
     // 1. prevent vite from obscuring rust errors