Window rendering in Javascript using WebGPU (#1027)

* Implement `runWindow` for JS

* Implement support functions

* Implement js bindings

* performance is lowercase in JS for some reason, breaking past conventions

* Wait for `document.body` to exist

* Fix typo in event name

* Fix workgroup size mangling hack

* Inline the buffer creation to make sure the same device instance is used to create them as the one associated with the canvas

* Fix typo on one buffer creation call

* Make the canvas actually use the real width and height, and resize appropriately

* This doesn't seem right, but let's see what it does

* Revert "This doesn't seem right, but let's see what it does"

This reverts commit ff048e0.

* Perhaps this is the source of the rendering glitch

* After reviewing the Rust code, I think this is the source of the error

* It shouldn't be printing 'bye' in the console

alan_compiler/src/std/root.ln

@@ -117,7 +117,7 @@ type Js = Env{"ALAN_OUTPUT_LANG"} == "js";
 
 // Importing the Root Scope backing implementation and supporting 3rd party libraries
 type{Rs} RootBacking = Rust{"alan_std" @ "https://github.com/alantech/alan.git"};
-type{Js} RootBacking = Node{"alan_std" @ "https://github.com/alantech/alan.git"};
+type{Js} RootBacking = Node{"alan_std" @ "https://github.com/alantech/alan.git#js-window"};
 
 // Defining derived types
 type void = ();
@@ -5376,22 +5376,37 @@ fn map{G, G2}(gb: GBuffer, f: (G, gu32) -> G2) {
 
 /// Window-related bindings
 type{Rs} Window = Binds{"alan_std::AlanWindowContext" <- RootBacking};
+type{Js} Window = Binds{"AlanWindowContext"};
 type{Rs} Frame = Binds{"alan_std::AlanWindowFrame" <- RootBacking};
+type{Js} Frame = Binds{"AlanWindowFrame"};
 
 fn{Rs} window "alan_std::run_window" <- RootBacking :: (Mut{(Mut{Window}) -> ()}, Mut{(Mut{Window}) -> u32[]}, (Frame) -> GPGPU[]) -> ()!;
+fn{Js} window "alan_std.runWindow" <- RootBacking :: (Window -> (), Window -> u32[], Frame -> GPGPU[]) -> ()!;
 fn{Rs} width Method{"width"} :: Window -> u32;
+fn{Js} width "alan_std.contextWidth" <- RootBacking :: Window -> u32;
 fn{Rs} height Method{"height"} :: Window -> u32;
+fn{Js} height "alan_std.contextHeight" <- RootBacking :: Window -> u32;
 fn{Rs} bufferWidth Method{"buffer_width"} :: Window -> u32;
+fn{Js} bufferWidth "alan_std.contextBufferWidth" <- RootBacking :: Window -> u32;
 fn{Rs} mouseX Method{"mouse_x"} :: Mut{Window} -> u32;
+fn{Js} mouseX "alan_std.contextMouseX" <- RootBacking :: Window -> u32;
 fn{Rs} mouseY Method{"mouse_y"} :: Mut{Window} -> u32;
+fn{Js} mouseY "alan_std.contextMouseY" <- RootBacking :: Window -> u32;
 fn{Rs} cursorVisible Method{"cursor_visible"} :: Mut{Window} -> ();
+fn{Js} cursorVisible "alan_std.contextCursorVisible" <- RootBacking :: Window -> ();
 fn{Rs} cursorInvisible Method{"cursor_invisible"} :: Mut{Window} -> ();
+fn{Js} cursorInvisible "alan_std.contextCursorInvisible" <- RootBacking :: Window -> ();
 fn{Rs} transparent Method{"transparent"} :: Mut{Window} -> ();
+fn{Js} transparent "alan_std.contextTransparent" <- RootBacking :: Window -> ();
 fn{Rs} opaque Method{"opaque"} :: Mut{Window} -> ();
+fn{Js} opaque "alan_std.contextOpaque" <- RootBacking :: Window -> ();
 fn{Rs} runtime Method{"runtime"} :: Window -> u32;
+fn{Js} runtime "alan_std.contextRuntime" <- RootBacking :: Window -> u32;
 fn{Rs} context Property{"context.clone()"} :: Frame -> GBuffer;
+fn{Js} context "alan_std.frameContext" <- RootBacking :: Frame -> GBuffer;
 fn{Rs} framebuffer Property{"framebuffer.clone()"} :: Frame -> GBuffer;
-fn{Rs} pixel Frame = gFor(-1, -2); // Magic numbers for the binding
+fn{Js} framebuffer "alan_std.frameFramebuffer" <- RootBacking :: Frame -> GBuffer;
+fn pixel Frame = gFor(-1, -2); // Magic numbers for the binding
 
 /// Process exit-related bindings
 fn{Rs} ExitCode "std::process::ExitCode::from" :: Own{u8} -> ExitCode;

alan_std.js

@@ -970,3 +970,266 @@ export async function replaceBuffer(b, v) {
   g.queue.submit([encoder.finish()]);
   tempBuffer.destroy();
 }
+
+/// Window-related types and functions
+
+export function contextWidth(context) {
+  return new U32(context.canvas.width);
+}
+
+export function contextHeight(context) {
+  return new U32(context.canvas.height);
+}
+
+export function contextBufferWidth(context) {
+  return new U32(context.bufferWidth / 4);
+}
+
+export function contextMouseX(context) {
+  if (typeof(context.mouseX) === "undefined") {
+    context.mouseX = new U32(0);
+    context.mouseY = new U32(0);
+  }
+  return context.mouseX;
+}
+
+export function contextMouseY(context) {
+  if (typeof(context.mouseY) === "undefined") {
+    context.mouseX = new U32(0);
+    context.mouseY = new U32(0);
+  }
+  return context.mouseY;
+}
+
+export function contextCursorVisible(context) {
+  context.cursorVisible = true;
+}
+
+export function contextCursorInvisible(context) {
+  context.cursorVisible = false;
+}
+
+export function contextTransparent(context) {
+  context.transparent = true;
+}
+
+export function contextOpaque(context) {
+  context.transparent = false;
+}
+
+export function contextRuntime(context) {
+  return f32AsU32(new F32((performance.now() - context.start) / 1000.0));
+}
+
+export function frameContext(frame) {
+  return frame.context;
+}
+
+export function frameFramebuffer(frame) {
+  return frame.framebuffer;
+}
+
+export async function runWindow(initialContextFn, contextFn, gpgpuShaderFn) {
+  // None of this can run before `document.body` exists, so let's wait for that
+  await new Promise((r) => document.addEventListener("DOMContentLoaded", () => r()));
+  let context = {
+    canvas: undefined,
+    start: undefined,
+    bufferWidth: undefined,
+    mouseX: undefined,
+    mouseY: undefined,
+    cursorVisible: true,
+    transparent: false,
+  };
+  await initialContextFn(context);
+  context.start = performance.now();
+  // If the `initialContextFn` doesn't attach a canvas, we make one to take up the whole window
+  if (!context.canvas) {
+    let canvas = document.createElement('canvas');
+    canvas.setAttribute('id', 'AlanWindow');
+    document.body.appendChild(canvas);
+    canvas.style['z-index'] = 9001;
+    canvas.style['position'] = 'absolute';
+    canvas.style['left'] = '0px';
+    canvas.style['top'] = '0px';
+    canvas.style['width'] = '100%';
+    canvas.style['height'] = '100%';
+    canvas.width = window.innerWidth;
+    canvas.height = window.innerHeight;
+    document.body.addEventListener("resize", () => {
+      canvas.width = window.innerWidth;
+      canvas.height = window.innerHeight;
+    });
+    if (!context.cursorVisible) {
+      canvas.style['cursor'] = 'none';
+    }
+    context.canvas = canvas;
+  }
+  context.canvas.addEventListener("mousemove", (event) => {
+    if (typeof(context.mouseX) !== "undefined" || typeof(context.mouseY) !== "undefined") {
+      context.mouseX = new U32(event.offsetX);
+      context.mouseY = new U32(event.offsetY);
+    }
+  });
+  let surface = context.canvas.getContext('webgpu');
+  let adapter = await navigator.gpu.requestAdapter();
+  let device = await adapter.requestDevice();
+  let queue = device.queue;
+  surface.configure({
+    device,
+    format: 'bgra8unorm',
+    alphaMode: context.transparent ? 'premultiplied' : 'opaque',
+    usage: GPUTextureUsage.COPY_DST | GPUTextureUsage.RENDER_ATTACHMENT,
+    viewFormats: ['bgra8unorm'],
+  });
+  let contextBuffer = await device.createBuffer({
+    size: 16,
+    usage: storageBufferType(),
+    label: `buffer_${uuidv4().replaceAll('-', '_')}`,
+  });
+  contextBuffer.ValKind = U32;
+  let width = Math.max(1, context.canvas.width);
+  let height = Math.max(1, context.canvas.height);
+  context.bufferWidth = (4 * width) % 256 === 0 ? 4 * width : 4 * width + (256 - ((4 * width) % 256));
+  let bufferHeight = height;
+  let bufferSize = context.bufferWidth * bufferHeight;
+  let buffer = await device.createBuffer({
+    size: bufferSize,
+    usage: storageBufferType(),
+    label: `buffer_${uuidv4().replaceAll('-', '_')}`,
+  });
+  buffer.ValKind = U32;
+  let gpgpuShaders = await gpgpuShaderFn({ context: contextBuffer, framebuffer: buffer });
+  let redraw = async function() {
+    // First resize things if necessary
+    if (width !== context.canvas.width || height !== context.canvas.height) {
+      width = context.canvas.width;
+      height = context.canvas.height;
+      context.bufferWidth = (4 * width) % 256 === 0 ? 4 * width : 4 * width + (256 - ((4 * width) % 256));
+      bufferHeight = height;
+      bufferSize = context.bufferWidth * bufferHeight;
+      let oldBufferId = buffer.label;
+      let newBuffer = await device.createBuffer({
+        size: bufferSize,
+        usage: storageBufferType(),
+        label: `buffer_${uuidv4().replaceAll('-', '_')}`,
+      });
+      newBuffer.ValKind = U32;
+      for (let shader of gpgpuShaders) {
+        for (let group of shader.buffers) {
+          let idx = undefined;
+          for (let i = 0; i < group.length; i++) {
+            let buffer = group[i];
+            if (buffer.label == oldBufferId) {
+              idx = i;
+              break;
+            }
+          }
+          if (typeof(idx) !== 'undefined') {
+            group[idx] = newBuffer;
+          }
+        }
+      }
+      buffer.destroy();
+      buffer = newBuffer;
+    }
+    // Now, actually start drawing
+    let oldContextBufferId = contextBuffer.label;
+    let frame = surface.getCurrentTexture();
+    let encoder = device.createCommandEncoder();
+    let contextArray = await contextFn(context);
+    let newContextBuffer = await device.createBuffer({
+      mappedAtCreation: true,
+      size: contextArray.length * 4,
+      usage: storageBufferType(),
+      label: `buffer_${uuidv4().replaceAll('-', '_')}`,
+    });
+    let ab = newContextBuffer.getMappedRange();
+    let v = new Uint32Array(ab);
+    for (let i = 0; i < contextArray.length; i++) {
+      v[i] = contextArray[i].valueOf();
+    }
+    newContextBuffer.unmap();
+    newContextBuffer.ValType = U32;
+    for (let gg of gpgpuShaders) {
+      if (typeof(gg.module) === "undefined") {
+        gg.module = device.createShaderModule({
+          code: gg.source,
+        });
+      }
+      if (typeof(gg.computePipeline) === "undefined") {
+        gg.computePipeline = device.createComputePipeline({
+          compute: {
+            module: gg.module,
+            entryPoint: gg.entryPoint,
+          },
+          layout: 'auto',
+        });
+      }
+      let bindGroups = [];
+      let cpass = encoder.beginComputePass();
+      cpass.setPipeline(gg.computePipeline);
+      for (let i = 0; i < gg.buffers.length; i++) {
+        let bindGroupLayout = gg.computePipeline.getBindGroupLayout(i);
+        let bindGroupBuffers = gg.buffers[i];
+        for (let j = 0; j < bindGroupBuffers.length; j++) {
+          if (bindGroupBuffers[j].label === oldContextBufferId) {
+            bindGroupBuffers[j] = newContextBuffer;
+          }
+        }
+        let bindGroupEntries = [];
+        for (let j = 0; j < bindGroupBuffers.length; j++) {
+          bindGroupEntries.push({
+            binding: j,
+            resource: { buffer: bindGroupBuffers[j] }
+          });
+        }
+        let bindGroup = device.createBindGroup({
+          layout: bindGroupLayout,
+          entries: bindGroupEntries,
+        });
+        bindGroups.push(bindGroup)
+      }
+      for (let i = 0; i < gg.buffers.length; i++) {
+        cpass.setBindGroup(i, bindGroups[i]);
+      }
+      let x = 0;
+      let y = 0;
+      switch (gg.workgroupSizes[0].val) {
+      case -1:
+        x = width;
+        break;
+      case -2:
+        x = height;
+        break;
+      default:
+        x = gg.workgroupSizes[0].val;
+      }
+      switch (gg.workgroupSizes[1].val) {
+      case -1:
+        y = width;
+        break;
+      case -2:
+        y = height;
+        break;
+      default:
+        y = gg.workgroupSizes[1].val;
+      }
+      let z = gg.workgroupSizes[2].val;
+      cpass.dispatchWorkgroups(x, y, z);
+      cpass.end();
+    }
+    contextBuffer.destroy();
+    contextBuffer = newContextBuffer;
+    encoder.copyBufferToTexture({
+      buffer,
+      bytesPerRow: context.bufferWidth,
+    }, {
+      texture: frame,
+    }, [width, height, 1]);
+    queue.submit([encoder.finish()]);
+    requestAnimationFrame(redraw);
+  }
+  requestAnimationFrame(redraw);
+  await new Promise((r) => {}); // Block this path forever to match Rust behavior
+}