Write some very basic codegen using walrus

crates/wasm/Cargo.toml

@@ -6,6 +6,4 @@ edition = "2021"
 
 [dependencies]
 rose = { path = "../core" }
-
-[dev-dependencies]
-wasmtime = "9"
+walrus = "0.20"

crates/wasm/src/lib.rs

@@ -1,41 +1,90 @@
-#[cfg(test)]
-mod tests {
-    use wasmtime::{Caller, Engine, Func, Instance, Module, Store};
+use walrus::{ir::BinaryOp, FunctionBuilder, LocalId, Module, ModuleConfig, ValType};
 
-    // https://docs.rs/wasmtime/9.0.2/wasmtime/
-    #[test]
-    fn test_foo() -> wasmtime::Result<()> {
-        // Modules can be compiled through either the text or binary format
-        let engine = Engine::default();
-        let wat = r#"
-        (module
-            (import "host" "host_func" (func $host_hello (param i32)))
+pub const EXPORT_NAME: &str = "";
 
-            (func (export "hello")
-                i32.const 3
-                call $host_hello)
-        )
-    "#;
-        let module = Module::new(&engine, wat)?;
-
-        // All wasm objects operate within the context of a "store". Each
-        // `Store` has a type parameter to store host-specific data, which in
-        // this case we're using `4` for.
-        let mut store = Store::new(&engine, 4);
-        let host_func = Func::wrap(&mut store, |caller: Caller<'_, u32>, param: i32| {
-            println!("Got {} from WebAssembly", param);
-            println!("my host state is: {}", caller.data());
-        });
-
-        // Instantiation of a module requires specifying its imports and then
-        // afterwards we can fetch exports by name, as well as asserting the
-        // type signature of the function with `get_typed_func`.
-        let instance = Instance::new(&mut store, &module, &[host_func.into()])?;
-        let hello = instance.get_typed_func::<(), ()>(&mut store, "hello")?;
-
-        // And finally we can call the wasm!
-        hello.call(&mut store, ())?;
+fn type_to_wasm(t: rose::Type) -> ValType {
+    match t {
+        rose::Type::Bool => ValType::I32,
+        rose::Type::Int => ValType::I32,
+        rose::Type::Real => ValType::F64,
+        rose::Type::Size { val: _ } => todo!(),
+        rose::Type::Nat { bound: _ } => todo!(),
+        rose::Type::Var { id: _ } => todo!(),
+    }
+}
 
-        Ok(())
+pub fn compile(f: &rose::Def<rose::Function>, _generics: &[usize]) -> Vec<u8> {
+    let config = ModuleConfig::new();
+    let mut module = Module::with_config(config);
+    let params: Vec<ValType> = f.def.params.iter().map(|&t| type_to_wasm(t)).collect();
+    let results: Vec<ValType> = f.def.ret.iter().map(|&t| type_to_wasm(t)).collect();
+    let mut func = FunctionBuilder::new(&mut module.types, &params, &results);
+    let params: Vec<LocalId> = params.into_iter().map(|t| module.locals.add(t)).collect();
+    let mut body = func.func_body();
+    for &id in params.iter() {
+        body.local_get(id);
+    }
+    let locals: Vec<LocalId> = f
+        .def
+        .locals
+        .iter()
+        .map(|&t| module.locals.add(type_to_wasm(t)))
+        .collect();
+    for &instr in &f.def.body {
+        match instr {
+            rose::Instr::Generic { id: _ } => todo!(),
+            rose::Instr::Get { id } => {
+                body.local_get(locals[id.0]);
+            }
+            rose::Instr::Set { id } => {
+                body.local_set(locals[id.0]);
+            }
+            rose::Instr::Bool { val: _ } => todo!(),
+            rose::Instr::Int { val: _ } => todo!(),
+            rose::Instr::Real { val: _ } => todo!(),
+            rose::Instr::Vector { id: _ } => todo!(),
+            rose::Instr::Tuple { id: _ } => todo!(),
+            rose::Instr::Index => todo!(),
+            rose::Instr::Member { id: _ } => todo!(),
+            rose::Instr::Call { id: _ } => todo!(),
+            rose::Instr::Unary { op: _ } => todo!(),
+            rose::Instr::Binary { op } => match op {
+                rose::Binop::And => todo!(),
+                rose::Binop::Or => todo!(),
+                rose::Binop::EqBool => todo!(),
+                rose::Binop::NeqBool => todo!(),
+                rose::Binop::NeqInt => todo!(),
+                rose::Binop::LtInt => todo!(),
+                rose::Binop::LeqInt => todo!(),
+                rose::Binop::EqInt => todo!(),
+                rose::Binop::GtInt => todo!(),
+                rose::Binop::GeqInt => todo!(),
+                rose::Binop::NeqReal => todo!(),
+                rose::Binop::LtReal => todo!(),
+                rose::Binop::LeqReal => todo!(),
+                rose::Binop::EqReal => todo!(),
+                rose::Binop::GtReal => todo!(),
+                rose::Binop::GeqReal => todo!(),
+                rose::Binop::AddInt => todo!(),
+                rose::Binop::SubInt => todo!(),
+                rose::Binop::MulInt => todo!(),
+                rose::Binop::DivInt => todo!(),
+                rose::Binop::Mod => todo!(),
+                rose::Binop::AddReal => {
+                    body.binop(BinaryOp::F64Add);
+                }
+                rose::Binop::SubReal => todo!(),
+                rose::Binop::MulReal => todo!(),
+                rose::Binop::DivReal => todo!(),
+            },
+            rose::Instr::If => todo!(),
+            rose::Instr::Else => todo!(),
+            rose::Instr::End => todo!(),
+            rose::Instr::For { limit: _ } => todo!(),
+        }
     }
+    module
+        .exports
+        .add(EXPORT_NAME, func.finish(params, &mut module.funcs));
+    module.emit_wasm()
 }

crates/web/Cargo.toml

@@ -11,6 +11,7 @@ crate-type = ["cdylib"]
 console_error_panic_hook = "0.1"
 rose = { path = "../core", features = ["serde"] }
 rose-interp = { path = "../interp", features = ["serde"] }
+rose-wasm = { path = "../wasm" }
 serde = { version = "1", features = ["derive"] }
 serde-wasm-bindgen = "0.4"
 ts-rs = "6"

crates/web/src/lib.rs

@@ -460,3 +460,13 @@ pub fn interp(Func(f): &Func, args: JsValue) -> Result<JsValue, serde_wasm_bindg
     assert_eq!(ret.len(), 1);
     to_js_value(&ret[0])
 }
