From 02c6cf2f3bbb0c3542375cf742918877db42f31e Mon Sep 17 00:00:00 2001
From: Mikko Rasa <tdb@tdb.fi>
Date: Sun, 9 Oct 2022 12:05:50 +0300
Subject: [PATCH] Add shader functions to get more exact environment samples

The environment map depth buffer is now stored in a cube texture and
passed to the shader so it's possible to raymarch through it.
---
 shaderlib/environment.glsl        | 63 +++++++++++++++++++++++++++++--
 source/effects/environmentmap.cpp | 18 +++++++--
 source/effects/environmentmap.h   |  3 +-
 3 files changed, 75 insertions(+), 9 deletions(-)

diff --git a/shaderlib/environment.glsl b/shaderlib/environment.glsl
index b5f0cdfb..0b619d0e 100644
--- a/shaderlib/environment.glsl
+++ b/shaderlib/environment.glsl
@@ -1,16 +1,67 @@
 layout(set=0) uniform EnvMap
 {
-	mat3 env_world_matrix;
+	mat4 env_world_matrix;
+	mat4 world_env_matrix;
+	mat4 env_projection_matrix;
+	mat4 env_projection_inverse;
 };
 
 layout(set=0) uniform samplerCube environment_map;
+layout(set=0) uniform samplerCube environment_depth;
 layout(set=0) uniform samplerCube irradiance_map;
 
 #pragma MSP stage(fragment)
 virtual vec3 get_environment_sample(vec3 direction, float roughness)
 {
 	float lod = (2-roughness)*roughness*(textureQueryLevels(environment_map)-1);
-	return textureLod(environment_map, env_world_matrix*direction, lod).rgb;
+	return textureLod(environment_map, mat3(env_world_matrix)*direction, lod).rgb;
+}
+
+bool is_env_obscured(vec3 pos)
+{
+	// Note: column-major notation
+	float proj22 = env_projection_matrix[2][2];
+	float proj32 = env_projection_matrix[3][2];
+
+	vec3 abs_pos = abs(pos);
+	float depth = max(max(abs_pos.x, abs_pos.y), abs_pos.z);
+	float env_depth = proj32/(texture(environment_depth, pos).r+proj22);
+	return env_depth<depth;
+}
+
+virtual vec3 get_environment_sample(vec3 position, vec3 direction, float roughness)
+{
+	vec3 local_pos = (env_world_matrix*vec4(position, 1.0)).xyz;
+	if(is_env_obscured(local_pos))
+		return textureLod(environment_map, local_pos, roughness).rgb;
+
+	vec3 local_dir = mat3(env_world_matrix)*direction;
+	float max_arc = 0.06;
+	int max_steps = int(3.14159/max_arc);
+
+	for(int i=0; i<max_steps; ++i)
+	{
+		vec3 prev_pos = local_pos;
+		float dist_sq = dot(local_pos, local_pos);
+		vec3 tangent = local_dir-local_pos*(dot(local_dir, local_pos)/dist_sq);
+		float arc = sqrt(dot(tangent, tangent)/dist_sq);
+		local_pos += local_dir*(max_arc/arc);
+
+		if(is_env_obscured(local_pos))
+		{
+			for(int j=0; j<5; ++j)
+			{
+				vec3 mid_pos = (prev_pos+local_pos)/2.0;
+				if(is_env_obscured(mid_pos))
+					local_pos = mid_pos;
+				else
+					prev_pos = mid_pos;
+			}
+			return textureLod(environment_map, local_pos, roughness).rgb;
+		}
+	}
+
+	return textureLod(environment_map, local_dir, roughness).rgb;
 }
 
 virtual vec3 get_reflection(vec3 normal, vec3 look)
@@ -19,8 +70,14 @@ virtual vec3 get_reflection(vec3 normal, vec3 look)
 	return get_environment_sample(reflect_dir, 0.0);
 }
 
+virtual vec3 get_reflection(vec3 position, vec3 normal, vec3 look)
+{
+	vec3 reflect_dir = reflect(look, normal);
+	return get_environment_sample(position, reflect_dir, 0.0);
+}
+
 virtual vec3 get_irradiance_sample(vec3 normal)
 {
-	return texture(irradiance_map, env_world_matrix*normal).rgb;
+	return texture(irradiance_map, mat3(env_world_matrix)*normal).rgb;
 }
 
