From 1dea6c0ef07e68c174bd742793a773532c75bb02 Mon Sep 17 00:00:00 2001
From: Arman Uguray <armansito@google.com>
Date: Tue, 27 Jun 2023 23:21:05 -0700
Subject: [PATCH] Fix invalid buffer access errors caught by shader validation

Fixed several other shader validation errors caught when running
vello_shaders natively on Metal.

These were primarily caused by reading an invalid drawtag while
accessing the scene buffer. Scene buffer access in the offending
pipelines now initialize the draw tag to DRAWTAG_NOP if an
invocation ID would land beyond the valid index range of encoded draw
objects.
---
 .vscode/settings.json   |  1 +
 shader/binning.wgsl     |  5 ++++-
 shader/draw_leaf.wgsl   |  4 +++-
 shader/draw_reduce.wgsl |  4 +++-
 shader/shared/util.wgsl | 23 +++++++++++++++++++++++
 src/shaders.rs          |  1 +
 6 files changed, 35 insertions(+), 3 deletions(-)
 create mode 100644 shader/shared/util.wgsl

diff --git a/.vscode/settings.json b/.vscode/settings.json
index c7bfa32..201cb77 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -12,6 +12,7 @@
     "segment": "${workspaceFolder}/shader/shared/segment.wgsl",
     "tile": "${workspaceFolder}/shader/shared/tile.wgsl",
     "transform": "${workspaceFolder}/shader/shared/transform.wgsl"
+    "util": "${workspaceFolder}/shader/shared/util.wgsl",
   },
   "wgsl-analyzer.diagnostics.nagaVersion": "main",
   "wgsl-analyzer.preprocessor.shaderDefs": [
diff --git a/shader/binning.wgsl b/shader/binning.wgsl
index 77ab462..2672bb8 100644
--- a/shader/binning.wgsl
+++ b/shader/binning.wgsl
@@ -74,7 +74,10 @@ fn main(
         let draw_monoid = draw_monoids[element_ix];
         var clip_bbox = vec4(-1e9, -1e9, 1e9, 1e9);
         if draw_monoid.clip_ix > 0u {
-            clip_bbox = clip_bbox_buf[draw_monoid.clip_ix - 1u];
+            // TODO: `clip_ix` should always be valid as long as the monoids are correct. Leaving
+            // the bounds check in here for correctness but we should assert this condition instead
+            // once there is a debug-assertion mechanism.
+            clip_bbox = clip_bbox_buf[min(draw_monoid.clip_ix - 1u, config.n_clip - 1u)];
         }
         // For clip elements, clip_box is the bbox of the clip path,
         // intersected with enclosing clips.
diff --git a/shader/draw_leaf.wgsl b/shader/draw_leaf.wgsl
index 3879afc..6154b92 100644
--- a/shader/draw_leaf.wgsl
+++ b/shader/draw_leaf.wgsl
@@ -29,6 +29,8 @@ var<storage, read_write> info: array<u32>;
 @group(0) @binding(6)
 var<storage, read_write> clip_inp: array<ClipInp>;
 
+#import util
+
 let WG_SIZE = 256u;
 
 fn read_transform(transform_base: u32, ix: u32) -> Transform {
@@ -73,7 +75,7 @@ fn main(
     workgroupBarrier();
     var m = sh_scratch[0];
     workgroupBarrier();
-    let tag_word = scene[config.drawtag_base + ix];
+    let tag_word = read_draw_tag_from_scene(ix);
     agg = map_draw_tag(tag_word);
     sh_scratch[local_id.x] = agg;
     for (var i = 0u; i < firstTrailingBit(WG_SIZE); i += 1u) {
diff --git a/shader/draw_reduce.wgsl b/shader/draw_reduce.wgsl
index af17d78..051d8f8 100644
--- a/shader/draw_reduce.wgsl
+++ b/shader/draw_reduce.wgsl
@@ -16,13 +16,15 @@ let WG_SIZE = 256u;
 
 var<workgroup> sh_scratch: array<DrawMonoid, WG_SIZE>;
 
+#import util
+
 @compute @workgroup_size(256)
 fn main(
     @builtin(global_invocation_id) global_id: vec3<u32>,
     @builtin(local_invocation_id) local_id: vec3<u32>,
 ) {
     let ix = global_id.x;
-    let tag_word = scene[config.drawtag_base + ix];
+    let tag_word = read_draw_tag_from_scene(ix);
     var agg = map_draw_tag(tag_word);
     sh_scratch[local_id.x] = agg;
     for (var i = 0u; i < firstTrailingBit(WG_SIZE); i += 1u) {
diff --git a/shader/shared/util.wgsl b/shader/shared/util.wgsl
new file mode 100644
index 0000000..1e40fc4
--- /dev/null
+++ b/shader/shared/util.wgsl
@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: Apache-2.0 OR MIT OR Unlicense
+
+// This file defines utility functions that interact with host-shareable buffer objects. It should
+// be imported once following the resource binding declarations in the shader module that access
+// them.
+
+// Reads a draw tag from the scene buffer, defaulting to DRAWTAG_NOP if the given `ix` is beyond the
+// range of valid draw objects (e.g this can happen if `ix` is derived from an invocation ID in a
+// workgroup that partially spans valid range).
+//
+// This function depends on the following global declarations:
+//    * `scene`: array<u32>
+//    * `config`: Config (see config.wgsl)
+fn read_draw_tag_from_scene(ix: u32) -> u32 {
+    let tag_ix = config.drawtag_base + ix;
+    var tag_word: u32;
+    if tag_ix < config.drawtag_base + config.n_drawobj {
+        tag_word = scene[tag_ix];
+    } else {
+        tag_word = DRAWTAG_NOP;
+    }
+    return tag_word;
+}
diff --git a/src/shaders.rs b/src/shaders.rs
index 8da1807..1083374 100644
--- a/src/shaders.rs
+++ b/src/shaders.rs
@@ -319,4 +319,5 @@ const SHARED_SHADERS: &[(&str, &str)] = &[
     shared_shader!("segment"),
     shared_shader!("tile"),
     shared_shader!("transform"),
+    shared_shader!("util"),
 ];