Fix descriptor set allocation and update

- Descriptor set pool allocation did not take the layout pointer into
  account.
- Allocating descriptor sets from the pool was not tracking the first
  node.
- The pointer to a second allocation node was offset from null instead
  of the pool base address.
- Added assert that the descriptor update type matches the layout.
- Refactoring to avoid duplicate casting of handles.

Bug b/123244275

Tests: dEQP-VK.api.object_management.*
Change-Id: Idf7aeb8d2597b30038ba1e9e371d99f09639f13c
Reviewed-on: https://swiftshader-review.googlesource.com/c/SwiftShader/+/28230
Presubmit-Ready: Nicolas Capens <nicolascapens@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
Kokoro-Presubmit: kokoro <noreply+kokoro@google.com>
Tested-by: Nicolas Capens <nicolascapens@google.com>
diff --git a/src/Vulkan/VkDescriptorPool.cpp b/src/Vulkan/VkDescriptorPool.cpp
index 1de234e..a05b5a2 100644
--- a/src/Vulkan/VkDescriptorPool.cpp
+++ b/src/Vulkan/VkDescriptorPool.cpp
@@ -51,7 +51,7 @@
 	for(uint32_t i = 0; i < descriptorSetCount; i++)
 	{
 		pDescriptorSets[i] = VK_NULL_HANDLE;
-		layoutSizes[i] = Cast(pSetLayouts[i])->getSize();
+		layoutSizes[i] = Cast(pSetLayouts[i])->getDescriptorSetAllocationSize();
 	}
 
 	VkResult result = allocateSets(&(layoutSizes[0]), descriptorSetCount, pDescriptorSets);
@@ -79,7 +79,7 @@
 	size_t freeSpace = poolSize - nextItemStart;
 	if(freeSpace >= size)
 	{
-		return reinterpret_cast<VkDescriptorSet>(nextItemStart);
+		return reinterpret_cast<VkDescriptorSet>(reinterpret_cast<char*>(pool) + nextItemStart);
 	}
 
 	// Second, look for space at the beginning of the pool
@@ -121,17 +121,19 @@
 	}
 
 	// Attempt to allocate single chunk of memory
-	VkDescriptorSet memory = findAvailableMemory(totalSize);
-	if(memory != VK_NULL_HANDLE)
 	{
-		pDescriptorSets[0] = memory;
-		for(uint32_t i = 1; i < numAllocs; i++)
+		VkDescriptorSet memory = findAvailableMemory(totalSize);
+		if(memory != VK_NULL_HANDLE)
 		{
-			pDescriptorSets[i] =
-				reinterpret_cast<VkDescriptorSet>(reinterpret_cast<char*>(memory) + sizes[i - 1]);
-			nodes.insert(Node(pDescriptorSets[i], sizes[i]));
+			for(uint32_t i = 0; i < numAllocs; i++)
+			{
+				pDescriptorSets[i] = memory;
+				nodes.insert(Node(pDescriptorSets[i], sizes[i]));
+				memory = reinterpret_cast<VkDescriptorSet>(reinterpret_cast<char*>(memory) + sizes[i]);
+			}
+
+			return VK_SUCCESS;
 		}
-		return VK_SUCCESS;
 	}
 
 	// Atttempt to allocate each descriptor set separately
