AO support for GIProbe (right on time for Godot Sprint!)

11 files changed, 93 insertions, 42 deletions

diff --git a/servers/visual/rasterizer.h b/servers/visual/rasterizer.h
index 51acd3a797..f35b68db47 100644
--- a/servers/visual/rasterizer.h
+++ b/servers/visual/rasterizer.h
@@ -508,6 +508,12 @@ public:
 	virtual void gi_probe_set_energy(RID p_gi_probe, float p_energy) = 0;
 	virtual float gi_probe_get_energy(RID p_gi_probe) const = 0;
 
+	virtual void gi_probe_set_ao(RID p_gi_probe, float p_ao) = 0;
+	virtual float gi_probe_get_ao(RID p_gi_probe) const = 0;
+
+	virtual void gi_probe_set_ao_strength(RID p_gi_probe, float p_strength) = 0;
+	virtual float gi_probe_get_ao_strength(RID p_gi_probe) const = 0;
+
 	virtual void gi_probe_set_bias(RID p_gi_probe, float p_bias) = 0;
 	virtual float gi_probe_get_bias(RID p_gi_probe) const = 0;
 
diff --git a/servers/visual/rasterizer_rd/rasterizer_scene_forward_rd.cpp b/servers/visual/rasterizer_rd/rasterizer_scene_forward_rd.cpp
index 867120db04..17f6bf8cdb 100644
--- a/servers/visual/rasterizer_rd/rasterizer_scene_forward_rd.cpp
+++ b/servers/visual/rasterizer_rd/rasterizer_scene_forward_rd.cpp
@@ -1555,6 +1555,8 @@ void RasterizerSceneForwardRD::_setup_gi_probes(RID *p_gi_probe_probe_cull_resul
 		gi_probe_ubo.blend_ambient = !storage->gi_probe_is_interior(base_probe);
 		gi_probe_ubo.texture_slot = gi_probe_instance_get_slot(rpi);
 		gi_probe_ubo.anisotropy_strength = storage->gi_probe_get_anisotropy_strength(base_probe);
+		gi_probe_ubo.ao = storage->gi_probe_get_ao(base_probe);
+		gi_probe_ubo.ao_strength = storage->gi_probe_get_ao_strength(base_probe);
 
 		if (gi_probe_is_anisotropic()) {
 			gi_probe_ubo.texture_slot *= 3;
diff --git a/servers/visual/rasterizer_rd/rasterizer_scene_forward_rd.h b/servers/visual/rasterizer_rd/rasterizer_scene_forward_rd.h
index 6e88a04e6d..ce335287ce 100644
--- a/servers/visual/rasterizer_rd/rasterizer_scene_forward_rd.h
+++ b/servers/visual/rasterizer_rd/rasterizer_scene_forward_rd.h
@@ -262,7 +262,9 @@ class RasterizerSceneForwardRD : public RasterizerSceneRD {
 		uint32_t texture_slot;
 
 		float anisotropy_strength;
-		uint32_t pad[3];
+		float ao;
+		float ao_strength;
+		uint32_t pad[1];
 	};
 
 	enum {
diff --git a/servers/visual/rasterizer_rd/rasterizer_scene_rd.h b/servers/visual/rasterizer_rd/rasterizer_scene_rd.h
index 4712ad92c9..c3a647ddb1 100644
--- a/servers/visual/rasterizer_rd/rasterizer_scene_rd.h
+++ b/servers/visual/rasterizer_rd/rasterizer_scene_rd.h
@@ -229,6 +229,7 @@ private:
 	bool gi_probe_use_anisotropy = false;
 	GIProbeQuality gi_probe_quality = GIPROBE_QUALITY_MEDIUM;
 	bool gi_probe_slots_dirty = true;
+
 	Vector<RID> gi_probe_slots;
 
 	enum {
diff --git a/servers/visual/rasterizer_rd/rasterizer_storage_rd.cpp b/servers/visual/rasterizer_rd/rasterizer_storage_rd.cpp
index fa4cd39049..b3f186240a 100644
--- a/servers/visual/rasterizer_rd/rasterizer_storage_rd.cpp
+++ b/servers/visual/rasterizer_rd/rasterizer_storage_rd.cpp
@@ -3771,6 +3771,32 @@ float RasterizerStorageRD::gi_probe_get_energy(RID p_gi_probe) const {
 	return gi_probe->energy;
 }
 
+void RasterizerStorageRD::gi_probe_set_ao(RID p_gi_probe, float p_ao) {
+	GIProbe *gi_probe = gi_probe_owner.getornull(p_gi_probe);
+	ERR_FAIL_COND(!gi_probe);
+
+	gi_probe->ao = p_ao;
+}
+float RasterizerStorageRD::gi_probe_get_ao(RID p_gi_probe) const {
+	GIProbe *gi_probe = gi_probe_owner.getornull(p_gi_probe);
+	ERR_FAIL_COND_V(!gi_probe, 0);
+	return gi_probe->ao;
+}
+
+void RasterizerStorageRD::gi_probe_set_ao_strength(RID p_gi_probe, float p_strength) {
+
+	GIProbe *gi_probe = gi_probe_owner.getornull(p_gi_probe);
+	ERR_FAIL_COND(!gi_probe);
+
+	gi_probe->ao_strength = p_strength;
+}
+
+float RasterizerStorageRD::gi_probe_get_ao_strength(RID p_gi_probe) const {
+	GIProbe *gi_probe = gi_probe_owner.getornull(p_gi_probe);
+	ERR_FAIL_COND_V(!gi_probe, 0);
+	return gi_probe->ao_strength;
+}
+
 void RasterizerStorageRD::gi_probe_set_bias(RID p_gi_probe, float p_bias) {
 	GIProbe *gi_probe = gi_probe_owner.getornull(p_gi_probe);
 	ERR_FAIL_COND(!gi_probe);
diff --git a/servers/visual/rasterizer_rd/rasterizer_storage_rd.h b/servers/visual/rasterizer_rd/rasterizer_storage_rd.h
index f24933a349..0708c8d071 100644
--- a/servers/visual/rasterizer_rd/rasterizer_storage_rd.h
+++ b/servers/visual/rasterizer_rd/rasterizer_storage_rd.h
@@ -423,6 +423,8 @@ private:
 
 		float dynamic_range = 4.0;
 		float energy = 1.0;
+		float ao = 0.0;
+		float ao_strength = 0.5;
 		float bias = 1.4;
 		float normal_bias = 0.0;
 		float propagation = 0.7;
@@ -987,6 +989,12 @@ public:
 	void gi_probe_set_energy(RID p_gi_probe, float p_energy);
 	float gi_probe_get_energy(RID p_gi_probe) const;
 
+	void gi_probe_set_ao(RID p_gi_probe, float p_ao);
+	float gi_probe_get_ao(RID p_gi_probe) const;
+
+	void gi_probe_set_ao_strength(RID p_gi_probe, float p_strength);
+	float gi_probe_get_ao_strength(RID p_gi_probe) const;
+
 	void gi_probe_set_bias(RID p_gi_probe, float p_bias);
 	float gi_probe_get_bias(RID p_gi_probe) const;
 
diff --git a/servers/visual/rasterizer_rd/shaders/scene_forward.glsl b/servers/visual/rasterizer_rd/shaders/scene_forward.glsl
index 1ffccdc4ee..8c122e6fcb 100644
--- a/servers/visual/rasterizer_rd/shaders/scene_forward.glsl
+++ b/servers/visual/rasterizer_rd/shaders/scene_forward.glsl
@@ -956,37 +956,7 @@ vec4 voxel_cone_trace(texture3D probe, vec3 cell_size, vec3 pos, vec3 direction,
 
 	return color;
 }
-#if 0
-vec4 voxel_cone_trace_skiplod(texture3D probe, vec3 cell_size, vec3 pos, vec3 direction, float tan_half_angle, float max_distance, float p_bias) {
 
-	float dist = p_bias;
-	vec4 color = vec4(0.0);
-	float skip_lod = 1.0;
-
-	while (dist < max_distance && color.a < 0.95) {
-		float diameter = max(1.0, 2.0 * tan_half_angle * dist);
-		vec3 uvw_pos = (pos + dist * direction) * cell_size;
-		float half_diameter = diameter * 0.5;
-		//check if outside, then break
-		if ( any(greaterThan(abs(uvw_pos - 0.5),vec3(0.5f + half_diameter * cell_size)) ) ) {
-			break;
-		}
-		vec4 scolor = textureLod(sampler3D(probe,material_samplers[SAMPLER_LINEAR_WITH_MIPMAPS_CLAMP]), uvw_pos, log2(diameter));
-		float a = (1.0 - color.a);
-		color += a * scolor;
-
-		float upper_opacity = textureLod(sampler3D(probe,material_samplers[SAMPLER_LINEAR_WITH_MIPMAPS_CLAMP]), uvw_pos, skip_lod).a;
-		float skip_factor = exp2( max( 0.0f, skip_lod * 0.5f - 1.0f ) ) * (1.0f - upper_opacity) + upper_opacity;
-
-		skip_factor = mix( skip_factor, 1.0f, min( -1.0 + upper_opacity * probeParams.vctSpecularSdfFactor + tan_half_angle * 50.0f, 1.0f ) );
-		skip_lod = clamp( skip_lod + (1.0f - upper_opacity) * 2.0f - 1.0f, 1.0f, probeParams.vctSpecSdfMaxMip );
-
-		dist += half_diameter * skip_factor;
-	}
-
-	return color;
-}
-#endif
 
 #ifndef GI_PROBE_HIGH_QUALITY
 //faster version for 45 degrees
@@ -1015,7 +985,8 @@ vec4 voxel_cone_trace_anisotropic_45_degrees(texture3D probe,texture3D aniso_pos
 		lod_level+=1.0;
 
 		float a = (1.0 - color.a);
-		color += a * scolor;
+		scolor *= a;
+		color += scolor;
 		dist += radius;
 		radius = max(0.5, tan_half_angle * dist);
 
@@ -1026,7 +997,7 @@ vec4 voxel_cone_trace_anisotropic_45_degrees(texture3D probe,texture3D aniso_pos
 }
 #else
 
-vec4 voxel_cone_trace_45_degrees(texture3D probe, vec3 cell_size, vec3 pos, vec3 direction, float tan_half_angle, float max_distance, float p_bias) {
+vec4 voxel_cone_trace_45_degrees(texture3D probe, vec3 cell_size, vec3 pos, vec3 direction, float tan_half_angle, float max_distance, float p_bias ) {
 
 	float dist = p_bias;
 	vec4 color = vec4(0.0);
@@ -1044,7 +1015,8 @@ vec4 voxel_cone_trace_45_degrees(texture3D probe, vec3 cell_size, vec3 pos, vec3
 		lod_level+=1.0;
 
 		float a = (1.0 - color.a);
-		color += a * scolor;
+		scolor *= a;
+		color += scolor;
 		dist += radius;
 		radius = max(0.5, tan_half_angle * dist);
 
@@ -1060,7 +1032,7 @@ vec4 voxel_cone_trace_45_degrees(texture3D probe, vec3 cell_size, vec3 pos, vec3
 
 
 //standard voxel cone trace
-vec4 voxel_cone_trace_anisotropic(texture3D probe,texture3D aniso_pos,texture3D aniso_neg,vec3 normal, vec3 cell_size, vec3 pos, vec3 direction, float tan_half_angle, float max_distance, float p_bias) {
+vec4 voxel_cone_trace_anisotropic(texture3D probe,texture3D aniso_pos,texture3D aniso_neg,vec3 normal, vec3 cell_size, vec3 pos, vec3 direction, float tan_half_angle, float max_distance, float p_bias ) {
 
 	float dist = p_bias;
 	vec4 color = vec4(0.0);
@@ -1081,7 +1053,8 @@ vec4 voxel_cone_trace_anisotropic(texture3D probe,texture3D aniso_pos,texture3D
 		scolor.rgb*=dot(max(vec3(0.0),(normal * aniso_pos)),vec3(1.0)) + dot(max(vec3(0.0),(-normal * aniso_neg)),vec3(1.0));
 
 		float a = (1.0 - color.a);
-		color += a * scolor;
+		scolor *= a;
+		color += scolor;
 		dist += half_diameter;
 
 	}
@@ -1155,8 +1128,8 @@ void gi_probe_compute(uint index, vec3 position, vec3 normal,vec3 ref_vec, mat3
 
 #endif
 	vec3 light = vec3(0.0);
-	for (int i = 0; i < MAX_CONE_DIRS; i++) {
 
+	for (int i = 0; i < MAX_CONE_DIRS; i++) {
 
 		vec3 dir = normalize((gi_probes.data[index].xform * vec4(normal_xform * cone_dirs[i], 0.0)).xyz);
 
@@ -1165,11 +1138,14 @@ void gi_probe_compute(uint index, vec3 position, vec3 normal,vec3 ref_vec, mat3
 #ifdef GI_PROBE_USE_ANISOTROPY
 		vec4 cone_light = voxel_cone_trace_anisotropic(gi_probe_textures[gi_probes.data[index].texture_slot],gi_probe_textures[gi_probes.data[index].texture_slot+1],gi_probe_textures[gi_probes.data[index].texture_slot+2],normalize(mix(dir,normal,gi_probes.data[index].anisotropy_strength)),cell_size, position, dir, cone_angle_tan, max_distance, gi_probes.data[index].bias);
 #else
+
 		vec4 cone_light = voxel_cone_trace(gi_probe_textures[gi_probes.data[index].texture_slot], cell_size, position, dir, cone_angle_tan, max_distance, gi_probes.data[index].bias);
+
 #endif // GI_PROBE_USE_ANISOTROPY
 
 #else
 
+
 #ifdef GI_PROBE_USE_ANISOTROPY
 		vec4 cone_light = voxel_cone_trace_anisotropic_45_degrees(gi_probe_textures[gi_probes.data[index].texture_slot],gi_probe_textures[gi_probes.data[index].texture_slot+1],gi_probe_textures[gi_probes.data[index].texture_slot+2],normalize(mix(dir,normal,gi_probes.data[index].anisotropy_strength)),cell_size, position, dir, cone_angle_tan, max_distance, gi_probes.data[index].bias);
 #else
@@ -1180,11 +1156,24 @@ void gi_probe_compute(uint index, vec3 position, vec3 normal,vec3 ref_vec, mat3
 		if (gi_probes.data[index].blend_ambient) {
 			cone_light.rgb = mix(ambient, cone_light.rgb, min(1.0, cone_light.a / 0.95));
 		}
+
 		light+=cone_weights[i] * cone_light.rgb;
+
+
 	}
 
 	light *= gi_probes.data[index].dynamic_range;
 
+	if (gi_probes.data[index].ambient_occlusion > 0.01) {
+		float ao = 0.0;
+
+		for (int i=0;i<5;i++) {
+			vec3 ofs = (position + normal * float(1<<i) * 1.5) * cell_size;
+			ao += textureLod(sampler3D(gi_probe_textures[gi_probes.data[index].texture_slot],material_samplers[SAMPLER_LINEAR_WITH_MIPMAPS_CLAMP]),ofs,float(i) ).a;
+			light *= mix(max(0.0,1.0-(ao * gi_probes.data[index].ambient_occlusion )),1.0,gi_probes.data[index].ambient_occlusion_strength);
+		}
+	}
+
 	out_diff += vec4(light * blend, blend);
 
 	//irradiance
@@ -1415,13 +1404,12 @@ FRAGMENT_SHADER_CODE
 
 			vec4 amb_accum = vec4(0.0);
 			vec4 spec_accum = vec4(0.0);
-
-			gi_probe_compute(index1, vertex, normal, ref_vec,normal_mat, roughness * roughness, ambient_light, specular_light, spec_accum, amb_accum );
+			gi_probe_compute(index1, vertex, normal, ref_vec,normal_mat, roughness * roughness, ambient_light, specular_light, spec_accum, amb_accum);
 
 			uint index2 = instances.data[instance_index].gi_offset>>16;
 
 			if (index2!=0xFFFF) {
-				gi_probe_compute(index2, vertex, normal, ref_vec,normal_mat, roughness * roughness, ambient_light, specular_light, spec_accum, amb_accum );
+				gi_probe_compute(index2, vertex, normal, ref_vec,normal_mat, roughness * roughness, ambient_light, specular_light, spec_accum, amb_accum);
 			}
 
 			if (amb_accum.a > 0.0) {
diff --git a/servers/visual/rasterizer_rd/shaders/scene_forward_inc.glsl b/servers/visual/rasterizer_rd/shaders/scene_forward_inc.glsl
index c59d5ed756..9e97f76c57 100644
--- a/servers/visual/rasterizer_rd/shaders/scene_forward_inc.glsl
+++ b/servers/visual/rasterizer_rd/shaders/scene_forward_inc.glsl
@@ -207,8 +207,8 @@ struct GIProbeData {
 	uint texture_slot;
 
 	float anisotropy_strength;
-	uint pad0;
-	uint pad1;
+	float ambient_occlusion;
+	float ambient_occlusion_strength;
 	uint pad2;
 };
 
diff --git a/servers/visual/visual_server_raster.h b/servers/visual/visual_server_raster.h
index 0edb2762fa..cd1b805448 100644
--- a/servers/visual/visual_server_raster.h
+++ b/servers/visual/visual_server_raster.h
@@ -363,6 +363,12 @@ public:
 	BIND2(gi_probe_set_energy, RID, float)
 	BIND1RC(float, gi_probe_get_energy, RID)
 
+	BIND2(gi_probe_set_ao, RID, float)
+	BIND1RC(float, gi_probe_get_ao, RID)
+
+	BIND2(gi_probe_set_ao_strength, RID, float)
+	BIND1RC(float, gi_probe_get_ao_strength, RID)
+
 	BIND2(gi_probe_set_bias, RID, float)
 	BIND1RC(float, gi_probe_get_bias, RID)
 
diff --git a/servers/visual/visual_server_wrap_mt.h b/servers/visual/visual_server_wrap_mt.h
index 29f2f20941..34e830e30f 100644
--- a/servers/visual/visual_server_wrap_mt.h
+++ b/servers/visual/visual_server_wrap_mt.h
@@ -287,6 +287,12 @@ public:
 	FUNC2(gi_probe_set_energy, RID, float)
 	FUNC1RC(float, gi_probe_get_energy, RID)
 
+	FUNC2(gi_probe_set_ao, RID, float)
+	FUNC1RC(float, gi_probe_get_ao, RID)
+
+	FUNC2(gi_probe_set_ao_strength, RID, float)
+	FUNC1RC(float, gi_probe_get_ao_strength, RID)
+
 	FUNC2(gi_probe_set_bias, RID, float)
 	FUNC1RC(float, gi_probe_get_bias, RID)
 
diff --git a/servers/visual_server.h b/servers/visual_server.h
index 9c86456804..3867797960 100644
--- a/servers/visual_server.h
+++ b/servers/visual_server.h
@@ -484,6 +484,12 @@ public:
 	virtual void gi_probe_set_energy(RID p_gi_probe, float p_energy) = 0;
 	virtual float gi_probe_get_energy(RID p_gi_probe) const = 0;
 
+	virtual void gi_probe_set_ao(RID p_gi_probe, float p_ao) = 0;
+	virtual float gi_probe_get_ao(RID p_gi_probe) const = 0;
+
+	virtual void gi_probe_set_ao_strength(RID p_gi_probe, float p_strength) = 0;
+	virtual float gi_probe_get_ao_strength(RID p_gi_probe) const = 0;
+
 	virtual void gi_probe_set_bias(RID p_gi_probe, float p_bias) = 0;
 	virtual float gi_probe_get_bias(RID p_gi_probe) const = 0;