Support int and uint as flattened UBO types.

reference/shaders/flatten/push-constant.flatten.vert

@@ -0,0 +1,13 @@
+#version 310 es
+
+uniform vec4 PushMe[5];
+layout(location = 1) in vec4 Pos;
+layout(location = 0) out vec2 vRot;
+layout(location = 0) in vec2 Rot;
+
+void main()
+{
+    gl_Position = mat4(PushMe[0], PushMe[1], PushMe[2], PushMe[3]) * Pos;
+    vRot = mat2(PushMe[4].xy, PushMe[4].zw) * Rot;
+}
+

reference/shaders/flatten/types.flatten.frag

@@ -0,0 +1,14 @@
+#version 310 es
+precision mediump float;
+precision highp int;
+
+uniform mediump ivec4 UBO1[2];
+uniform mediump uvec4 UBO2[2];
+uniform vec4 UBO0[2];
+layout(location = 0) out vec4 FragColor;
+
+void main()
+{
+    FragColor = ((((vec4(UBO1[0]) + vec4(UBO1[1])) + vec4(UBO2[0])) + vec4(UBO2[1])) + UBO0[0]) + UBO0[1];
+}
+

shaders/flatten/push-constant.flatten.vert

@@ -0,0 +1,16 @@
+#version 310 es
+
+layout(push_constant, std430) uniform PushMe
+{
+   mat4 MVP;
+   mat2 Rot; // The MatrixStride will be 8 here.
+} registers;
+
+layout(location = 0) in vec2 Rot;
+layout(location = 1) in vec4 Pos;
+layout(location = 0) out vec2 vRot;
+void main()
+{
+   gl_Position = registers.MVP * Pos;
+   vRot = registers.Rot * Rot;
+}

shaders/flatten/types.flatten.frag

@@ -0,0 +1,27 @@
+#version 310 es
+precision mediump float;
+
+layout(std140, binding = 0) uniform UBO0
+{
+   vec4 a;
+   vec4 b;
+};
+
+layout(std140, binding = 0) uniform UBO1
+{
+   ivec4 c;
+   ivec4 d;
+};
+
+layout(std140, binding = 0) uniform UBO2
+{
+   uvec4 e;
+   uvec4 f;
+};
+
+layout(location = 0) out vec4 FragColor;
+
+void main()
+{
+   FragColor = vec4(c) + vec4(d) + vec4(e) + vec4(f) + a + b; 
+}

spirv_cross.cpp

@@ -3056,3 +3056,28 @@ uint64_t Compiler::get_buffer_block_flags(const SPIRVariable &var)
 
 	return base_flags | (~all_members_flag_mask);
 }
+
+bool Compiler::get_common_basic_type(const SPIRType &type, SPIRType::BaseType &base_type)
+{
+	if (type.basetype == SPIRType::Struct)
+	{
+		base_type = SPIRType::Unknown;
+		for (auto &member_type : type.member_types)
+		{
+			SPIRType::BaseType member_base;
+			if (!get_common_basic_type(get<SPIRType>(member_type), member_base))
+				return false;
+
+			if (base_type == SPIRType::Unknown)
+				base_type = member_base;
+			else if (base_type != member_base)
+				return false;
+		}
+		return true;
+	}
+	else
+	{
+		base_type = type.basetype;
+		return true;
+	}
+}

spirv_cross.hpp

@@ -579,6 +579,7 @@ protected:
 	VariableTypeRemapCallback variable_remap_callback;
 
 	uint64_t get_buffer_block_flags(const SPIRVariable &var);
+	bool get_common_basic_type(const SPIRType &type, SPIRType::BaseType &base_type);
 };
 }
 

spirv_glsl.cpp

@@ -1094,7 +1094,33 @@ void CompilerGLSL::emit_buffer_block_flattened(const SPIRVariable &var)
 	auto buffer_name = to_name(type.self, false);
 	size_t buffer_size = (get_declared_struct_size(type) + 15) / 16;
 
-	statement("uniform vec4 ", buffer_name, "[", buffer_size, "];");
+	SPIRType::BaseType basic_type;
+	if (get_common_basic_type(type, basic_type))
+	{
+		SPIRType tmp;
+		tmp.basetype = basic_type;
+
+		const char *flat_type = nullptr;
+		switch (basic_type)
+		{
+		case SPIRType::Float:
+			flat_type = "vec4 ";
+			break;
+		case SPIRType::Int:
+			flat_type = "ivec4 ";
+			break;
+		case SPIRType::UInt:
+			flat_type = "uvec4 ";
+			break;
+		default:
+			SPIRV_CROSS_THROW("Basic types in a flattened UBO must be float, int or uint.");
+		}
+
+		auto flags = get_buffer_block_flags(var);
+		statement("uniform ", flags_to_precision_qualifiers_glsl(tmp, flags), flat_type, buffer_name, "[", buffer_size, "];");
+	}
+	else
+		SPIRV_CROSS_THROW("All basic types in a flattened block must be the same.");
 }
 
 void CompilerGLSL::emit_interface_block(const SPIRVariable &var)