+
+#[wasm_bindgen(js_name = "wasmExportName")]
+pub fn wasm_export_name() -> String {
+    rose_wasm::EXPORT_NAME.to_owned()
+}
+
+#[wasm_bindgen]
+pub fn compile(Func(f): &Func, generics: &[usize]) -> Vec<u8> {
+    rose_wasm::compile(f, generics)
+}

packages/core/src/compile.ts

@@ -0,0 +1,18 @@
+import * as ffi from "./ffi.js";
+import { Fn, ToJs } from "./fn.js";
+import { Real } from "./real.js";
+
+/**
+ * Converts an abstract function into a concrete function using the interpreter.
+ */
+export const compile = async <const T extends readonly Real[]>(
+  f: Fn & ((...args: T) => Real),
+  generics: number[]
+): Promise<(...args: ToJs<T>) => number> => {
+  const bytes = ffi.compile(f.f, generics);
+  console.log(bytes);
+  const instance = await WebAssembly.instantiate(
+    await WebAssembly.compile(bytes)
+  );
+  return instance.exports[ffi.wasmExportName] as any;
+};

packages/core/src/ffi.ts

@@ -52,4 +52,9 @@ export const bake = (ctx: Context): Fn => {
 
 export const interp = (f: Fn, args: Val[]): Val => wasm.interp(f.f, args);
 
+export const wasmExportName = wasm.wasmExportName();
+
+export const compile = (f: Fn, generics: number[]) =>
+  wasm.compile(f.f, new Uint32Array(generics));
+
 export type { Type, Val };

packages/core/src/fn.ts

@@ -20,10 +20,14 @@ export interface Fn {
   f: ffi.Fn;
 }
 
-type Resolve<T extends readonly Reals[]> = {
+type FromType<T extends readonly Reals[]> = {
   [K in keyof T]: Real;
 };
 
+export type ToJs<T extends readonly Real[]> = {
+  [K in keyof T]: number;
+};
+
 const ffiType = (t: Type): ffi.Type => {
   switch (t.tag) {
     case "Bool":
@@ -41,8 +45,8 @@ const ffiType = (t: Type): ffi.Type => {
 // TODO: allow args other than `Real`
 export const fn = <const T extends readonly Reals[]>(
   types: T,
-  f: (...args: Resolve<T>) => Real // TODO: allow return other than `Real`
-): Fn & ((...args: Resolve<T>) => Real) => {
+  f: (...args: FromType<T>) => Real // TODO: allow return other than `Real`
+): Fn & ((...args: FromType<T>) => Real) => {
   // TODO: support closures
   if (context !== undefined)
     throw Error("can't define a function while defining another function");
@@ -55,7 +59,7 @@ export const fn = <const T extends readonly Reals[]>(
     // reverse because stack is LIFO
     for (let i = types.length - 1; i >= 0; --i)
       params.push({ ctx, id: ctx.set(paramTypes[i]) });
-    getReal(ctx, f(...(params.reverse() as Resolve<T>)));
+    getReal(ctx, f(...(params.reverse() as FromType<T>)));
     func = ffi.bake(ctx);
   } catch (e) {
     // `ctx` points into Wasm memory, so if we didn't finish the `ffi.bake` call
@@ -65,7 +69,7 @@ export const fn = <const T extends readonly Reals[]>(
   } finally {
     setCtx(undefined);
   }
-  const g = (...args: Resolve<T>): Real => {
+  const g = (...args: FromType<T>): Real => {
     throw Error("TODO");
   };
   g.params = types as unknown as Type[];

packages/core/src/index.test.ts

@@ -1,5 +1,5 @@
 import { expect, test } from "vitest";
-import { Real, add, div, fn, interp, mul, sub } from "./index.js";
+import { Real, add, compile, div, fn, interp, mul, sub } from "./index.js";
 
 test("2 + 2 = 4", () => {
   const f = fn([Real, Real], (x, y) => add(x, y));
@@ -12,3 +12,9 @@ test("basic arithmetic", () => {
   const g = interp(f);
   expect(g()).toBe(6);
 });
+
+test("compile", async () => {
+  const f = fn([Real, Real], (x, y) => add(x, y));
+  const g = await compile(f, []);
+  expect(g(2, 2)).toBe(4);
+});

packages/core/src/index.ts

@@ -17,6 +17,7 @@ export type Vec<T> = vec.Vec<T>;
 export const Vec = <T>(t: T, n: int.Int): Vecs<T> => ({ tag: "Vec", t, n });
 
 export { and, cond, iff, not, or, xor } from "./bool.js";
+export { compile } from "./compile.js";
 export { fn } from "./fn.js";
 export { int, mod } from "./int.js";
 export { interp } from "./interp.js";

packages/core/src/interp.ts

@@ -1,20 +1,16 @@
 import * as ffi from "./ffi.js";
-import { Fn } from "./fn.js";
+import { Fn, ToJs } from "./fn.js";
 import { Real } from "./real.js";
 
-type Resolve<T extends readonly Real[]> = {
-  [K in keyof T]: number;
-};
-
 /**
  * Converts an abstract function into a concrete function using the interpreter.
  */
 export const interp =
   <const T extends readonly Real[]>(
     f: Fn & ((...args: T) => Real)
-  ): ((...args: Resolve<T>) => number) =>
+  ): ((...args: ToJs<T>) => number) =>
   // just return a closure
-  (...args: Resolve<T>) => {
+  (...args: ToJs<T>) => {
     // that calls the interpreter
     const x = ffi.interp(
       f.f,

packages/core/tsconfig.json

@@ -1,9 +1,12 @@
 {
   "compilerOptions": {
     // basic stuff, or everything breaks
-    "module": "ESNext",
+    "module": "ES2022", // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Operators/await#top_level_await
     "moduleResolution": "node16",
-    "lib": ["ESNext"],
+    "lib": [
+      "DOM", // https://github.com/microsoft/TypeScript-DOM-lib-generator/issues/826
+      "ES2021" // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/FinalizationRegistry
+    ],
     "skipLibCheck": true,
 
     // build config