From bc0f419fffedde7ea896e6bf2a4d98cf5c9618b0 Mon Sep 17 00:00:00 2001
From: Magnus Norddahl <dpjudas@users.noreply.github.com>
Date: Sun, 3 Jul 2022 13:22:17 +0200
Subject: [PATCH] Add ambient occlusion

---
 src/lightmap/glsl_frag.h       | 68 ++++++++++++++++++++++++++++++++++
 src/lightmap/gpuraytracer2.cpp | 19 ++++++++--
 2 files changed, 84 insertions(+), 3 deletions(-)

diff --git a/src/lightmap/glsl_frag.h b/src/lightmap/glsl_frag.h
index 99a45da..b59f731 100644
--- a/src/lightmap/glsl_frag.h
+++ b/src/lightmap/glsl_frag.h
@@ -61,6 +61,9 @@ layout(push_constant) uniform PushConstants
 layout(location = 0) in vec3 worldpos;
 layout(location = 0) out vec4 fragcolor;
 
+vec2 Hammersley(uint i, uint N);
+float RadicalInverse_VdC(uint bits);
+
 void main()
 {
 	const float minDistance = 0.01;
@@ -137,7 +140,72 @@ void main()
 		}
 	}
 
+	// Ambient occlusion
+	if (SurfaceIndex >= 0)
+	{
+		const float minDistance = 0.05;
+		const float aoDistance = 100;
+		const int SampleCount = 2048;
+
+		vec3 N = normal;
+		vec3 up = abs(N.x) < abs(N.y) ? vec3(1.0, 0.0, 0.0) : vec3(0.0, 1.0, 0.0);
+		vec3 tangent = normalize(cross(up, N));
+		vec3 bitangent = cross(N, tangent);
+
+		float ambience = 0.0f;
+		for (uint i = 0; i < SampleCount; i++)
+		{
+			vec2 Xi = Hammersley(i, SampleCount);
+			vec3 H = normalize(vec3(Xi.x * 2.0f - 1.0f, Xi.y * 2.0f - 1.0f, 1.5 - length(Xi)));
+			vec3 L = H.x * tangent + H.y * bitangent + H.z * N;
+
+			rayQueryEXT rayQuery;
+			rayQueryInitializeEXT(rayQuery, acc, gl_RayFlagsTerminateOnFirstHitEXT, 0xFF, origin, minDistance, L, aoDistance);
+
+			while(rayQueryProceedEXT(rayQuery))
+			{
+				if (rayQueryGetIntersectionTypeEXT(rayQuery, false) == gl_RayQueryCommittedIntersectionTriangleEXT)
+				{
+					rayQueryConfirmIntersectionEXT(rayQuery);
+				}
+			}
+
+			if (rayQueryGetIntersectionTypeEXT(rayQuery, true) == gl_RayQueryCommittedIntersectionTriangleEXT)
+			{
+				int primitiveID = rayQueryGetIntersectionPrimitiveIndexEXT(rayQuery, true);
+				SurfaceInfo surface = surfaces[surfaceIndices[primitiveID]];
+				if (surface.Sky == 0.0)
+				{
+					float hitDistance = rayQueryGetIntersectionTEXT(rayQuery, true);
+					ambience += clamp(hitDistance / aoDistance, 0.0, 1.0);
+				}
+			}
+			else
+			{
+				ambience += 1.0;
+			}
+		}
+		ambience /= float(SampleCount);
+
+		incoming.rgb = incoming.rgb * ambience;
+	}
+
 	fragcolor = vec4(incoming, 1.0);
 }
 
+vec2 Hammersley(uint i, uint N)
+{
+	return vec2(float(i) / float(N), RadicalInverse_VdC(i));
+}
+
+float RadicalInverse_VdC(uint bits)
+{
+	bits = (bits << 16u) | (bits >> 16u);
+	bits = ((bits & 0x55555555u) << 1u) | ((bits & 0xAAAAAAAAu) >> 1u);
+	bits = ((bits & 0x33333333u) << 2u) | ((bits & 0xCCCCCCCCu) >> 2u);
+	bits = ((bits & 0x0F0F0F0Fu) << 4u) | ((bits & 0xF0F0F0F0u) >> 4u);
+	bits = ((bits & 0x00FF00FFu) << 8u) | ((bits & 0xFF00FF00u) >> 8u);
+	return float(bits) * 2.3283064365386963e-10f; // / 0x100000000
+}
+
 )glsl";
diff --git a/src/lightmap/gpuraytracer2.cpp b/src/lightmap/gpuraytracer2.cpp
index 41db2fb..2369c23 100644
--- a/src/lightmap/gpuraytracer2.cpp
+++ b/src/lightmap/gpuraytracer2.cpp
@@ -65,8 +65,15 @@ void GPURaytracer2::Raytrace(LevelMesh* level)
 		Surface* surface = mesh->surfaces[i].get();
 		int sampleWidth = surface->lightmapDims[0];
 		int sampleHeight = surface->lightmapDims[1];
+		surfaceImages.push_back(CreateImage(sampleWidth, sampleHeight));
+	}
 
-		LightmapImage img = CreateImage(sampleWidth, sampleHeight);
+	for (size_t i = 0; i < mesh->surfaces.size(); i++)
+	{
+		Surface* surface = mesh->surfaces[i].get();
+		int sampleWidth = surface->lightmapDims[0];
+		int sampleHeight = surface->lightmapDims[1];
+		LightmapImage& img = surfaceImages[i];
 
 		RenderPassBegin()
 			.RenderPass(renderPass.get())
@@ -96,6 +103,14 @@ void GPURaytracer2::Raytrace(LevelMesh* level)
 		cmdbuffer->draw(4, 1, 0, 0);
 
 		cmdbuffer->endRenderPass();
+	}
+
+	for (size_t i = 0; i < mesh->surfaces.size(); i++)
+	{
+		Surface* surface = mesh->surfaces[i].get();
+		int sampleWidth = surface->lightmapDims[0];
+		int sampleHeight = surface->lightmapDims[1];
+		LightmapImage& img = surfaceImages[i];
 
 		PipelineBarrier()
 			.AddImage(img.Image.get(), VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT)
@@ -108,8 +123,6 @@ void GPURaytracer2::Raytrace(LevelMesh* level)
 		region.imageExtent.height = sampleHeight;
 		region.imageExtent.depth = 1;
 		cmdbuffer->copyImageToBuffer(img.Image->image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, img.Transfer->buffer, 1, &region);
-
-		surfaceImages.push_back(std::move(img));
 	}
 
 	FinishCommands();