@@ -3346,15 +3372,18 @@ std::string CompilerGLSL::flattened_access_chain_matrix(uint32_t base, const uin
 {
 	std::string expr;
 
+	uint32_t matrix_stride = 0;
+	flattened_access_chain_offset(base, indices, count, offset, nullptr, &matrix_stride);
+
 	expr += type_to_glsl_constructor(target_type);
 	expr += "(";
 
-	for (uint32_t i = 0; i < target_type.columns; ++i)
+	for (uint32_t i = 0; i < target_type.columns; i++)
 	{
 		if (i != 0)
 			expr += ", ";
 
-		expr += flattened_access_chain_vector_scalar(base, indices, count, target_type, offset + i * 16);
+		expr += flattened_access_chain_vector_scalar(base, indices, count, target_type, offset + i * matrix_stride);
 	}
 
 	expr += ")";
@@ -3365,13 +3394,10 @@ std::string CompilerGLSL::flattened_access_chain_matrix(uint32_t base, const uin
 std::string CompilerGLSL::flattened_access_chain_vector_scalar(uint32_t base, const uint32_t *indices, uint32_t count,
                                                                const SPIRType &target_type, uint32_t offset)
 {
-	if (target_type.basetype != SPIRType::Float)
-		SPIRV_CROSS_THROW("Access chains that use non-floating-point base types can not be flattened");
-
 	auto result = flattened_access_chain_offset(base, indices, count, offset);
 
-	assert(result.second % 4 == 0);
-	uint32_t index = result.second / 4;
+	assert(result.second % (target_type.width / 8) == 0);
+	uint32_t index = 8 * result.second / target_type.width;
 
 	auto buffer_name = to_name(expression_type(base).self);
 
@@ -3379,7 +3405,7 @@ std::string CompilerGLSL::flattened_access_chain_vector_scalar(uint32_t base, co
 
 	expr += buffer_name;
 	expr += "[";
-	expr += result.first; // this is a series of N1*k1+N2*k2+... that is either empty or ends with a +
+	expr += result.first; // this is a series of N1 * k1 + N2 * k2 + ... that is either empty or ends with a +
 	expr += convert_to_string(index / 4);
 	expr += "]";
 
@@ -3390,7 +3416,7 @@ std::string CompilerGLSL::flattened_access_chain_vector_scalar(uint32_t base, co
 
 std::pair<std::string, uint32_t> CompilerGLSL::flattened_access_chain_offset(uint32_t base, const uint32_t *indices,
                                                                              uint32_t count, uint32_t offset,
-                                                                             bool *need_transpose)
+                                                                             bool *need_transpose, uint32_t *out_matrix_stride)
 {
 	const auto *type = &expression_type(base);
 	uint32_t current_type = type->self;
@@ -3410,10 +3436,28 @@ std::pair<std::string, uint32_t> CompilerGLSL::flattened_access_chain_offset(uin
 			if (!array_stride)
 				SPIRV_CROSS_THROW("SPIR-V does not define ArrayStride for buffer block.");
 
-			expr += to_expression(index);
-			expr += " * ";
-			expr += convert_to_string(array_stride / 16);
-			expr += " + ";
+			auto *constant = maybe_get<SPIRConstant>(index);
+			if (constant)
+			{
+				// Constant array access.
+				offset += constant->scalar() * array_stride;
+			}
+			else
+			{
+				// Dynamic array access.
+				// FIXME: This will need to change if we support other flattening types than 32-bit.
+				const uint32_t word_stride = 16;
+				if (array_stride % word_stride)
+				{
+					SPIRV_CROSS_THROW("Array stride for dynamic indexing must be divisible by the size of a 4-component vector. "
+						                  "Likely culprit here is a float or vec2 array inside a push constant block. This cannot be flattened.");
+				}
+
+				expr += to_expression(index);
+				expr += " * ";
+				expr += convert_to_string(array_stride / word_stride);
+				expr += " + ";
+			}
 
 			uint32_t parent_type = type->parent_type;
 			type = &get<SPIRType>(parent_type);
@@ -3476,6 +3520,8 @@ std::pair<std::string, uint32_t> CompilerGLSL::flattened_access_chain_offset(uin
 
 	if (need_transpose)
 		*need_transpose = row_major_matrix_needs_conversion;
+	if (out_matrix_stride)
+		*out_matrix_stride = matrix_stride;
 
 	return std::make_pair(expr, offset);
 }

spirv_glsl.hpp

@@ -329,7 +329,7 @@ protected:
 	                                                 const SPIRType &target_type, uint32_t offset);
 	std::pair<std::string, uint32_t> flattened_access_chain_offset(uint32_t base, const uint32_t *indices,
 	                                                               uint32_t count, uint32_t offset,
-	                                                               bool *need_transpose = nullptr);
+	                                                               bool *need_transpose = nullptr, uint32_t *matrix_stride = nullptr);
 
 	const char *index_to_swizzle(uint32_t index);
 	std::string remap_swizzle(uint32_t result_type, uint32_t input_components, uint32_t expr);