From 8e96085caf4f4c2ab812516f607fabb7e908db17 Mon Sep 17 00:00:00 2001
From: Dzmitry Malyshau <kvarkus@gmail.com>
Date: Sat, 13 Feb 2021 21:33:46 -0500
Subject: [PATCH] [spv-in] support image load/store ops

---
 src/back/spv/writer.rs |  9 ++-----
 src/front/spv/error.rs |  2 +-
 src/front/spv/image.rs | 60 +++++++++++++++++++++++++++++++++++++++---
 src/front/spv/mod.rs   | 10 +++++--
 4 files changed, 67 insertions(+), 14 deletions(-)

diff --git a/src/back/spv/writer.rs b/src/back/spv/writer.rs
index f5d7fa1e67..f9b05376f0 100644
--- a/src/back/spv/writer.rs
+++ b/src/back/spv/writer.rs
@@ -1801,13 +1801,8 @@ impl Writer {
                 let reject_id =
                     self.write_expression(ir_module, ir_function, reject, block, function)?;
 
-                let instruction = Instruction::select(
-                    result_type_id,
-                    id,
-                    condition_id,
-                    accept_id,
-                    reject_id,
-                );
+                let instruction =
+                    Instruction::select(result_type_id, id, condition_id, accept_id, reject_id);
                 block.body.push(instruction);
                 RawExpression::Value(id)
             }
diff --git a/src/front/spv/error.rs b/src/front/spv/error.rs
index 194c279b8d..c7bac4b176 100644
--- a/src/front/spv/error.rs
+++ b/src/front/spv/error.rs
@@ -40,7 +40,7 @@ pub enum Error {
     InvalidBinding(spirv::Word),
     InvalidGlobalVar(crate::Expression),
     InvalidImageBaseType(Handle<crate::Type>),
-    InvalidSampleImage(Handle<crate::Type>),
+    InvalidImage(Handle<crate::Type>),
     InvalidDepthReference(Handle<crate::Type>),
     InvalidAsType(Handle<crate::Type>),
     InvalidVectorType(Handle<crate::Type>),
diff --git a/src/front/spv/image.rs b/src/front/spv/image.rs
index 3453e7048a..6a01e7df36 100644
--- a/src/front/spv/image.rs
+++ b/src/front/spv/image.rs
@@ -182,7 +182,59 @@ impl<I: Iterator<Item = u32>> super::Parser<I> {
         Ok(())
     }
 
-    pub(super) fn parse_image_fetch(
+    pub(super) fn parse_image_write(
+        &mut self,
+        words_left: u16,
+        type_arena: &Arena<crate::Type>,
+        global_arena: &Arena<crate::GlobalVariable>,
+        expressions: &mut Arena<crate::Expression>,
+    ) -> Result<crate::Statement, Error> {
+        let image_id = self.next()?;
+        let coordinate_id = self.next()?;
+        let value_id = self.next()?;
+
+        if words_left != 0 {
+            let image_ops = self.next()?;
+            let other = spirv::ImageOperands::from_bits_truncate(image_ops);
+            log::warn!("Skipping {:?}", other);
+            for _ in 1..words_left {
+                self.next()?;
+            }
+        }
+
+        let image_lexp = self.lookup_expression.lookup(image_id)?;
+        let image_var_handle = expressions[image_lexp.handle].as_global_var()?;
+        let image_var = &global_arena[image_var_handle];
+
+        let coord_lexp = self.lookup_expression.lookup(coordinate_id)?;
+        let coord_type_handle = self.lookup_type.lookup(coord_lexp.type_id)?.handle;
+        let (coordinate, array_index) = match type_arena[image_var.ty].inner {
+            crate::TypeInner::Image {
+                dim,
+                arrayed,
+                class: _,
+            } => extract_image_coordinates(
+                dim,
+                arrayed,
+                coord_lexp.handle,
+                coord_type_handle,
+                type_arena,
+                expressions,
+            ),
+            _ => return Err(Error::InvalidImage(image_var.ty)),
+        };
+
+        let value_lexp = self.lookup_expression.lookup(value_id)?;
+
+        Ok(crate::Statement::ImageStore {
+            image: image_lexp.handle,
+            coordinate,
+            array_index,
+            value: value_lexp.handle,
+        })
+    }
+
+    pub(super) fn parse_image_load(
         &mut self,
         mut words_left: u16,
         type_arena: &Arena<crate::Type>,
@@ -240,7 +292,7 @@ impl<I: Iterator<Item = u32>> super::Parser<I> {
                 type_arena,
                 expressions,
             ),
-            _ => return Err(Error::InvalidSampleImage(image_var.ty)),
+            _ => return Err(Error::InvalidImage(image_var.ty)),
         };
 
         let expr = crate::Expression::ImageLoad {
@@ -328,7 +380,7 @@ impl<I: Iterator<Item = u32>> super::Parser<I> {
                 type_arena,
                 expressions,
             ),
-            _ => return Err(Error::InvalidSampleImage(image_var.ty)),
+            _ => return Err(Error::InvalidImage(image_var.ty)),
         };
 
         let expr = crate::Expression::ImageSample {
@@ -429,7 +481,7 @@ impl<I: Iterator<Item = u32>> super::Parser<I> {
                 type_arena,
                 expressions,
             ),
-            _ => return Err(Error::InvalidSampleImage(image_var.ty)),
+            _ => return Err(Error::InvalidImage(image_var.ty)),
         };
 
         let expr = crate::Expression::ImageSample {
diff --git a/src/front/spv/mod.rs b/src/front/spv/mod.rs
index 3de5bebf4f..2bcad16251 100644
--- a/src/front/spv/mod.rs
+++ b/src/front/spv/mod.rs
@@ -1100,9 +1100,15 @@ impl<I: Iterator<Item = u32>> Parser<I> {
                     inst.expect(5)?;
                     self.parse_image_couple()?;
                 }
-                Op::ImageFetch => {
+                Op::ImageWrite => {
+                    let extra = inst.expect_at_least(4)?;
+                    let stmt =
+                        self.parse_image_write(extra, type_arena, global_arena, expressions)?;
+                    block.push(stmt);
+                }
+                Op::ImageFetch | Op::ImageRead => {
                     let extra = inst.expect_at_least(5)?;
-                    self.parse_image_fetch(extra, type_arena, global_arena, expressions)?;
+                    self.parse_image_load(extra, type_arena, global_arena, expressions)?;
                 }
                 Op::ImageSampleImplicitLod
                 | Op::ImageSampleExplicitLod