Merge pull request #99562 from badsectoracula/improve_voxelgi_bake_ui

Add VoxelGI bake cancelling and progress UI improvement

6 files changed, 116 insertions, 36 deletions

diff --git a/editor/plugins/voxel_gi_editor_plugin.cpp b/editor/plugins/voxel_gi_editor_plugin.cpp
index 68fe013c08..527138e060 100644
--- a/editor/plugins/voxel_gi_editor_plugin.cpp
+++ b/editor/plugins/voxel_gi_editor_plugin.cpp
@@ -146,15 +146,15 @@ void VoxelGIEditorPlugin::make_visible(bool p_visible) {
 
 EditorProgress *VoxelGIEditorPlugin::tmp_progress = nullptr;
 
-void VoxelGIEditorPlugin::bake_func_begin(int p_steps) {
+void VoxelGIEditorPlugin::bake_func_begin() {
 	ERR_FAIL_COND(tmp_progress != nullptr);
 
-	tmp_progress = memnew(EditorProgress("bake_gi", TTR("Bake VoxelGI"), p_steps));
+	tmp_progress = memnew(EditorProgress("bake_gi", TTR("Bake VoxelGI"), 1000, true));
 }
 
-void VoxelGIEditorPlugin::bake_func_step(int p_step, const String &p_description) {
-	ERR_FAIL_NULL(tmp_progress);
-	tmp_progress->step(p_description, p_step, false);
+bool VoxelGIEditorPlugin::bake_func_step(int p_progress, const String &p_description) {
+	ERR_FAIL_NULL_V(tmp_progress, false);
+	return tmp_progress->step(p_description, p_progress, false);
 }
 
 void VoxelGIEditorPlugin::bake_func_end() {
diff --git a/editor/plugins/voxel_gi_editor_plugin.h b/editor/plugins/voxel_gi_editor_plugin.h
index d09822dda6..01a2ab4bd1 100644
--- a/editor/plugins/voxel_gi_editor_plugin.h
+++ b/editor/plugins/voxel_gi_editor_plugin.h
@@ -50,8 +50,8 @@ class VoxelGIEditorPlugin : public EditorPlugin {
 	EditorFileDialog *probe_file = nullptr;
 
 	static EditorProgress *tmp_progress;
-	static void bake_func_begin(int p_steps);
-	static void bake_func_step(int p_step, const String &p_description);
+	static void bake_func_begin();
+	static bool bake_func_step(int p_progress, const String &p_description);
 	static void bake_func_end();
 
 	void _bake();
diff --git a/scene/3d/voxel_gi.cpp b/scene/3d/voxel_gi.cpp
index 80ff176a98..5fec7021eb 100644
--- a/scene/3d/voxel_gi.cpp
+++ b/scene/3d/voxel_gi.cpp
@@ -389,6 +389,17 @@ VoxelGI::BakeBeginFunc VoxelGI::bake_begin_function = nullptr;
 VoxelGI::BakeStepFunc VoxelGI::bake_step_function = nullptr;
 VoxelGI::BakeEndFunc VoxelGI::bake_end_function = nullptr;
 
+static int voxelizer_plot_bake_base = 0;
+static int voxelizer_plot_bake_total = 0;
+
+static bool voxelizer_plot_bake_step_function(int current, int) {
+	return VoxelGI::bake_step_function((voxelizer_plot_bake_base + current) * 500 / voxelizer_plot_bake_total, RTR("Plotting Meshes"));
+}
+
+static bool voxelizer_sdf_bake_step_function(int current, int total) {
+	return VoxelGI::bake_step_function(500 + current * 500 / total, RTR("Generating Distance Field"));
+}
+
 Vector3i VoxelGI::get_estimated_cell_size() const {
 	static const int subdiv_value[SUBDIV_MAX] = { 6, 7, 8, 9 };
 	int cell_subdiv = subdiv_value[subdiv];
@@ -432,22 +443,27 @@ void VoxelGI::bake(Node *p_from_node, bool p_create_visual_debug) {
 	_find_meshes(p_from_node, mesh_list);
 
 	if (bake_begin_function) {
-		bake_begin_function(mesh_list.size() + 1);
+		bake_begin_function();
 	}
 
-	int pmc = 0;
+	Voxelizer::BakeStepFunc voxelizer_step_func = bake_step_function != nullptr ? voxelizer_plot_bake_step_function : nullptr;
 
+	voxelizer_plot_bake_total = voxelizer_plot_bake_base = 0;
 	for (PlotMesh &E : mesh_list) {
-		if (bake_step_function) {
-			bake_step_function(pmc, RTR("Plotting Meshes") + " " + itos(pmc) + "/" + itos(mesh_list.size()));
+		voxelizer_plot_bake_total += baker.get_bake_steps(E.mesh);
+	}
+	for (PlotMesh &E : mesh_list) {
+		if (baker.plot_mesh(E.local_xform, E.mesh, E.instance_materials, E.override_material, voxelizer_step_func) != Voxelizer::BAKE_RESULT_OK) {
+			baker.end_bake();
+			if (bake_end_function) {
+				bake_end_function();
+			}
+			return;
 		}
-
-		pmc++;
-
-		baker.plot_mesh(E.local_xform, E.mesh, E.instance_materials, E.override_material);
+		voxelizer_plot_bake_base += baker.get_bake_steps(E.mesh);
 	}
 	if (bake_step_function) {
-		bake_step_function(pmc++, RTR("Finishing Plot"));
+		bake_step_function(500, RTR("Finishing Plot"));
 	}
 
 	baker.end_bake();
@@ -476,19 +492,22 @@ void VoxelGI::bake(Node *p_from_node, bool p_create_visual_debug) {
 		}
 
 		if (bake_step_function) {
-			bake_step_function(pmc++, RTR("Generating Distance Field"));
+			bake_step_function(500, RTR("Generating Distance Field"));
 		}
 
-		Vector<uint8_t> df = baker.get_sdf_3d_image();
+		voxelizer_step_func = bake_step_function != nullptr ? voxelizer_sdf_bake_step_function : nullptr;
 
-		RS::get_singleton()->voxel_gi_set_baked_exposure_normalization(probe_data_new->get_rid(), exposure_normalization);
+		Vector<uint8_t> df;
+		if (baker.get_sdf_3d_image(df, voxelizer_step_func) == Voxelizer::BAKE_RESULT_OK) {
+			RS::get_singleton()->voxel_gi_set_baked_exposure_normalization(probe_data_new->get_rid(), exposure_normalization);
 
-		probe_data_new->allocate(baker.get_to_cell_space_xform(), AABB(-size / 2, size), baker.get_voxel_gi_octree_size(), baker.get_voxel_gi_octree_cells(), baker.get_voxel_gi_data_cells(), df, baker.get_voxel_gi_level_cell_count());
+			probe_data_new->allocate(baker.get_to_cell_space_xform(), AABB(-size / 2, size), baker.get_voxel_gi_octree_size(), baker.get_voxel_gi_octree_cells(), baker.get_voxel_gi_data_cells(), df, baker.get_voxel_gi_level_cell_count());
 
-		set_probe_data(probe_data_new);
+			set_probe_data(probe_data_new);
 #ifdef TOOLS_ENABLED
-		probe_data_new->set_edited(true); //so it gets saved
+			probe_data_new->set_edited(true); //so it gets saved
 #endif
+		}
 	}
 
 	if (bake_end_function) {
diff --git a/scene/3d/voxel_gi.h b/scene/3d/voxel_gi.h
index 7d7787f721..d7e0d74d36 100644
--- a/scene/3d/voxel_gi.h
+++ b/scene/3d/voxel_gi.h
@@ -108,8 +108,8 @@ public:
 
 	};
 
-	typedef void (*BakeBeginFunc)(int);
-	typedef void (*BakeStepFunc)(int, const String &);
+	typedef void (*BakeBeginFunc)();
+	typedef bool (*BakeStepFunc)(int, const String &);
 	typedef void (*BakeEndFunc)();
 
 private:
diff --git a/scene/3d/voxelizer.cpp b/scene/3d/voxelizer.cpp
index 99392e9ba0..1074cad11e 100644
--- a/scene/3d/voxelizer.cpp
+++ b/scene/3d/voxelizer.cpp
@@ -382,8 +382,24 @@ Voxelizer::MaterialCache Voxelizer::_get_material_cache(Ref<Material> p_material
 	return mc;
 }
 
-void Voxelizer::plot_mesh(const Transform3D &p_xform, Ref<Mesh> &p_mesh, const Vector<Ref<Material>> &p_materials, const Ref<Material> &p_override_material) {
-	ERR_FAIL_COND_MSG(!p_xform.is_finite(), "Invalid mesh bake transform.");
+int Voxelizer::get_bake_steps(Ref<Mesh> &p_mesh) const {
+	int bake_total = 0;
+	for (int i = 0; i < p_mesh->get_surface_count(); i++) {
+		if (p_mesh->surface_get_primitive_type(i) != Mesh::PRIMITIVE_TRIANGLES) {
+			continue; // Only triangles.
+		}
+		Array a = p_mesh->surface_get_arrays(i);
+		Vector<Vector3> vertices = a[Mesh::ARRAY_VERTEX];
+		Vector<int> index = a[Mesh::ARRAY_INDEX];
+		bake_total += (index.size() > 0 ? index.size() : vertices.size()) / 3;
+	}
+	return bake_total;
+}
+
+Voxelizer::BakeResult Voxelizer::plot_mesh(const Transform3D &p_xform, Ref<Mesh> &p_mesh, const Vector<Ref<Material>> &p_materials, const Ref<Material> &p_override_material, BakeStepFunc p_bake_step_func) {
+	ERR_FAIL_COND_V_MSG(!p_xform.is_finite(), BAKE_RESULT_INVALID_PARAMETER, "Invalid mesh bake transform.");
+
+	int bake_total = get_bake_steps(p_mesh), bake_current = 0;
 
 	for (int i = 0; i < p_mesh->get_surface_count(); i++) {
 		if (p_mesh->surface_get_primitive_type(i) != Mesh::PRIMITIVE_TRIANGLES) {
@@ -428,6 +444,13 @@ void Voxelizer::plot_mesh(const Transform3D &p_xform, Ref<Mesh> &p_mesh, const V
 				Vector2 uvs[3];
 				Vector3 normal[3];
 
+				bake_current++;
+				if (p_bake_step_func != nullptr && (bake_current & 2047) == 1) {
+					if (p_bake_step_func(bake_current, bake_total)) {
+						return BAKE_RESULT_CANCELLED;
+					}
+				}
+
 				for (int k = 0; k < 3; k++) {
 					vtxs[k] = p_xform.xform(vr[ir[j * 3 + k]]);
 				}
@@ -460,6 +483,13 @@ void Voxelizer::plot_mesh(const Transform3D &p_xform, Ref<Mesh> &p_mesh, const V
 				Vector2 uvs[3];
 				Vector3 normal[3];
 
+				bake_current++;
+				if (p_bake_step_func != nullptr && (bake_current & 2047) == 1) {
+					if (p_bake_step_func(bake_current, bake_total)) {
+						return BAKE_RESULT_CANCELLED;
+					}
+				}
+
 				for (int k = 0; k < 3; k++) {
 					vtxs[k] = p_xform.xform(vr[j * 3 + k]);
 				}
@@ -487,6 +517,8 @@ void Voxelizer::plot_mesh(const Transform3D &p_xform, Ref<Mesh> &p_mesh, const V
 	}
 
 	max_original_cells = bake_cells.size();
+
+	return BAKE_RESULT_OK;
 }
 
 void Voxelizer::_sort() {
@@ -821,7 +853,7 @@ static void edt(float *f, int stride, int n) {
 
 #undef square
 
-Vector<uint8_t> Voxelizer::get_sdf_3d_image() const {
+Voxelizer::BakeResult Voxelizer::get_sdf_3d_image(Vector<uint8_t> &r_image, BakeStepFunc p_bake_step_function) const {
 	Vector3i octree_size = get_voxel_gi_octree_size();
 
 	uint32_t float_count = octree_size.x * octree_size.y * octree_size.z;
@@ -849,9 +881,17 @@ Vector<uint8_t> Voxelizer::get_sdf_3d_image() const {
 
 	//process in each direction
 
+	int bake_total = octree_size.x * 2 + octree_size.y, bake_current = 0;
+
 	//xy->z
 
-	for (int i = 0; i < octree_size.x; i++) {
+	for (int i = 0; i < octree_size.x; i++, bake_current++) {
+		if (p_bake_step_function) {
+			if (p_bake_step_function(bake_current, bake_total)) {
+				memdelete_arr(work_memory);
+				return BAKE_RESULT_CANCELLED;
+			}
+		}
 		for (int j = 0; j < octree_size.y; j++) {
 			edt(&work_memory[i + j * y_mult], z_mult, octree_size.z);
 		}
@@ -859,23 +899,34 @@ Vector<uint8_t> Voxelizer::get_sdf_3d_image() const {
 
 	//xz->y
 
-	for (int i = 0; i < octree_size.x; i++) {
+	for (int i = 0; i < octree_size.x; i++, bake_current++) {
+		if (p_bake_step_function) {
+			if (p_bake_step_function(bake_current, bake_total)) {
+				memdelete_arr(work_memory);
+				return BAKE_RESULT_CANCELLED;
+			}
+		}
 		for (int j = 0; j < octree_size.z; j++) {
 			edt(&work_memory[i + j * z_mult], y_mult, octree_size.y);
 		}
 	}
 
 	//yz->x
-	for (int i = 0; i < octree_size.y; i++) {
+	for (int i = 0; i < octree_size.y; i++, bake_current++) {
+		if (p_bake_step_function) {
+			if (p_bake_step_function(bake_current, bake_total)) {
+				memdelete_arr(work_memory);
+				return BAKE_RESULT_CANCELLED;
+			}
+		}
 		for (int j = 0; j < octree_size.z; j++) {
 			edt(&work_memory[i * y_mult + j * z_mult], 1, octree_size.x);
 		}
 	}
 
-	Vector<uint8_t> image3d;
-	image3d.resize(float_count);
+	r_image.resize(float_count);
 	{
-		uint8_t *w = image3d.ptrw();
+		uint8_t *w = r_image.ptrw();
 		for (uint32_t i = 0; i < float_count; i++) {
 			uint32_t d = uint32_t(Math::sqrt(work_memory[i]));
 			if (d == 0) {
@@ -888,7 +939,7 @@ Vector<uint8_t> Voxelizer::get_sdf_3d_image() const {
 
 	memdelete_arr(work_memory);
 
-	return image3d;
+	return BAKE_RESULT_OK;
 }
 
 #undef INF
diff --git a/scene/3d/voxelizer.h b/scene/3d/voxelizer.h
index 08d018eee9..41e7767308 100644
--- a/scene/3d/voxelizer.h
+++ b/scene/3d/voxelizer.h
@@ -34,6 +34,15 @@
 #include "scene/resources/multimesh.h"
 
 class Voxelizer {
+public:
+	enum BakeResult {
+		BAKE_RESULT_OK,
+		BAKE_RESULT_INVALID_PARAMETER,
+		BAKE_RESULT_CANCELLED,
+	};
+
+	typedef bool (*BakeStepFunc)(int, int);
+
 private:
 	enum : uint32_t {
 		CHILD_EMPTY = 0xFFFFFFFF
@@ -112,7 +121,8 @@ private:
 
 public:
 	void begin_bake(int p_subdiv, const AABB &p_bounds, float p_exposure_normalization);
-	void plot_mesh(const Transform3D &p_xform, Ref<Mesh> &p_mesh, const Vector<Ref<Material>> &p_materials, const Ref<Material> &p_override_material);
+	int get_bake_steps(Ref<Mesh> &p_mesh) const;
+	BakeResult plot_mesh(const Transform3D &p_xform, Ref<Mesh> &p_mesh, const Vector<Ref<Material>> &p_materials, const Ref<Material> &p_override_material, BakeStepFunc p_bake_step_function);
 	void end_bake();
 
 	int get_voxel_gi_octree_depth() const;
@@ -121,7 +131,7 @@ public:
 	Vector<uint8_t> get_voxel_gi_octree_cells() const;
 	Vector<uint8_t> get_voxel_gi_data_cells() const;
 	Vector<int> get_voxel_gi_level_cell_count() const;
-	Vector<uint8_t> get_sdf_3d_image() const;
+	BakeResult get_sdf_3d_image(Vector<uint8_t> &r_image, BakeStepFunc p_bake_step_function) const;
 
 	Ref<MultiMesh> create_debug_multimesh();