From cc8c2ebec8f160c5ad13754f08018547545a7231 Mon Sep 17 00:00:00 2001
From: charles-r-earp <charles.r.earp@gmail.com>
Date: Fri, 27 Nov 2020 15:46:57 -0800
Subject: [PATCH 1/6] Added block decoration. StorageBuffer and PushConstant
 now point to a Block<T>.

---
 crates/rustc_codegen_spirv/src/abi.rs         | 20 ++++++++++
 .../src/codegen_cx/constant.rs                |  1 +
 .../src/codegen_cx/type_.rs                   |  1 +
 crates/rustc_codegen_spirv/src/spirv_type.rs  |  8 ++++
 crates/rustc_codegen_spirv/src/symbols.rs     |  3 ++
 crates/spirv-std/src/lib.rs                   | 39 ++++++++++++++++++-
 6 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/crates/rustc_codegen_spirv/src/abi.rs b/crates/rustc_codegen_spirv/src/abi.rs
index 3c8ef3b34e..4a08a05fbc 100644
--- a/crates/rustc_codegen_spirv/src/abi.rs
+++ b/crates/rustc_codegen_spirv/src/abi.rs
@@ -227,6 +227,7 @@ impl<'tcx> ConvSpirvType<'tcx> for CastTarget {
             field_types: args,
             field_offsets,
             field_names: None,
+            is_block: false
         }
         .def(cx)
     }
@@ -311,6 +312,7 @@ fn trans_type_impl<'tcx>(cx: &CodegenCx<'tcx>, ty: TyAndLayout<'tcx>, is_immedia
             field_types: Vec::new(),
             field_offsets: Vec::new(),
             field_names: None,
+            is_block: false,
         }
         .def(cx),
         Abi::Scalar(ref scalar) => trans_scalar(cx, ty, scalar, None, is_immediate),
@@ -330,6 +332,7 @@ fn trans_type_impl<'tcx>(cx: &CodegenCx<'tcx>, ty: TyAndLayout<'tcx>, is_immedia
                 field_types: vec![one_spirv, two_spirv],
                 field_offsets: vec![one_offset, two_offset],
                 field_names: None,
+                is_block: false,
             }
             .def(cx)
         }
@@ -547,6 +550,20 @@ fn get_storage_class<'tcx>(cx: &CodegenCx<'tcx>, ty: TyAndLayout<'tcx>) -> Optio
     None
 }
 
+/// Handles `#[spirv(block)]`. Note this is only called in the scalar translation code, because this is only
+/// used for spooky builtin stuff, and we pinky promise to never have more than one pointer field in one of these.
+// TODO: Enforce this is only used in spirv-std.
+fn get_is_block_decorated<'tcx>(cx: &CodegenCx<'tcx>, ty: TyAndLayout<'tcx>) -> bool {
+    if let TyKind::Adt(adt, _substs) = ty.ty.kind() {
+        for attr in parse_attrs(cx, cx.tcx.get_attrs(adt.did)) {
+            if let SpirvAttribute::Block = attr {
+                return true;
+            }
+        }
+    }
+    false
+}
+
 fn trans_aggregate<'tcx>(cx: &CodegenCx<'tcx>, ty: TyAndLayout<'tcx>) -> Word {
     match ty.fields {
         FieldsShape::Primitive => cx.tcx.sess.fatal(&format!(
@@ -583,6 +600,7 @@ fn trans_aggregate<'tcx>(cx: &CodegenCx<'tcx>, ty: TyAndLayout<'tcx>) -> Word {
                     field_types: Vec::new(),
                     field_offsets: Vec::new(),
                     field_names: None,
+                    is_block: false,
                 }
                 .def(cx)
             } else {
@@ -676,6 +694,7 @@ fn trans_struct<'tcx>(cx: &CodegenCx<'tcx>, ty: TyAndLayout<'tcx>) -> Word {
             }
         };
     }
+    let is_block = get_is_block_decorated(cx, ty);
     SpirvType::Adt {
         name,
         size,
@@ -683,6 +702,7 @@ fn trans_struct<'tcx>(cx: &CodegenCx<'tcx>, ty: TyAndLayout<'tcx>) -> Word {
         field_types,
         field_offsets,
         field_names: Some(field_names),
+        is_block 
     }
     .def(cx)
 }