diff --git a/src/Vulkan/VkDescriptorSetLayout.cpp b/src/Vulkan/VkDescriptorSetLayout.cpp
index 3517768..4b65eca 100644
--- a/src/Vulkan/VkDescriptorSetLayout.cpp
+++ b/src/Vulkan/VkDescriptorSetLayout.cpp
@@ -21,17 +21,6 @@
 namespace
 {
 
-struct DescriptorSet
-{
-	vk::DescriptorSetLayout* layout;
-	uint8_t data[];
-};
-
-static inline DescriptorSet* Cast(VkDescriptorSet object)
-{
-	return reinterpret_cast<DescriptorSet*>(object);
-}
-
 static bool UsesImmutableSamplers(const VkDescriptorSetLayoutBinding& binding)
 {
 	return (((binding.descriptorType == VK_DESCRIPTOR_TYPE_SAMPLER) ||
@@ -118,13 +107,20 @@
 	return 0;
 }
 
-size_t DescriptorSetLayout::getSize() const
+size_t DescriptorSetLayout::getDescriptorSetAllocationSize() const
+{
+	// vk::DescriptorSet has a layout member field.
+	return sizeof(vk::DescriptorSetLayout*) + getDescriptorSetDataSize();
+}
+
+size_t DescriptorSetLayout::getDescriptorSetDataSize() const
 {
 	size_t size = 0;
 	for(uint32_t i = 0; i < bindingCount; i++)
 	{
 		size += bindings[i].descriptorCount * GetDescriptorSize(bindings[i].descriptorType);
 	}
+
 	return size;
 }
 
@@ -145,7 +141,7 @@
 void DescriptorSetLayout::initialize(VkDescriptorSet vkDescriptorSet)
 {
 	// Use a pointer to this descriptor set layout as the descriptor set's header
-	DescriptorSet* descriptorSet = ::Cast(vkDescriptorSet);
+	DescriptorSet* descriptorSet = vk::Cast(vkDescriptorSet);
 	descriptorSet->layout = this;
 	uint8_t* mem = descriptorSet->data;
 
@@ -174,13 +170,13 @@
 	return bindingOffsets[index] + OFFSET(DescriptorSet, data[0]);
 }
 
-uint8_t* DescriptorSetLayout::getOffsetPointer(VkDescriptorSet descriptorSet, uint32_t binding, uint32_t arrayElement, uint32_t count, size_t* typeSize) const
+uint8_t* DescriptorSetLayout::getOffsetPointer(DescriptorSet *descriptorSet, uint32_t binding, uint32_t arrayElement, uint32_t count, size_t* typeSize) const
 {
 	uint32_t index = getBindingIndex(binding);
 	*typeSize = GetDescriptorSize(bindings[index].descriptorType);
 	size_t byteOffset = bindingOffsets[index] + (*typeSize * arrayElement);
-	ASSERT(((*typeSize * count) + byteOffset) <= getSize()); // Make sure the operation will not go out of bounds
-	return &(::Cast(descriptorSet)->data[byteOffset]);
+	ASSERT(((*typeSize * count) + byteOffset) <= getDescriptorSetDataSize()); // Make sure the operation will not go out of bounds
+	return &descriptorSet->data[byteOffset];
 }
 
 const uint8_t* DescriptorSetLayout::GetInputData(const VkWriteDescriptorSet& descriptorWrites)
@@ -211,12 +207,13 @@
 
 void DescriptorSetLayout::WriteDescriptorSet(const VkWriteDescriptorSet& descriptorWrites)
 {
-	DescriptorSet* dstSet = ::Cast(descriptorWrites.dstSet);
+	DescriptorSet* dstSet = vk::Cast(descriptorWrites.dstSet);
 	DescriptorSetLayout* dstLayout = dstSet->layout;
 	ASSERT(dstLayout);
+	ASSERT(dstLayout->bindings[dstLayout->getBindingIndex(descriptorWrites.dstBinding)].descriptorType == descriptorWrites.descriptorType);
 
 	size_t typeSize = 0;
-	uint8_t* memToWrite = dstLayout->getOffsetPointer(descriptorWrites.dstSet, descriptorWrites.dstBinding, descriptorWrites.dstArrayElement, descriptorWrites.descriptorCount, &typeSize);
+	uint8_t* memToWrite = dstLayout->getOffsetPointer(dstSet, descriptorWrites.dstBinding, descriptorWrites.dstArrayElement, descriptorWrites.descriptorCount, &typeSize);
 
 	// If the dstBinding has fewer than descriptorCount array elements remaining
 	// starting from dstArrayElement, then the remainder will be used to update
@@ -230,19 +227,19 @@
 
 void DescriptorSetLayout::CopyDescriptorSet(const VkCopyDescriptorSet& descriptorCopies)
 {
-	DescriptorSet* srcSet = ::Cast(descriptorCopies.srcSet);
+	DescriptorSet* srcSet = vk::Cast(descriptorCopies.srcSet);
 	DescriptorSetLayout* srcLayout = srcSet->layout;
 	ASSERT(srcLayout);
 
-	DescriptorSet* dstSet = ::Cast(descriptorCopies.dstSet);
+	DescriptorSet* dstSet = vk::Cast(descriptorCopies.dstSet);
 	DescriptorSetLayout* dstLayout = dstSet->layout;
 	ASSERT(dstLayout);
 
 	size_t srcTypeSize = 0;
-	uint8_t* memToRead = srcLayout->getOffsetPointer(descriptorCopies.srcSet, descriptorCopies.srcBinding, descriptorCopies.srcArrayElement, descriptorCopies.descriptorCount, &srcTypeSize);
+	uint8_t* memToRead = srcLayout->getOffsetPointer(srcSet, descriptorCopies.srcBinding, descriptorCopies.srcArrayElement, descriptorCopies.descriptorCount, &srcTypeSize);
 
 	size_t dstTypeSize = 0;
-	uint8_t* memToWrite = dstLayout->getOffsetPointer(descriptorCopies.dstSet, descriptorCopies.dstBinding, descriptorCopies.dstArrayElement, descriptorCopies.descriptorCount, &dstTypeSize);
+	uint8_t* memToWrite = dstLayout->getOffsetPointer(dstSet, descriptorCopies.dstBinding, descriptorCopies.dstArrayElement, descriptorCopies.descriptorCount, &dstTypeSize);
 
 	ASSERT(srcTypeSize == dstTypeSize);
 	size_t writeSize = dstTypeSize * descriptorCopies.descriptorCount;
diff --git a/src/Vulkan/VkDescriptorSetLayout.hpp b/src/Vulkan/VkDescriptorSetLayout.hpp
index 3905a16..8cad112 100644
--- a/src/Vulkan/VkDescriptorSetLayout.hpp
+++ b/src/Vulkan/VkDescriptorSetLayout.hpp
@@ -20,6 +20,19 @@
 namespace vk
 {
 
+class DescriptorSetLayout;
+
+struct DescriptorSet
+{
+	vk::DescriptorSetLayout* layout;
+	uint8_t data[];
+};
+
+inline DescriptorSet* Cast(VkDescriptorSet object)
+{
+	return reinterpret_cast<DescriptorSet*>(object);
+}
+
 class DescriptorSetLayout : public Object<DescriptorSetLayout, VkDescriptorSetLayout>
 {
 public:
@@ -34,11 +47,13 @@
 	static void CopyDescriptorSet(const VkCopyDescriptorSet& descriptorCopies);
 
 	void initialize(VkDescriptorSet descriptorSet);
-	size_t getSize() const;
+	size_t getDescriptorSetAllocationSize() const;
+
 	size_t getBindingOffset(uint32_t binding) const;
-	uint8_t* getOffsetPointer(VkDescriptorSet descriptorSet, uint32_t binding, uint32_t arrayElement, uint32_t count, size_t* typeSize) const;
+	uint8_t* getOffsetPointer(DescriptorSet *descriptorSet, uint32_t binding, uint32_t arrayElement, uint32_t count, size_t* typeSize) const;
 
 private:
+	size_t getDescriptorSetDataSize() const;
 	uint32_t getBindingIndex(uint32_t binding) const;
 	static const uint8_t* GetInputData(const VkWriteDescriptorSet& descriptorWrites);
 
diff --git a/src/Vulkan/VkDescriptorUpdateTemplate.cpp b/src/Vulkan/VkDescriptorUpdateTemplate.cpp
index 903f5f2..1a814e8 100644
--- a/src/Vulkan/VkDescriptorUpdateTemplate.cpp
+++ b/src/Vulkan/VkDescriptorUpdateTemplate.cpp
@@ -18,26 +18,28 @@
 
 namespace vk
 {
-	DescriptorUpdateTemplate::DescriptorUpdateTemplate(const VkDescriptorUpdateTemplateCreateInfo* pCreateInfo, void* mem) :

-		descriptorUpdateEntryCount(pCreateInfo->descriptorUpdateEntryCount),

-		descriptorUpdateEntries(reinterpret_cast<VkDescriptorUpdateTemplateEntry*>(mem)),

-		descriptorSetLayout(Cast(pCreateInfo->descriptorSetLayout))

-	{

-		for(uint32_t i = 0; i < descriptorUpdateEntryCount; i++)

-		{

-			descriptorUpdateEntries[i] = pCreateInfo->pDescriptorUpdateEntries[i];

-		}

-	}

-

-	size_t DescriptorUpdateTemplate::ComputeRequiredAllocationSize(const VkDescriptorUpdateTemplateCreateInfo* info)

-	{

-		return info->descriptorUpdateEntryCount * sizeof(VkDescriptorUpdateTemplateEntry);

-	}

-

-	void DescriptorUpdateTemplate::updateDescriptorSet(VkDescriptorSet descriptorSet, const void* pData)

+	DescriptorUpdateTemplate::DescriptorUpdateTemplate(const VkDescriptorUpdateTemplateCreateInfo* pCreateInfo, void* mem) :
+		descriptorUpdateEntryCount(pCreateInfo->descriptorUpdateEntryCount),
+		descriptorUpdateEntries(reinterpret_cast<VkDescriptorUpdateTemplateEntry*>(mem)),
+		descriptorSetLayout(Cast(pCreateInfo->descriptorSetLayout))
 	{
 		for(uint32_t i = 0; i < descriptorUpdateEntryCount; i++)
 		{
+			descriptorUpdateEntries[i] = pCreateInfo->pDescriptorUpdateEntries[i];
+		}
+	}
+
+	size_t DescriptorUpdateTemplate::ComputeRequiredAllocationSize(const VkDescriptorUpdateTemplateCreateInfo* info)
+	{
+		return info->descriptorUpdateEntryCount * sizeof(VkDescriptorUpdateTemplateEntry);
+	}
+
+	void DescriptorUpdateTemplate::updateDescriptorSet(VkDescriptorSet vkDescriptorSet, const void* pData)
+	{
+		DescriptorSet* descriptorSet = vk::Cast(vkDescriptorSet);
+
+		for(uint32_t i = 0; i < descriptorUpdateEntryCount; i++)
+		{
 			for(uint32_t j = 0; j < descriptorUpdateEntries[i].descriptorCount; j++)
 			{
 				const char *memToRead = (const char *)pData + descriptorUpdateEntries[i].offset + j * descriptorUpdateEntries[i].stride;
@@ -46,6 +48,6 @@
 					descriptorSet, descriptorUpdateEntries[i].dstBinding, descriptorUpdateEntries[i].dstArrayElement, 1, &typeSize);
 				memcpy(memToWrite, memToRead, typeSize);
 			}
-		}

+		}
 	}
 }
\ No newline at end of file