3D rendering in WASM demo (#106)

mkeeter · May 17, 2024 · 4d80ae5 · 4d80ae5
1 parent 76ae930
commit 4d80ae5
Show file tree

Hide file tree

Showing 5 changed files with 210 additions and 17 deletions.
diff --git a/wasm-demo/index.html b/wasm-demo/index.html
@@ -1,44 +1,44 @@
 <html>
   <head>
+    <meta content="text/html;charset=utf-8" http-equiv="Content-Type" />
+    <title>Fidget Demo</title>
     <style>
       div#editor {
         height: 512;
         width: 100%;
         display: block;
         outline: 1px solid #bbb;
+        margin-bottom: 10px;
       }
       div#output-outer {
         width: 100%;
         display: block;
-        margin-top: 10px;
         outline: 1px solid #bbb;
       }
       div#canvascol {
-        width: 100%;
         padding-right: 10px;
+        padding-left: 10px;
       }
       canvas#glcanvas {
         display: block;
-        margin-left: 10px;
         aspect-ratio: 1;
         max-height: 512px;
-        width: calc(min(100%, 512px));
         outline: 1px solid #bbb;
+        margin-bottom: 12px;
       }
       span#status {
         display: block;
-        margin-left: 10px;
-        margin-top: 15px;
         font-family: monospace;
+        float: right;
       }
       .row {
         display: flex;
       }
       .column {
         flex: 50%;
+        max-width: calc(min(100%, 512px));
       }
     </style>
-    <meta content="text/html;charset=utf-8" http-equiv="Content-Type" />
   </head>
   <body>
     <div class="row">
@@ -49,6 +49,11 @@
       <div class="column" id="canvascol">
         <canvas id="glcanvas" width="512" height="512"></canvas>
         <span id="status">Loading...</span>
+        <select name="mode" id="mode">
+          <option value="bitmap">2D (bitmap)</option>
+          <option value="normals">3D (normals)</option>
+          <option value="heightmap">3D (heightmap)</option>
+        </select>
       </div>
     </div>
   </body>

diff --git a/wasm-demo/index.ts b/wasm-demo/index.ts
@@ -5,11 +5,12 @@ import { EditorState } from "@codemirror/state";
 import { defaultKeymap } from "@codemirror/commands";
 
 import {
+  RenderMode,
   ResponseKind,
   ScriptRequest,
+  ScriptResponse,
   ShapeRequest,
   StartRequest,
-  ScriptResponse,
   WorkerRequest,
   WorkerResponse,
 } from "./message";
@@ -53,12 +54,39 @@ class App {
       };
       this.workers.push(worker);
     }
+
+    // Also re-render if the mode changes
+    const select = document.getElementById("mode");
+    select.addEventListener("change", this.onModeChanged.bind(this), false);
+  }
+
+  onModeChanged() {
+    const text = this.editor.view.state.doc.toString();
+    this.onScriptChanged(text);
   }
 
   onScriptChanged(text: string) {
     this.workers[0].postMessage(new ScriptRequest(text));
   }
 