diff --git a/crates/rustc_codegen_spirv/src/codegen_cx/constant.rs b/crates/rustc_codegen_spirv/src/codegen_cx/constant.rs
index a538266752..14854f157b 100644
--- a/crates/rustc_codegen_spirv/src/codegen_cx/constant.rs
+++ b/crates/rustc_codegen_spirv/src/codegen_cx/constant.rs
@@ -183,6 +183,7 @@ impl<'tcx> ConstMethods<'tcx> for CodegenCx<'tcx> {
             field_types,
             field_offsets,
             field_names: None,
+            is_block: false,
         }
         .def(self);
         self.constant_composite(struct_ty, elts.iter().map(|f| f.def_cx(self)).collect())
diff --git a/crates/rustc_codegen_spirv/src/codegen_cx/type_.rs b/crates/rustc_codegen_spirv/src/codegen_cx/type_.rs
index ba59b55839..40248ab588 100644
--- a/crates/rustc_codegen_spirv/src/codegen_cx/type_.rs
+++ b/crates/rustc_codegen_spirv/src/codegen_cx/type_.rs
@@ -150,6 +150,7 @@ impl<'tcx> BaseTypeMethods<'tcx> for CodegenCx<'tcx> {
             field_types: els.to_vec(),
             field_offsets,
             field_names: None,
+            is_block: false,
         }
         .def(self)
     }
diff --git a/crates/rustc_codegen_spirv/src/spirv_type.rs b/crates/rustc_codegen_spirv/src/spirv_type.rs
index e0294d4026..e4b06dd802 100644
--- a/crates/rustc_codegen_spirv/src/spirv_type.rs
+++ b/crates/rustc_codegen_spirv/src/spirv_type.rs
@@ -31,6 +31,7 @@ pub enum SpirvType {
         field_types: Vec<Word>,
         field_offsets: Vec<Size>,
         field_names: Option<Vec<String>>,
+        is_block: bool
     },
     Opaque {
         name: String,
@@ -112,6 +113,7 @@ impl SpirvType {
                 ref field_types,
                 ref field_offsets,
                 ref field_names,
+                is_block,
             } => {
                 let mut emit = cx.emit_global();
                 // Ensure a unique struct is emitted each time, due to possibly having different OpMemberDecorates
@@ -132,6 +134,9 @@ impl SpirvType {
                         );
                     }
                 }
