Image queries

src/back/glsl/mod.rs

@@ -1338,6 +1338,74 @@ impl<'a, W: Write> Writer<'a, W> {
                 }
                 write!(self.out, ")")?;
             }
+            // Query translates into one of the:
+            // - textureSize/imageSize
+            // - textureQueryLevels
+            // - textureSamples/imageSamples
+            Expression::ImageQuery { image, query } => {
+                // This will only panic if the module is invalid
+                let (dim, class) = match ctx.typifier.get(image, &self.module.types) {
+                    TypeInner::Image {
+                        dim,
+                        arrayed: _,
+                        class,
+                    } => (dim, class),
+                    _ => unreachable!(),
+                };
+                let components = match dim {
+                    crate::ImageDimension::D1 => 1,
+                    crate::ImageDimension::D2 => 2,
+                    crate::ImageDimension::D3 => 3,
+                    crate::ImageDimension::Cube => 2,
+                };
+                match query {
+                    crate::ImageQuery::Size { level } => {
+                        match class {
+                            ImageClass::Sampled { .. } | ImageClass::Depth => {
+                                write!(self.out, "textureSize(")?;
+                                self.write_expr(image, ctx)?;
+                                write!(self.out, ",")?;
+                                if let Some(expr) = level {
+                                    self.write_expr(expr, ctx)?;
+                                } else {
+                                    write!(self.out, "0",)?;
+                                }
+                            }
+                            ImageClass::Storage(_) => {
+                                write!(self.out, "imageSize(")?;
+                                self.write_expr(image, ctx)?;
+                            }
+                        }
+                        write!(self.out, ").{}", &"xyz"[..components])?;
+                    }
+                    crate::ImageQuery::NumLevels => {
+                        write!(self.out, "textureQueryLevels(",)?;
+                        self.write_expr(image, ctx)?;
+                        write!(self.out, ")",)?;
+                    }
+                    crate::ImageQuery::NumLayers => {
+                        let selector = ['x', 'y', 'z', 'w'];
+                        let fun_name = match class {
+                            ImageClass::Sampled { .. } | ImageClass::Depth => "textureSize",
+                            ImageClass::Storage(_) => "imageSize",
+                        };
+                        write!(self.out, "{}(", fun_name)?;
+                        self.write_expr(image, ctx)?;
+                        write!(self.out, ",0).{}", selector[components])?;
+                    }
+                    crate::ImageQuery::NumSamples => {
+                        // assumes ARB_shader_texture_image_samples
+                        let fun_name = match class {
+                            ImageClass::Sampled { .. } | ImageClass::Depth => "textureSamples",
+                            ImageClass::Storage(_) => "imageSamples",
+                        };
+                        write!(self.out, "{}(", fun_name)?;
+                        self.write_expr(image, ctx)?;
+                        write!(self.out, ")",)?;
+                    }
+                }
+                return Err(Error::Custom("ImageQuery not implemented".to_string()));
+            }
             // `Unary` is pretty straightforward
             // "-" - for `Negate`
             // "~" - for `Not` if it's an integer

src/back/msl/writer.rs

@@ -143,6 +143,22 @@ impl<W: Write> Writer<W> {
         Ok(())
     }
 
+    fn put_image_query(
+        &mut self,
+        image: Handle<crate::Expression>,
+        query: &str,
+        level: Option<Handle<crate::Expression>>,
+        context: &ExpressionContext,
+    ) -> Result<(), Error> {
+        self.put_expression(image, context)?;
+        write!(self.out, ".get_{}(", query)?;
+        if let Some(expr) = level {
+            self.put_expression(expr, context)?;
+        }
+        write!(self.out, ")")?;
+        Ok(())
+    }
+
     fn put_expression(
         &mut self,
         expr_handle: Handle<crate::Expression>,
@@ -331,6 +347,59 @@ impl<W: Write> Writer<W> {
                 }
                 write!(self.out, ")")?;
             }
