From d5824f107a40d290a28deeae5f0fb287e3fe45d9 Mon Sep 17 00:00:00 2001
From: Eduard-Mihai Burtescu <eddyb@lyken.rs>
Date: Sun, 1 Oct 2023 17:12:52 +0300
Subject: [PATCH] qptr/lower: minimally support opaque handle `OpVariable`s.

---
 src/qptr/lower.rs | 48 +++++++++++++++++++++++++++++++++++++----------
 src/spv/spec.rs   |  4 ++++
 2 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/src/qptr/lower.rs b/src/qptr/lower.rs
index 7c674da..4095206 100644
--- a/src/qptr/lower.rs
+++ b/src/qptr/lower.rs
@@ -941,7 +941,7 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
 
                     // HACK(eddyb) this is for `aggregate_load_to_leaf_loads`,
                     // which gets used later, to replace uses of one of the
-                    // outputs ofthe original `OpLoad`, with uses of leaf loads.
+                    // outputs of the original `OpLoad`, with uses of leaf loads.
                     let mut leaf_loads = if spv_inst.opcode == wk.OpLoad {
                         Some(SmallVec::with_capacity(leaf_forms_and_extra_inputs.len()))
                     } else {
@@ -1029,15 +1029,18 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
 
             // FIXME(eddyb) this may need to automatically generate an
             // intermediary `QPtrOp::BufferData` when accessing buffers.
-            let mem_data_layout = match self.lowerer.layout_of(src_pointee_type)? {
-                TypeLayout::Concrete(mem) => mem,
-                _ => {
-                    return Err(LowerError(Diag::bug([
-                        "`OpCopyMemory` of data with non-memory type: ".into(),
-                        src_pointee_type.into(),
-                    ])));
-                }
-            };
+            let mem_data_layout_or_opaque_handle_type =
+                match self.lowerer.layout_of(src_pointee_type)? {
+                    TypeLayout::Concrete(mem) => Ok(mem),
+                    // HACK(eddyb) Rust-GPU generates `OpCopyMemory`s of handles.
+                    TypeLayout::Handle(shapes::Handle::Opaque(ty)) => Err(ty),
+                    _ => {
+                        return Err(LowerError(Diag::bug([
+                            "`OpCopyMemory` of data with non-memory type: ".into(),
+                            src_pointee_type.into(),
+                        ])));
+                    }
+                };
 
             let (dst_ptr, dst_base_offset) = flatten_offsets(dst_ptr);
             let (src_ptr, src_base_offset) = flatten_offsets(src_ptr);
@@ -1089,6 +1092,13 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
                 // be generating, because we don't know ahead of time whether we
                 // even want to expand the `OpCopyMemory`, at all.
                 let mut leaf_offsets_and_types = SmallVec::<[_; 8]>::new();
+                let mem_data_layout = match mem_data_layout_or_opaque_handle_type {
+                    Ok(mem_data_layout) => mem_data_layout,
+                    Err(opaque_handle_type) => {
+                        leaf_offsets_and_types.push((0, opaque_handle_type));
+                        return Some(leaf_offsets_and_types);
+                    }
+                };
                 mem_data_layout
                     .deeply_flatten_if(
                         0,
@@ -1231,6 +1241,7 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
         mut func_at_data_inst: FuncAtMut<'_, DataInst>,
         extra_error: Option<LowerError>,
     ) {
+        let wk = self.lowerer.wk;
         let cx = &self.lowerer.cx;
 
         let func_at_data_inst_frozen = func_at_data_inst.reborrow().freeze();
@@ -1282,6 +1293,23 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
                     continue;
                 }
 
+                // HACK(eddyb) avoid otherwise-unsupported instructions ending up
+                // with invalid address spaces (such cases should be errors,
+                // but Rust-GPU still emits `Generic` everywhere, and having
+                // def-vs-use type mismatches, instead, would also cause issues).
+                let addr_space = match &data_inst_form_def.kind {
+                    DataInstKind::SpvInst(spv_inst, _) => {
+                        if spv_inst.opcode == wk.OpVariable {
+                            AddrSpace::SpvStorageClass(wk.Function)
+                        } else if spv_inst.opcode == wk.OpImageTexelPointer {
+                            AddrSpace::SpvStorageClass(wk.Image)
+                        } else {
+                            addr_space
+                        }
+                    }
+                    _ => addr_space,
+                };
+
                 old_and_new_attrs.get_or_insert_with(get_old_attrs).attrs.insert(
                     QPtrAttr::FromSpvPtrOutput {
                         addr_space: OrdAssertEq(addr_space),
diff --git a/src/spv/spec.rs b/src/spv/spec.rs
index d5dc8f9..40c59de 100644
--- a/src/spv/spec.rs
+++ b/src/spv/spec.rs
@@ -153,6 +153,8 @@ def_well_known! {
 
         OpFunctionCall,
 
+        OpImageTexelPointer,
+
         OpLoad,
         OpStore,
         OpCopyMemory,
@@ -197,6 +199,8 @@ def_well_known! {
         Input,
         Output,
 
+        Image,
+
         IncomingRayPayloadKHR,
         IncomingCallableDataKHR,
         HitAttributeKHR,