diff --git a/source/effects/environmentmap.cpp b/source/effects/environmentmap.cpp
index a59c1431..a6cb6a36 100644
--- a/source/effects/environmentmap.cpp
+++ b/source/effects/environmentmap.cpp
@@ -31,20 +31,21 @@ EnvironmentMap::EnvironmentMap(unsigned s, PixelFormat f, unsigned l, Renderable
 		throw invalid_argument("EnvironmentMap::EnvironmentMap");
 
 	env_tex.storage(f, size, l);
-	depth_buf.storage(DEPTH_COMPONENT32F, size, size, 1);
+	depth_buf.storage(DEPTH_COMPONENT32F, size, 1);
 	for(unsigned i=0; i<6; ++i)
 	{
 		TextureCubeFace face = static_cast<TextureCubeFace>(i);
 		faces[i].fbo.set_format((COLOR_ATTACHMENT,f, DEPTH_ATTACHMENT,DEPTH_COMPONENT32F));
 		faces[i].fbo.attach(COLOR_ATTACHMENT, env_tex, face, 0);
-		faces[i].fbo.attach(DEPTH_ATTACHMENT, depth_buf);
+		faces[i].fbo.attach(DEPTH_ATTACHMENT, depth_buf, face, 0);
 		faces[i].camera.set_look_direction(TextureCube::get_face_direction(face));
 		faces[i].camera.set_up_direction(TextureCube::get_t_direction(face));
 		faces[i].camera.set_field_of_view(Geometry::Angle<float>::right());
 		faces[i].camera.set_aspect_ratio(1);
-		faces[i].camera.set_depth_clip(0.1, 100);
 	}
 
+	set_depth_clip(0.1f, 100.0f);
+
 	irradiance.storage(f, size/4, 1);
 	irradiance_fbo.set_format((COLOR_ATTACHMENT,f));
 	irradiance_fbo.attach_layered(COLOR_ATTACHMENT, irradiance);
@@ -74,7 +75,8 @@ EnvironmentMap::EnvironmentMap(unsigned s, PixelFormat f, unsigned l, Renderable
 		prefilter_shdata.uniform("roughness", 1.0f);
 	}
 
-	shdata.uniform("env_world_matrix", LinAl::Matrix<float, 3, 3>::identity());
+	shdata.uniform("env_world_matrix", Matrix::identity());
+	shdata.uniform("world_env_matrix", Matrix::identity());
 }
 
 void EnvironmentMap::set_fixed_position(const Vector3 &p)
@@ -87,6 +89,10 @@ void EnvironmentMap::set_depth_clip(float n, float f)
 {
 	for(unsigned i=0; i<6; ++i)
 		faces[i].camera.set_depth_clip(n, f);
+
+	const Matrix &proj_matrix = faces[0].camera.get_projection_matrix();
+	shdata.uniform("env_projection_matrix", proj_matrix);
+	shdata.uniform("env_projection_inverse", invert(proj_matrix));
 }
 
 void EnvironmentMap::set_update_interval(unsigned i)
@@ -131,6 +137,9 @@ void EnvironmentMap::setup_frame(Renderer &renderer)
 		center = matrix->column(3).slice<3>(0);
 	}
 
+	shdata.uniform("world_env_matrix", Matrix::translation(center));
+	shdata.uniform("env_world_matrix", Matrix::translation(-center));
+
 	Renderer::Push push(renderer);
 
 	for(unsigned i=0; i<6; ++i)
@@ -177,6 +186,7 @@ void EnvironmentMap::render(Renderer &renderer, Tag tag) const
 	Renderer::Push _push_rend(renderer);
 
 	renderer.set_texture("environment_map", &env_tex, &mip_sampler);
+	renderer.set_texture("environment_depth", &depth_buf, &sampler);
 	renderer.set_texture("irradiance_map", &irradiance, &sampler);
 	renderer.add_shader_data(shdata);
 	content.render(renderer, tag);
diff --git a/source/effects/environmentmap.h b/source/effects/environmentmap.h
index 9ef0e87d..e6ce5f81 100644
--- a/source/effects/environmentmap.h
+++ b/source/effects/environmentmap.h
@@ -6,7 +6,6 @@
 #include "effect.h"
 #include "framebuffer.h"
 #include "programdata.h"
-#include "texture2d.h"
 #include "texturecube.h"
 
 namespace Msp {
@@ -68,7 +67,7 @@ protected:
 	unsigned size;
 	Renderable &environment;
 	TextureCube env_tex;
-	Texture2D depth_buf;
+	TextureCube depth_buf;
 	Face faces[6];
 	Vector3 fixed_position;
 	bool use_fixed_pos = false;
-- 
2.45.2