+            //Note: for all the queries, the signed integers are expected,
+            // so a conversion is needed.
+            crate::Expression::ImageQuery { image, query } => match query {
+                crate::ImageQuery::Size { level } => {
+                    let dim = match *self.typifier.get(image, &context.module.types) {
+                        crate::TypeInner::Image { dim, .. } => dim,
+                        ref other => unreachable!("Unexpected type {:?}", other),
+                    };
+                    match dim {
+                        crate::ImageDimension::D1 => {
+                            write!(self.out, "int(")?;
+                            self.put_image_query(image, "width", level, context)?;
+                            write!(self.out, ")")?;
+                        }
+                        crate::ImageDimension::D2 => {
+                            write!(self.out, "int2(")?;
+                            self.put_image_query(image, "width", level, context)?;
+                            write!(self.out, ", ")?;
+                            self.put_image_query(image, "height", level, context)?;
+                            write!(self.out, ")")?;
+                        }
+                        crate::ImageDimension::D3 => {
+                            write!(self.out, "int3(")?;
+                            self.put_image_query(image, "width", level, context)?;
+                            write!(self.out, ", ")?;
+                            self.put_image_query(image, "height", level, context)?;
+                            write!(self.out, ", ")?;
+                            self.put_image_query(image, "depth", level, context)?;
+                            write!(self.out, ")")?;
+                        }
+                        crate::ImageDimension::Cube => {
+                            write!(self.out, "int(")?;
+                            self.put_image_query(image, "width", level, context)?;
+                            write!(self.out, ").xxx")?;
+                        }
+                    }
+                }
+                crate::ImageQuery::NumLevels => {
+                    write!(self.out, "int(")?;
+                    self.put_expression(image, context)?;
+                    write!(self.out, ".get_num_mip_levels())")?;
+                }
+                crate::ImageQuery::NumLayers => {
+                    write!(self.out, "int(")?;
+                    self.put_expression(image, context)?;
+                    write!(self.out, ".get_array_size())")?;
+                }
+                crate::ImageQuery::NumSamples => {
+                    write!(self.out, "int(")?;
+                    self.put_expression(image, context)?;
+                    write!(self.out, ".get_num_samples())")?;
+                }
+            },
             crate::Expression::Unary { op, expr } => {
                 let op_str = match op {
                     crate::UnaryOperator::Negate => "-",

src/lib.rs

@@ -621,6 +621,24 @@ pub enum SampleLevel {
     },
 }
 
+/// Type of an image query.
+#[derive(Clone, Copy, Debug, PartialEq)]
+#[cfg_attr(feature = "serialize", derive(Serialize))]
+#[cfg_attr(feature = "deserialize", derive(Deserialize))]
+pub enum ImageQuery {
+    /// Get the size at the specified level.
+    Size {
+        /// If `None`, the base level is considered.
+        level: Option<Handle<Expression>>,
+    },
+    /// Get the number of mipmap levels.
+    NumLevels,
+    /// Get the number of array layers.
+    NumLayers,
+    /// Get the number of samples.
+    NumSamples,
+}
+
 /// An expression that can be evaluated to obtain a value.
 #[derive(Clone, Debug)]
 #[cfg_attr(feature = "serialize", derive(Serialize))]
@@ -671,6 +689,11 @@ pub enum Expression {
         /// For multisampled images, this is Some(Sample).
         index: Option<Handle<Expression>>,
     },
+    /// Query information from an image.
+    ImageQuery {
+        image: Handle<Expression>,
+        query: ImageQuery,
+    },
     /// Apply an unary operator.
     Unary {
         op: UnaryOperator,

src/proc/constants.rs

@@ -120,8 +120,9 @@ impl<'a> ConstantSolver<'a> {
             Expression::Call { .. } => Err(ConstantSolvingError::Call),
             Expression::FunctionArgument(_) => Err(ConstantSolvingError::FunctionArg),
             Expression::GlobalVariable(_) => Err(ConstantSolvingError::GlobalVariable),
-            Expression::ImageSample { .. } => Err(ConstantSolvingError::ImageExpression),
-            Expression::ImageLoad { .. } => Err(ConstantSolvingError::ImageExpression),
+            Expression::ImageSample { .. }
+            | Expression::ImageLoad { .. }
+            | Expression::ImageQuery { .. } => Err(ConstantSolvingError::ImageExpression),
         }
     }
 

src/proc/interface.rs

@@ -85,6 +85,16 @@ where
                     self.traverse_expr(index);
                 }
             }
+            E::ImageQuery { image, query } => {
+                self.traverse_expr(image);
+                match query {
+                    crate::ImageQuery::Size { level: Some(expr) } => self.traverse_expr(expr),
+                    crate::ImageQuery::Size { .. }
+                    | crate::ImageQuery::NumLevels
+                    | crate::ImageQuery::NumLayers
+                    | crate::ImageQuery::NumSamples => (),
+                }
+            }
             E::Unary { expr, .. } => {
                 self.traverse_expr(expr);
             }

src/proc/typifier.rs

@@ -204,6 +204,35 @@ impl Typifier {
                 }),
                 _ => unreachable!(),
             },
+            crate::Expression::ImageQuery { image, query } => Resolution::Value(match query {
+                crate::ImageQuery::Size { level: _ } => match *self.get(image, types) {
+                    crate::TypeInner::Image { dim, .. } => match dim {
+                        crate::ImageDimension::D1 => crate::TypeInner::Scalar {
+                            kind: crate::ScalarKind::Sint,
+                            width: 4,
+                        },
+                        crate::ImageDimension::D2 => crate::TypeInner::Vector {
+                            size: crate::VectorSize::Bi,
+                            kind: crate::ScalarKind::Sint,
+                            width: 4,
+                        },
+                        crate::ImageDimension::D3 | crate::ImageDimension::Cube => {
+                            crate::TypeInner::Vector {
+                                size: crate::VectorSize::Tri,
+                                kind: crate::ScalarKind::Sint,
+                                width: 4,
+                            }
+                        }
+                    },
+                    _ => unreachable!(),
+                },
+                crate::ImageQuery::NumLevels
+                | crate::ImageQuery::NumLayers
+                | crate::ImageQuery::NumSamples => crate::TypeInner::Scalar {
+                    kind: crate::ScalarKind::Sint,
+                    width: 4,
+                },
+            }),
             crate::Expression::Unary { expr, .. } => self.resolutions[expr.index()].clone(),
             crate::Expression::Binary { op, left, right } => match op {
                 crate::BinaryOperator::Add