[naga wgsl-in] Proper singular generic in vec and matrix (#6189)

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

Author: sagudev

naga/src/front/wgsl/lower/construction.rs

@@ -578,8 +578,13 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                 Constructor::Type(ty)
             }
             ast::ConstructorType::PartialVector { size } => Constructor::PartialVector { size },
-            ast::ConstructorType::Vector { size, scalar } => {
-                let ty = ctx.ensure_type_exists(scalar.to_inner_vector(size));
+            ast::ConstructorType::Vector { size, ty, ty_span } => {
+                let ty = self.resolve_ast_type(ty, &mut ctx.as_global())?;
+                let scalar = match ctx.module.types[ty].inner {
+                    crate::TypeInner::Scalar(sc) => sc,
+                    _ => return Err(Error::UnknownScalarType(ty_span)),
+                };
+                let ty = ctx.ensure_type_exists(crate::TypeInner::Vector { size, scalar });
                 Constructor::Type(ty)
             }
             ast::ConstructorType::PartialMatrix { columns, rows } => {
@@ -588,13 +593,22 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             ast::ConstructorType::Matrix {
                 rows,
                 columns,
-                width,
+                ty,
+                ty_span,
             } => {
-                let ty = ctx.ensure_type_exists(crate::TypeInner::Matrix {
-                    columns,
-                    rows,
-                    scalar: crate::Scalar::float(width),
-                });
+                let ty = self.resolve_ast_type(ty, &mut ctx.as_global())?;
+                let scalar = match ctx.module.types[ty].inner {
+                    crate::TypeInner::Scalar(sc) => sc,
+                    _ => return Err(Error::UnknownScalarType(ty_span)),
+                };
+                let ty = match scalar.kind {
+                    crate::ScalarKind::Float => ctx.ensure_type_exists(crate::TypeInner::Matrix {
+                        columns,
+                        rows,
+                        scalar,
+                    }),
+                    _ => return Err(Error::BadMatrixScalarKind(ty_span, scalar)),
+                };
                 Constructor::Type(ty)
             }
             ast::ConstructorType::PartialArray => Constructor::PartialArray,

naga/src/front/wgsl/lower/mod.rs

@@ -2965,16 +2965,34 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
     ) -> Result<Handle<crate::Type>, Error<'source>> {
         let inner = match ctx.types[handle] {
             ast::Type::Scalar(scalar) => scalar.to_inner_scalar(),
-            ast::Type::Vector { size, scalar } => scalar.to_inner_vector(size),
+            ast::Type::Vector { size, ty, ty_span } => {
+                let ty = self.resolve_ast_type(ty, ctx)?;
+                let scalar = match ctx.module.types[ty].inner {
+                    crate::TypeInner::Scalar(sc) => sc,
+                    _ => return Err(Error::UnknownScalarType(ty_span)),
+                };
+                crate::TypeInner::Vector { size, scalar }
+            }
             ast::Type::Matrix {
                 rows,
                 columns,
-                width,
-            } => crate::TypeInner::Matrix {
-                columns,
-                rows,
-                scalar: crate::Scalar::float(width),
-            },
+                ty,
+                ty_span,
+            } => {
+                let ty = self.resolve_ast_type(ty, ctx)?;
+                let scalar = match ctx.module.types[ty].inner {
+                    crate::TypeInner::Scalar(sc) => sc,
+                    _ => return Err(Error::UnknownScalarType(ty_span)),
+                };
+                match scalar.kind {
+                    crate::ScalarKind::Float => crate::TypeInner::Matrix {
+                        columns,
+                        rows,
+                        scalar,
+                    },
+                    _ => return Err(Error::BadMatrixScalarKind(ty_span, scalar)),
+                }
+            }
             ast::Type::Atomic(scalar) => scalar.to_inner_atomic(),
             ast::Type::Pointer { base, space } => {
                 let base = self.resolve_ast_type(base, ctx)?;

naga/src/front/wgsl/parse/ast.rs

@@ -198,12 +198,14 @@ pub enum Type<'a> {
     Scalar(Scalar),
     Vector {
         size: crate::VectorSize,
-        scalar: Scalar,
+        ty: Handle<Type<'a>>,
+        ty_span: Span,
     },
     Matrix {
         columns: crate::VectorSize,
         rows: crate::VectorSize,
-        width: crate::Bytes,
+        ty: Handle<Type<'a>>,
+        ty_span: Span,
     },
     Atomic(Scalar),
     Pointer {
@@ -330,7 +332,8 @@ pub enum ConstructorType<'a> {
     /// `vec3<f32>(1.0)`.
     Vector {
         size: crate::VectorSize,
-        scalar: Scalar,
+        ty: Handle<Type<'a>>,
+        ty_span: Span,
     },
 
     /// A matrix construction whose component type is inferred from the
@@ -345,7 +348,8 @@ pub enum ConstructorType<'a> {
     Matrix {
         columns: crate::VectorSize,
         rows: crate::VectorSize,
-        width: crate::Bytes,
+        ty: Handle<Type<'a>>,
+        ty_span: Span,
     },
 
     /// An array whose component type and size are inferred from the arguments: