From 405413452d9b6cdcd0865a6fc5f866a0deef7cf3 Mon Sep 17 00:00:00 2001
From: teoxoy <28601907+teoxoy@users.noreply.github.com>
Date: Thu, 11 May 2023 21:39:49 +0200
Subject: [PATCH] Correct type of index and flush the emitter

---
 src/front/spv/error.rs |  2 ++
 src/front/spv/mod.rs   | 58 ++++++++++++++++++++++++++++++++++--------
 2 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/src/front/spv/error.rs b/src/front/spv/error.rs
index 6c9cf384f1..2f9bf2d1bc 100644
--- a/src/front/spv/error.rs
+++ b/src/front/spv/error.rs
@@ -78,6 +78,8 @@ pub enum Error {
     InvalidAccess(crate::Expression),
     #[error("invalid access index %{0}")]
     InvalidAccessIndex(spirv::Word),
+    #[error("invalid index type %{0}")]
+    InvalidIndexType(spirv::Word),
     #[error("invalid binding %{0}")]
     InvalidBinding(spirv::Word),
     #[error("invalid global var {0:?}")]
diff --git a/src/front/spv/mod.rs b/src/front/spv/mod.rs
index 14bae2a9cd..65423db03e 100644
--- a/src/front/spv/mod.rs
+++ b/src/front/spv/mod.rs
@@ -1678,18 +1678,26 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                     let root_type_lookup = self.lookup_type.lookup(root_lexp.type_id)?;
                     let index_lexp = self.lookup_expression.lookup(index_id)?;
                     let index_handle = get_expr_handle!(index_id, index_lexp);
+                    let index_type = self.lookup_type.lookup(index_lexp.type_id)?.handle;
 
                     let num_components = match ctx.type_arena[root_type_lookup.handle].inner {
                         crate::TypeInner::Vector { size, .. } => size as u32,
                         _ => return Err(Error::InvalidVectorType(root_type_lookup.handle)),
                     };
 
-                    let make_index = |ctx: &mut BlockContext, index: u32| {
-                        ctx.expressions
-                            .append(crate::Expression::Literal(crate::Literal::U32(index)), span)
+                    let mut make_index = |ctx: &mut BlockContext, index: u32| {
+                        make_index_literal(
+                            ctx,
+                            index,
+                            &mut block,
+                            &mut emitter,
+                            index_type,
+                            index_lexp.type_id,
+                            span,
+                        )
                     };
 
-                    let index_expr = make_index(ctx, 0);
+                    let index_expr = make_index(ctx, 0)?;
                     let mut handle = ctx.expressions.append(
                         crate::Expression::Access {
                             base: root_handle,
@@ -1698,7 +1706,7 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                         span,
                     );
                     for index in 1..num_components {
-                        let index_expr = make_index(ctx, index);
+                        let index_expr = make_index(ctx, index)?;
                         let access_expr = ctx.expressions.append(
                             crate::Expression::Access {
                                 base: root_handle,
@@ -1749,20 +1757,24 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                     let root_type_lookup = self.lookup_type.lookup(root_lexp.type_id)?;
                     let index_lexp = self.lookup_expression.lookup(index_id)?;
                     let index_handle = get_expr_handle!(index_id, index_lexp);
+                    let index_type = self.lookup_type.lookup(index_lexp.type_id)?.handle;
 
                     let num_components = match ctx.type_arena[root_type_lookup.handle].inner {
                         crate::TypeInner::Vector { size, .. } => size as u32,
                         _ => return Err(Error::InvalidVectorType(root_type_lookup.handle)),
                     };
 
-                    let make_index = |ctx: &mut BlockContext, index: u32| {
-                        ctx.expressions
-                            .append(crate::Expression::Literal(crate::Literal::U32(index)), span)
-                    };
-
                     let mut components = Vec::with_capacity(num_components as usize);
                     for index in 0..num_components {
-                        let index_expr = make_index(ctx, index);
+                        let index_expr = make_index_literal(
+                            ctx,
+                            index,
+                            &mut block,
+                            &mut emitter,
+                            index_type,
+                            index_lexp.type_id,
+                            span,
+                        )?;
                         let access_expr = ctx.expressions.append(
                             crate::Expression::Access {
                                 base: root_handle,
@@ -5145,6 +5157,30 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
     }
 }
 
+fn make_index_literal(
+    ctx: &mut BlockContext,
+    index: u32,
+    block: &mut crate::Block,
+    emitter: &mut super::Emitter,
+    index_type: Handle<crate::Type>,
+    index_type_id: spirv::Word,
+    span: crate::Span,
+) -> Result<Handle<crate::Expression>, Error> {
+    block.extend(emitter.finish(ctx.expressions));
+
+    let literal = match ctx.type_arena[index_type].inner.scalar_kind() {
+        Some(crate::ScalarKind::Uint) => crate::Literal::U32(index),
+        Some(crate::ScalarKind::Sint) => crate::Literal::I32(index as i32),
+        _ => return Err(Error::InvalidIndexType(index_type_id)),
+    };
+    let expr = ctx
+        .expressions
+        .append(crate::Expression::Literal(literal), span);
+
+    emitter.start(ctx.expressions);
+    Ok(expr)
+}
+
 pub fn parse_u8_slice(data: &[u8], options: &Options) -> Result<crate::Module, Error> {
     if data.len() % 4 != 0 {
         return Err(Error::IncompleteData);