+                if is_block {
+                    emit.decorate(id, Decoration::Block, None);
+                }
                 if let Some(field_names) = field_names {
                     for (index, field_name) in field_names.iter().enumerate() {
                         emit.member_name(result, index as u32, field_name);
@@ -313,6 +318,7 @@ impl fmt::Debug for SpirvTypePrinter<'_, '_> {
                 ref field_types,
                 ref field_offsets,
                 ref field_names,
+                is_block
             } => {
                 let fields = field_types
                     .iter()
@@ -326,6 +332,7 @@ impl fmt::Debug for SpirvTypePrinter<'_, '_> {
                     .field("field_types", &fields)
                     .field("field_offsets", field_offsets)
                     .field("field_names", field_names)
+                    .field("is_block", &is_block)
                     .finish()
             }
             SpirvType::Opaque { ref name } => f
@@ -432,6 +439,7 @@ impl SpirvTypePrinter<'_, '_> {
                 ref field_types,
                 field_offsets: _,
                 ref field_names,
+                is_block: _,
             } => {
                 write!(f, "struct {} {{ ", name)?;
                 for (index, &field) in field_types.iter().enumerate() {
diff --git a/crates/rustc_codegen_spirv/src/symbols.rs b/crates/rustc_codegen_spirv/src/symbols.rs
index 7b81336052..496408f1a9 100644
--- a/crates/rustc_codegen_spirv/src/symbols.rs
+++ b/crates/rustc_codegen_spirv/src/symbols.rs
@@ -4,6 +4,7 @@ use rspirv::spirv::{BuiltIn, ExecutionMode, ExecutionModel, StorageClass};
 use rustc_ast::ast::{AttrKind, Attribute, Lit, LitIntType, LitKind, NestedMetaItem};
 use rustc_span::symbol::{Ident, Symbol};
 use std::collections::HashMap;
+use std::iter::once;
 
 /// Various places in the codebase (mostly attribute parsing) need to compare rustc Symbols to particular keywords.
 /// Symbols are interned, as in, they don't actually store the string itself inside them, but rather an index into a
@@ -324,6 +325,7 @@ impl Symbols {
         let attributes_iter = builtins
             .chain(storage_classes)
             .chain(execution_models)
+            .chain(once(("block", SpirvAttribute::Block)))
             .map(|(a, b)| (Symbol::intern(a), b));
         let mut attributes = HashMap::new();
         for (a, b) in attributes_iter {
@@ -411,6 +413,7 @@ impl From<ExecutionModel> for Entry {
 pub enum SpirvAttribute {
     Builtin(BuiltIn),
     StorageClass(StorageClass),
+    Block,
     Entry(Entry),
     DescriptorSet(u32),
     Binding(u32),
diff --git a/crates/spirv-std/src/lib.rs b/crates/spirv-std/src/lib.rs
index d70e3a64ad..f0660147bb 100644
--- a/crates/spirv-std/src/lib.rs
+++ b/crates/spirv-std/src/lib.rs
@@ -42,8 +42,16 @@ pub use math_ext::MathExt;
 
 pub use glam;
 
+#[doc(hidden)]
+#[allow(unused_attributes)]
+#[spirv(block)]
+#[repr(transparent)]
+#[derive(Clone, Copy)]
+pub struct Block<T>(T); 
+
 macro_rules! pointer_addrspace_write {
     (false) => {};
+    (false, block) => {};
     (true) => {
         #[inline]
         #[allow(unused_attributes)]
@@ -52,6 +60,14 @@ macro_rules! pointer_addrspace_write {
             *self.x = v
         }
     };
+    (true, block) => {
+        #[inline]
+        #[allow(unused_attributes)]
+        #[spirv(really_unsafe_ignore_bitcasts)]
+        pub fn store(&mut self, v: T) {
+            *self.x = Block(v);
+        }
+    };
 }
 
 macro_rules! pointer_addrspace {
@@ -73,6 +89,25 @@ macro_rules! pointer_addrspace {
             pointer_addrspace_write!($writeable);
         }
     };
+    ($storage_class:ident, $type_name:ident, $writeable:tt, block) => {
+        #[allow(unused_attributes)]
+        #[spirv($storage_class)]
+        pub struct $type_name<'a, T> {
+            x: &'a mut Block<T>,
+        }
+
+        impl<'a, T: Copy> $type_name<'a, T> {
+            #[inline]
+            #[allow(unused_attributes)]
+            #[spirv(really_unsafe_ignore_bitcasts)]
+            pub fn load(&self) -> T {
+                let block = *self.x;
+                block.0
+            }
+
+            pointer_addrspace_write!($writeable, block);
+        }
+    };
 }
 
 // Make sure these strings stay synced with symbols.rs
@@ -86,10 +121,10 @@ pointer_addrspace!(cross_workgroup, CrossWorkgroup, true);
 pointer_addrspace!(private, Private, true);
 pointer_addrspace!(function, Function, true);
 pointer_addrspace!(generic, Generic, true);
-pointer_addrspace!(push_constant, PushConstant, false);
+pointer_addrspace!(push_constant, PushConstant, false, block);
 pointer_addrspace!(atomic_counter, AtomicCounter, true);
 pointer_addrspace!(image, Image, true);
-pointer_addrspace!(storage_buffer, StorageBuffer, true);
+pointer_addrspace!(storage_buffer, StorageBuffer, true, block);
 pointer_addrspace!(callable_data_khr, CallableDataKHR, true);
 pointer_addrspace!(incoming_callable_data_khr, IncomingCallableDataKHR, true);
 pointer_addrspace!(ray_payload_khr, RayPayloadKHR, true);

From 3574288415230bbc6fa6c88d2da3bdde78b72e7d Mon Sep 17 00:00:00 2001
From: charles-r-earp <charles.r.earp@gmail.com>
Date: Fri, 27 Nov 2020 15:47:24 -0800
Subject: [PATCH 2/6] fmt

---
 crates/rustc_codegen_spirv/src/abi.rs        | 4 ++--
 crates/rustc_codegen_spirv/src/spirv_type.rs | 4 ++--
 crates/spirv-std/src/lib.rs                  | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/crates/rustc_codegen_spirv/src/abi.rs b/crates/rustc_codegen_spirv/src/abi.rs
index 4a08a05fbc..469670095b 100644
--- a/crates/rustc_codegen_spirv/src/abi.rs
+++ b/crates/rustc_codegen_spirv/src/abi.rs
@@ -227,7 +227,7 @@ impl<'tcx> ConvSpirvType<'tcx> for CastTarget {
             field_types: args,
             field_offsets,
             field_names: None,
-            is_block: false
+            is_block: false,
         }
         .def(cx)
     }
@@ -702,7 +702,7 @@ fn trans_struct<'tcx>(cx: &CodegenCx<'tcx>, ty: TyAndLayout<'tcx>) -> Word {
         field_types,
         field_offsets,
         field_names: Some(field_names),
-        is_block 
+        is_block,
     }
     .def(cx)
 }
diff --git a/crates/rustc_codegen_spirv/src/spirv_type.rs b/crates/rustc_codegen_spirv/src/spirv_type.rs
index e4b06dd802..4e5311352b 100644
--- a/crates/rustc_codegen_spirv/src/spirv_type.rs
+++ b/crates/rustc_codegen_spirv/src/spirv_type.rs
@@ -31,7 +31,7 @@ pub enum SpirvType {
         field_types: Vec<Word>,
         field_offsets: Vec<Size>,
         field_names: Option<Vec<String>>,
-        is_block: bool
+        is_block: bool,
     },
     Opaque {
         name: String,
@@ -318,7 +318,7 @@ impl fmt::Debug for SpirvTypePrinter<'_, '_> {
                 ref field_types,
                 ref field_offsets,
                 ref field_names,
-                is_block
+                is_block,
             } => {
                 let fields = field_types
                     .iter()
diff --git a/crates/spirv-std/src/lib.rs b/crates/spirv-std/src/lib.rs
index f0660147bb..2bc66e4b7a 100644
--- a/crates/spirv-std/src/lib.rs
+++ b/crates/spirv-std/src/lib.rs
@@ -47,7 +47,7 @@ pub use glam;
 #[spirv(block)]
 #[repr(transparent)]
 #[derive(Clone, Copy)]
-pub struct Block<T>(T); 
+pub struct Block<T>(T);
 
 macro_rules! pointer_addrspace_write {
     (false) => {};

From 7b3020ed7d12f98f8d8a07cf1849be23414c95d8 Mon Sep 17 00:00:00 2001
From: charles-r-earp <charles.r.earp@gmail.com>
Date: Fri, 27 Nov 2020 23:51:42 -0800
Subject: [PATCH 3/6] Merged with upstream updates. Matched storage classes in
 pointercast.

---
 .../src/builder/builder_methods.rs            | 20 +++++++-
 crates/spirv-std/src/storage_class.rs         | 51 ++++++++++++++++++-
 2 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/crates/rustc_codegen_spirv/src/builder/builder_methods.rs b/crates/rustc_codegen_spirv/src/builder/builder_methods.rs
index 6fa76bcb83..d692acd5d9 100644
--- a/crates/rustc_codegen_spirv/src/builder/builder_methods.rs
+++ b/crates/rustc_codegen_spirv/src/builder/builder_methods.rs
@@ -1220,8 +1220,11 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
     }
 
     fn pointercast(&mut self, val: Self::Value, dest_ty: Self::Type) -> Self::Value {
-        let val_pointee = match self.lookup_type(val.ty) {
-            SpirvType::Pointer { pointee, .. } => pointee,
+        let (storage_class, val_pointee) = match self.lookup_type(val.ty) {
+            SpirvType::Pointer {
+                storage_class,
+                pointee,
+            } => (storage_class, pointee),
             other => self.fatal(&format!(
                 "pointercast called on non-pointer source type: {:?}",
                 other
@@ -1237,6 +1240,19 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
         if val.ty == dest_ty {
             val
         } else if let Some(indices) = self.try_pointercast_via_gep(val_pointee, dest_pointee) {
+            let dest_ty = if self
+                .really_unsafe_ignore_bitcasts
+                .borrow()
+                .contains(&self.current_fn)
+            {
+                SpirvType::Pointer {
+                    storage_class,
+                    pointee: dest_pointee,
+                }
+                .def(self)
+            } else {
+                dest_ty
+            };
             let indices = indices
                 .into_iter()
                 .map(|idx| self.constant_u32(idx).def(self))
diff --git a/crates/spirv-std/src/storage_class.rs b/crates/spirv-std/src/storage_class.rs
index 077d310a31..7ab2d394fb 100644
--- a/crates/spirv-std/src/storage_class.rs
+++ b/crates/spirv-std/src/storage_class.rs
@@ -48,6 +48,53 @@ macro_rules! storage_class {
         }
     };
 
+    // Interior Block
+    ($(#[$($meta:meta)+])* block $block:ident storage_class $name:ident ; $($tt:tt)*) => {
+
+        $(#[$($meta)+])*
+        #[spirv(block)]
+        #[allow(unused_attributes)]
+        pub struct $block <T>(T);
+
+        $(#[$($meta)+])*
+        #[allow(unused_attributes)]
+        pub struct $name<'value, T> {
+            value: &'value mut $block <T>,
+        }
+
+        impl<T: Copy> $name<'_, T> {
+            /// Load the value into memory.
+            #[inline]
+            #[allow(unused_attributes)]
+            #[spirv(really_unsafe_ignore_bitcasts)]
+            pub fn load(&self) -> T {
+                self.value.0
+            }
+        }
+
+        storage_class!($($tt)*);
+    };
+
+    // Methods available on writeable storage classes.
+    ($(#[$($meta:meta)+])* writeable block $block:ident storage_class $name:ident $($tt:tt)+) => {
+        storage_class!($(#[$($meta)+])* block $block storage_class $name $($tt)+);
+
+        impl <T: Copy> $name<'_, T> {
+            /// Store the value in storage.
+            #[inline]
+            #[allow(unused_attributes)]
+            #[spirv(really_unsafe_ignore_bitcasts)]
+            pub fn store(&mut self, v: T) {
+                self.value.0 = v
+            }
+
+            /// A convenience function to load a value into memory and store it.
+            pub fn then(&mut self, then: impl FnOnce(T) -> T) {
+                self.store((then)(self.load()));
+            }
+        }
+    };
+
     (;) => {};
     () => {};
 }
@@ -112,7 +159,7 @@ storage_class! {
     /// Intended to contain a small bank of values pushed from the client API.
     /// Variables declared with this storage class are read-only, and must not
     /// have initializers.
-    #[spirv(push_constant)] storage_class PushConstant;
+    #[spirv(push_constant)] block PushConstantBlock storage_class PushConstant;
 
     /// Atomic counter-specific memory.
     ///
@@ -131,7 +178,7 @@ storage_class! {
     ///
     /// Shared externally, readable and writable, visible across all functions
     /// in all invocations in all work groups.
-    #[spirv(storage_buffer)] writeable storage_class StorageBuffer;
+    #[spirv(storage_buffer)] writeable block StorageBufferBlock storage_class StorageBuffer;
 
     /// Used for storing arbitrary data associated with a ray to pass
     /// to callables. (Requires `SPV_KHR_ray_tracing` extension)

From 0d3ec6d2d6a8362d1355dea327d30b5b7a0bbde8 Mon Sep 17 00:00:00 2001
From: charles-r-earp <charles.r.earp@gmail.com>
Date: Mon, 30 Nov 2020 13:38:38 -0800
Subject: [PATCH 4/6] Cleanup of unused imports. Removed meta attributes on
 Block.

---
 crates/rustc_codegen_spirv/src/symbols.rs | 1 -
 crates/spirv-std/src/storage_class.rs     | 1 -
 2 files changed, 2 deletions(-)

diff --git a/crates/rustc_codegen_spirv/src/symbols.rs b/crates/rustc_codegen_spirv/src/symbols.rs
index beec3d38f2..63af6518a8 100644
--- a/crates/rustc_codegen_spirv/src/symbols.rs
+++ b/crates/rustc_codegen_spirv/src/symbols.rs
@@ -6,7 +6,6 @@ use rspirv::spirv::{
 use rustc_ast::ast::{AttrKind, Attribute, Lit, LitIntType, LitKind, NestedMetaItem};
 use rustc_span::symbol::{Ident, Symbol};
 use std::collections::HashMap;
-use std::iter::once;
 
 /// Various places in the codebase (mostly attribute parsing) need to compare rustc Symbols to particular keywords.
 /// Symbols are interned, as in, they don't actually store the string itself inside them, but rather an index into a
diff --git a/crates/spirv-std/src/storage_class.rs b/crates/spirv-std/src/storage_class.rs
index 7ab2d394fb..8421d8d790 100644
--- a/crates/spirv-std/src/storage_class.rs
+++ b/crates/spirv-std/src/storage_class.rs
@@ -51,7 +51,6 @@ macro_rules! storage_class {
     // Interior Block
     ($(#[$($meta:meta)+])* block $block:ident storage_class $name:ident ; $($tt:tt)*) => {
 
-        $(#[$($meta)+])*
         #[spirv(block)]
         #[allow(unused_attributes)]
         pub struct $block <T>(T);

From 127200f5c0c01ac6e850fee844090dcb834d0cbb Mon Sep 17 00:00:00 2001
From: charles-r-earp <charles.r.earp@gmail.com>
Date: Mon, 30 Nov 2020 14:32:52 -0800
Subject: [PATCH 5/6] Renamed fields of block storage classes self.value.0 to
 self.block.value

---
 crates/spirv-std/src/storage_class.rs | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/crates/spirv-std/src/storage_class.rs b/crates/spirv-std/src/storage_class.rs
index 8421d8d790..6efb31660d 100644
--- a/crates/spirv-std/src/storage_class.rs
+++ b/crates/spirv-std/src/storage_class.rs
@@ -53,12 +53,14 @@ macro_rules! storage_class {
 
         #[spirv(block)]
         #[allow(unused_attributes)]
-        pub struct $block <T>(T);
+        pub struct $block <T> {
+            value: T
+        }
 
         $(#[$($meta)+])*
         #[allow(unused_attributes)]
-        pub struct $name<'value, T> {
-            value: &'value mut $block <T>,
+        pub struct $name<'block, T> {
+            block: &'block mut $block <T>,
         }
 
         impl<T: Copy> $name<'_, T> {
@@ -67,7 +69,7 @@ macro_rules! storage_class {
             #[allow(unused_attributes)]
             #[spirv(really_unsafe_ignore_bitcasts)]
             pub fn load(&self) -> T {
-                self.value.0
+                self.block.value
             }
         }
 
@@ -84,9 +86,9 @@ macro_rules! storage_class {
             #[allow(unused_attributes)]
             #[spirv(really_unsafe_ignore_bitcasts)]
             pub fn store(&mut self, v: T) {
-                self.value.0 = v
+                self.block.value = v;
             }
-
+            
             /// A convenience function to load a value into memory and store it.
             pub fn then(&mut self, then: impl FnOnce(T) -> T) {
                 self.store((then)(self.load()));

From 512a2cf31588baf124e85b107c6f151aacb263b3 Mon Sep 17 00:00:00 2001
From: charles-r-earp <charles.r.earp@gmail.com>
Date: Mon, 30 Nov 2020 15:53:25 -0800
Subject: [PATCH 6/6] fmt

---
 crates/spirv-std/src/storage_class.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/spirv-std/src/storage_class.rs b/crates/spirv-std/src/storage_class.rs
index 6efb31660d..4fc5740978 100644
--- a/crates/spirv-std/src/storage_class.rs
+++ b/crates/spirv-std/src/storage_class.rs
@@ -88,7 +88,7 @@ macro_rules! storage_class {
             pub fn store(&mut self, v: T) {
                 self.block.value = v;
             }
-            
+
             /// A convenience function to load a value into memory and store it.
             pub fn then(&mut self, then: impl FnOnce(T) -> T) {
                 self.store((then)(self.load()));