From c7d1d96540688b7a853fa69d4ca45d968755a89c Mon Sep 17 00:00:00 2001
From: Mikko Rasa <tdb@tdb.fi>
Date: Thu, 11 Nov 2021 12:56:28 +0200
Subject: [PATCH] Handle descriptor sets when allocating uniform locations

---
 source/glsl/finalize.cpp | 28 +++++++++++++++++++++++++---
 source/glsl/finalize.h   |  3 ++-
 source/glsl/validate.cpp |  5 +++++
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/source/glsl/finalize.cpp b/source/glsl/finalize.cpp
index cb304aa3..60312be8 100644
--- a/source/glsl/finalize.cpp
+++ b/source/glsl/finalize.cpp
@@ -62,8 +62,9 @@ void StructOrganizer::visit(VariableDeclaration &var)
 }
 
 
-void LocationAllocator::apply(Module &module, const Features &features, bool a)
+void LocationAllocator::apply(Module &module, const Features &f, bool a)
 {
+	features = f;
 	alloc_new = a;
 	for(Stage &s: module.stages)
 		apply(s);
@@ -144,11 +145,16 @@ void LocationAllocator::allocate_locations(const string &iface)
 void LocationAllocator::bind_uniform(RefPtr<Layout> &layout, const string &name, unsigned range)
 {
 	auto i = uniforms.find(name);
+
+	int desc_set = (i!=uniforms.end() ? i->second.desc_set : 0);
+	if(features.target_api==VULKAN && get_layout_value(layout.get(), "set")<0)
+		add_layout_qualifier(layout, Layout::Qualifier("set", desc_set));
+
 	if(i!=uniforms.end() && i->second.bind_point>=0)
 		add_layout_qualifier(layout, Layout::Qualifier("binding", i->second.bind_point));
 	else if(alloc_new)
 	{
-		set<unsigned> &used = used_bindings[0];
+		set<unsigned> &used = used_bindings[desc_set];
 
 		unsigned bind_point = fold32(hash64(name))%range;
 		while(used.count(bind_point))
@@ -162,12 +168,28 @@ void LocationAllocator::bind_uniform(RefPtr<Layout> &layout, const string &name,
 
 bool LocationAllocator::visit_uniform(const string &name, RefPtr<Layout> &layout)
 {
+	int desc_set = 0;
 	int bind_point = get_layout_value(layout.get(), "binding");
+
+	if(features.target_api==VULKAN)
+	{
+		desc_set = get_layout_value(layout.get(), "set");
+		if(desc_set<0 && bind_point>=0)
+		{
+			desc_set = 0;
+			add_layout_qualifier(layout, Layout::Qualifier("set", desc_set));
+		}
+
+		if(desc_set>=0)
+			uniforms[name].desc_set = desc_set;
+	}
+
 	if(bind_point>=0)
 	{
-		used_bindings[0].insert(bind_point);
+		used_bindings[desc_set].insert(bind_point);
 		uniforms[name].bind_point = bind_point;
 	}
+
 	return bind_point>=0;
 }
 
diff --git a/source/glsl/finalize.h b/source/glsl/finalize.h
index ab4276f5..19ad5947 100644
--- a/source/glsl/finalize.h
+++ b/source/glsl/finalize.h
@@ -33,9 +33,10 @@ private:
 		int desc_set;
 		int bind_point;
 
-		Uniform(): location(-1), desc_set(-1), bind_point(-1) { }
+		Uniform(): location(-1), desc_set(0), bind_point(-1) { }
 	};
 
+	Features features;
 	bool alloc_new = true;
 	std::map<std::string, std::set<unsigned> > used_locations;
 	std::map<std::string, Uniform> uniforms;
diff --git a/source/glsl/validate.cpp b/source/glsl/validate.cpp
index 6a4b8eb0..7d3d20d1 100644
--- a/source/glsl/validate.cpp
+++ b/source/glsl/validate.cpp
@@ -834,6 +834,11 @@ void GlobalInterfaceValidator::check_uniform(const Uniform &uni)
 			error(*uni.node, format("Mismatched location %d for uniform '%s'", uni.location, uni.name));
 			add_info(*i->second->node, format("Previously declared here with location %d", i->second->location));
 		}
+		if(i->second->desc_set!=uni.desc_set)
+		{
+			error(*uni.node, format("Mismatched descriptor set %d for uniform '%s'", uni.desc_set, uni.name));
+			add_info(*i->second->node, format("Previously declared here with descriptor set %d", i->second->desc_set));
+		}
 		if(uni.bind_point>=0 && i->second->bind_point>=0 && i->second->bind_point!=uni.bind_point)
 		{
 			error(*uni.node, format("Mismatched binding %d for uniform '%s'", uni.bind_point, uni.name));
-- 
2.45.2