+  getMode() {
+    const e = document.getElementById("mode") as HTMLSelectElement;
+    switch (e.value) {
+      case "bitmap": {
+        return RenderMode.Bitmap;
+      }
+      case "normals": {
+        return RenderMode.Normals;
+      }
+      case "heightmap": {
+        return RenderMode.Heightmap;
+      }
+      default: {
+        return null;
+      }
+    }
+  }
+
   onWorkerMessage(i: number, req: WorkerResponse) {
     switch (req.kind) {
       case ResponseKind.Image: {
@@ -90,8 +118,9 @@ class App {
         if (r.tape) {
           this.start_time = performance.now();
           this.workers_done = 0;
+          const mode = this.getMode();
           this.workers.forEach((w) => {
-            w.postMessage(new ShapeRequest(r.tape));
+            w.postMessage(new ShapeRequest(r.tape, mode));
           });
         }
         break;

diff --git a/wasm-demo/message.ts b/wasm-demo/message.ts
@@ -24,13 +24,21 @@ export class StartRequest {
   }
 }
 
+export enum RenderMode {
+  Bitmap,
+  Heightmap,
+  Normals,
+}
+
 export class ShapeRequest {
   kind: RequestKind.Shape;
   tape: Uint8Array;
+  mode: RenderMode;
 
-  constructor(tape: Uint8Array) {
+  constructor(tape: Uint8Array, mode: RenderMode) {
     this.tape = tape;
     this.kind = RequestKind.Shape;
+    this.mode = mode;
   }
 }
 

diff --git a/wasm-demo/src/lib.rs b/wasm-demo/src/lib.rs
@@ -46,7 +46,7 @@ pub fn deserialize_tape(data: Vec<u8>) -> Result<JsVmShape, String> {
 /// The image has a total size of `image_size` (on each side) and is divided
 /// into `0 <= pos < workers_per_side^2` tiles.
 #[wasm_bindgen]
-pub fn render_region(
+pub fn render_region_2d(
     shape: JsVmShape,
     image_size: usize,
     index: usize,
@@ -102,3 +102,129 @@ pub fn render_region(
     inner(shape.0, image_size, index, workers_per_side)
         .map_err(|e| format!("{e}"))
 }
+
+/// Renders a subregion of a heightmap, for webworker-based multithreading
+///
+/// The image has a total size of `image_size` (on each side) and is divided
+/// into `0 <= pos < workers_per_side^2` tiles.
+#[wasm_bindgen]
+pub fn render_region_heightmap(
+    shape: JsVmShape,
+    image_size: usize,
+    index: usize,
+    workers_per_side: usize,
+) -> Result<Vec<u8>, String> {
+    if index >= workers_per_side.pow(2) {
+        return Err("invalid index".to_owned());
+    }
+    if image_size % workers_per_side != 0 {
+        return Err(
+            "image_size must be divisible by workers_per_side".to_owned()
+        );
+    }
+    let (depth, _norm) =
+        render_3d_inner(shape.0, image_size, index, workers_per_side)
+            .map_err(|e| format!("{e}"))?;
+
+    // Convert into an image
+    Ok(depth
+        .into_iter()
+        .flat_map(|v| {
+            let d = (v as usize * 255 / image_size) as u8;
+            [d, d, d, 255]
+        })
+        .collect())
+}
+
+/// Renders a subregion with normals, for webworker-based multithreading
+///
+/// The image has a total size of `image_size` (on each side) and is divided
+/// into `0 <= pos < workers_per_side^2` tiles.
+#[wasm_bindgen]
+pub fn render_region_normals(
+    shape: JsVmShape,
+    image_size: usize,
+    index: usize,
+    workers_per_side: usize,
+) -> Result<Vec<u8>, String> {
+    if index >= workers_per_side.pow(2) {
+        return Err("invalid index".to_owned());
+    }
+    if image_size % workers_per_side != 0 {
+        return Err(
+            "image_size must be divisible by workers_per_side".to_owned()
+        );
+    }
+    let (_depth, norm) =
+        render_3d_inner(shape.0, image_size, index, workers_per_side)
+            .map_err(|e| format!("{e}"))?;
+
+    // Convert into an image
+    Ok(norm
+        .into_iter()
+        .flat_map(|[r, g, b]| [r, g, b, 255])
+        .collect())
+}
+
+fn render_3d_inner(
+    shape: VmShape,
+    image_size: usize,
+    index: usize,
+    workers_per_side: usize,
+) -> Result<(Vec<u32>, Vec<[u8; 3]>), Error> {
+    let mut current_depth = vec![];
+    let mut current_norm = vec![];
+
+    // Work from front to back, so we can bail out early if the image is full
+    for z in (0..workers_per_side).rev() {
+        // Corner position in [0, workers_per_side] coordinates
+        let mut corner = nalgebra::Vector3::new(
+            index / workers_per_side,
+            index % workers_per_side,
+            z,
+        )
+        .cast::<f32>();
+        // Corner position in [-1, 1] coordinates
+        corner = (corner * 2.0 / workers_per_side as f32).add_scalar(-1.0);
+
+        // Scale of each tile
+        let scale = 2.0 / workers_per_side as f32;
+
+        // Tile center
+        let center = corner.add_scalar(scale / 2.0);
+
+        let cfg = RenderConfig::<3> {
+            image_size: image_size / workers_per_side,
+            bounds: Bounds {
+                center,
+                size: scale / 2.0,
+            },
+            ..RenderConfig::default()
+        };
+
+        // Special case for the first tile, which can be copied over
+        let (mut depth, norm) = cfg.run(shape.clone())?;
+        for d in &mut depth {
+            if *d > 0 {
+                *d += (z * image_size / workers_per_side) as u32;
+            }
+        }
+        if current_depth.is_empty() {
+            current_depth = depth;
+            current_norm = norm;
+        } else {
+            let mut all = true;
+            for i in 0..depth.len() {
+                if depth[i] > current_depth[i] {
+                    current_depth[i] = depth[i];
+                    current_norm[i] = norm[i];
+                }
+                all &= current_depth[i] == 0;
+            }
+            if all {
+                break;
+            }
+        }
+    }
+    Ok((current_depth, current_norm))
+}
diff --git a/wasm-demo/worker.ts b/wasm-demo/worker.ts
@@ -1,5 +1,6 @@
 import {
   ImageResponse,
+  RenderMode,
   RequestKind,
   ScriptRequest,
   ScriptResponse,
@@ -23,12 +24,36 @@ class Worker {
 
   render(s: ShapeRequest) {
     const shape = fidget.deserialize_tape(s.tape);
-    const out = fidget.render_region(
-      shape,
-      RENDER_SIZE,
-      this.index,
-      WORKERS_PER_SIDE,
-    );
+    let out: Uint8Array;
+    switch (s.mode) {
+      case RenderMode.Bitmap: {
+        out = fidget.render_region_2d(
+          shape,
+          RENDER_SIZE,
+          this.index,
+          WORKERS_PER_SIDE,
+        );
+        break;
+      }
+      case RenderMode.Heightmap: {
+        out = fidget.render_region_heightmap(
+          shape,
+          RENDER_SIZE,
+          this.index,
+          WORKERS_PER_SIDE,
+        );
+        break;
+      }
+      case RenderMode.Normals: {
+        out = fidget.render_region_normals(
+          shape,
+          RENDER_SIZE,
+          this.index,
+          WORKERS_PER_SIDE,
+        );
+        break;
+      }
+    }
     postMessage(new ImageResponse(out), { transfer: [out.buffer] });
   }