Fix SoftDynamicBody3D memory corruption when switching mesh at runtime

When switching the mesh at runtime, the physics server wasn't properly
updated with the new mesh. Now we keep track of the soft body mesh to
make sure everything is properly initialized on pre-draw.

Also cleaned a few things around private methods.

2 files changed, 50 insertions, 47 deletions

diff --git a/scene/3d/soft_dynamic_body_3d.cpp b/scene/3d/soft_dynamic_body_3d.cpp
index 9fceb21790..43292d42f1 100644
--- a/scene/3d/soft_dynamic_body_3d.cpp
+++ b/scene/3d/soft_dynamic_body_3d.cpp
@@ -250,7 +250,7 @@ void SoftDynamicBody3D::_notification(int p_what) {
 
 			RID space = get_world_3d()->get_space();
 			PhysicsServer3D::get_singleton()->soft_body_set_space(physics_rid, space);
-			prepare_physics_server();
+			_prepare_physics_server();
 		} break;
 
 		case NOTIFICATION_READY: {
@@ -284,13 +284,13 @@ void SoftDynamicBody3D::_notification(int p_what) {
 
 		case NOTIFICATION_DISABLED: {
 			if (is_inside_tree() && (disable_mode == DISABLE_MODE_REMOVE)) {
-				prepare_physics_server();
+				_prepare_physics_server();
 			}
 		} break;
 
 		case NOTIFICATION_ENABLED: {
 			if (is_inside_tree() && (disable_mode == DISABLE_MODE_REMOVE)) {
-				prepare_physics_server();
+				_prepare_physics_server();
 			}
 		} break;
 
@@ -378,7 +378,7 @@ void SoftDynamicBody3D::_bind_methods() {
 TypedArray<String> SoftDynamicBody3D::get_configuration_warnings() const {
 	TypedArray<String> warnings = Node::get_configuration_warnings();
 
-	if (get_mesh().is_null()) {
+	if (mesh.is_null()) {
 		warnings.push_back(TTR("This body will be ignored until you set a mesh."));
 	}
 
@@ -407,11 +407,17 @@ void SoftDynamicBody3D::_update_physics_server() {
 }
 
 void SoftDynamicBody3D::_draw_soft_mesh() {
-	if (get_mesh().is_null()) {
+	if (mesh.is_null()) {
 		return;
 	}
 
-	const RID mesh_rid = get_mesh()->get_rid();
+	RID mesh_rid = mesh->get_rid();
+	if (owned_mesh != mesh_rid) {
+		_become_mesh_owner();
+		mesh_rid = mesh->get_rid();
+		PhysicsServer3D::get_singleton()->soft_body_set_mesh(physics_rid, mesh_rid);
+	}
+
 	if (!rendering_server_handler.is_ready(mesh_rid)) {
 		rendering_server_handler.prepare(mesh_rid, 0);
 
@@ -430,11 +436,11 @@ void SoftDynamicBody3D::_draw_soft_mesh() {
 	rendering_server_handler.commit_changes();
 }
 
-void SoftDynamicBody3D::prepare_physics_server() {
+void SoftDynamicBody3D::_prepare_physics_server() {
 #ifdef TOOLS_ENABLED
 	if (Engine::get_singleton()->is_editor_hint()) {
-		if (get_mesh().is_valid()) {
-			PhysicsServer3D::get_singleton()->soft_body_set_mesh(physics_rid, get_mesh()->get_rid());
+		if (mesh.is_valid()) {
+			PhysicsServer3D::get_singleton()->soft_body_set_mesh(physics_rid, mesh->get_rid());
 		} else {
 			PhysicsServer3D::get_singleton()->soft_body_set_mesh(physics_rid, RID());
 		}
@@ -443,9 +449,13 @@ void SoftDynamicBody3D::prepare_physics_server() {
 	}
 #endif
 
-	if (get_mesh().is_valid() && (is_enabled() || (disable_mode != DISABLE_MODE_REMOVE))) {
-		become_mesh_owner();
-		PhysicsServer3D::get_singleton()->soft_body_set_mesh(physics_rid, get_mesh()->get_rid());
+	if (mesh.is_valid() && (is_enabled() || (disable_mode != DISABLE_MODE_REMOVE))) {
+		RID mesh_rid = mesh->get_rid();
+		if (owned_mesh != mesh_rid) {
+			_become_mesh_owner();
+			mesh_rid = mesh->get_rid();
+		}
+		PhysicsServer3D::get_singleton()->soft_body_set_mesh(physics_rid, mesh_rid);
 		RS::get_singleton()->connect("frame_pre_draw", callable_mp(this, &SoftDynamicBody3D::_draw_soft_mesh));
 	} else {
 		PhysicsServer3D::get_singleton()->soft_body_set_mesh(physics_rid, RID());
@@ -455,38 +465,32 @@ void SoftDynamicBody3D::prepare_physics_server() {
 	}
 }
 
-void SoftDynamicBody3D::become_mesh_owner() {
-	if (mesh.is_null()) {
-		return;
-	}
-
-	if (!mesh_owner) {
-		mesh_owner = true;
+void SoftDynamicBody3D::_become_mesh_owner() {
+	Vector<Ref<Material>> copy_materials;
+	copy_materials.append_array(surface_override_materials);
 
-		Vector<Ref<Material>> copy_materials;
-		copy_materials.append_array(surface_override_materials);
+	ERR_FAIL_COND(!mesh->get_surface_count());
 
-		ERR_FAIL_COND(!mesh->get_surface_count());
+	// Get current mesh array and create new mesh array with necessary flag for SoftDynamicBody
+	Array surface_arrays = mesh->surface_get_arrays(0);
+	Array surface_blend_arrays = mesh->surface_get_blend_shape_arrays(0);
+	Dictionary surface_lods = mesh->surface_get_lods(0);
+	uint32_t surface_format = mesh->surface_get_format(0);
 
-		// Get current mesh array and create new mesh array with necessary flag for SoftDynamicBody
-		Array surface_arrays = mesh->surface_get_arrays(0);
-		Array surface_blend_arrays = mesh->surface_get_blend_shape_arrays(0);
-		Dictionary surface_lods = mesh->surface_get_lods(0);
-		uint32_t surface_format = mesh->surface_get_format(0);
+	surface_format |= Mesh::ARRAY_FLAG_USE_DYNAMIC_UPDATE;
 
-		surface_format |= Mesh::ARRAY_FLAG_USE_DYNAMIC_UPDATE;
+	Ref<ArrayMesh> soft_mesh;
+	soft_mesh.instantiate();
+	soft_mesh->add_surface_from_arrays(Mesh::PRIMITIVE_TRIANGLES, surface_arrays, surface_blend_arrays, surface_lods, surface_format);
+	soft_mesh->surface_set_material(0, mesh->surface_get_material(0));
 
-		Ref<ArrayMesh> soft_mesh;
-		soft_mesh.instantiate();
-		soft_mesh->add_surface_from_arrays(Mesh::PRIMITIVE_TRIANGLES, surface_arrays, surface_blend_arrays, surface_lods, surface_format);
-		soft_mesh->surface_set_material(0, mesh->surface_get_material(0));
+	set_mesh(soft_mesh);
 
-		set_mesh(soft_mesh);
-
-		for (int i = copy_materials.size() - 1; 0 <= i; --i) {
-			set_surface_override_material(i, copy_materials[i]);
-		}
+	for (int i = copy_materials.size() - 1; 0 <= i; --i) {
+		set_surface_override_material(i, copy_materials[i]);
 	}
+
+	owned_mesh = soft_mesh->get_rid();
 }
 
 void SoftDynamicBody3D::set_collision_mask(uint32_t p_mask) {
@@ -551,13 +555,13 @@ void SoftDynamicBody3D::set_disable_mode(DisableMode p_mode) {
 	bool inside_tree = is_inside_tree();
 
 	if (inside_tree && (disable_mode == DISABLE_MODE_REMOVE)) {
-		prepare_physics_server();
+		_prepare_physics_server();
 	}
 
 	disable_mode = p_mode;
 
 	if (inside_tree && (disable_mode == DISABLE_MODE_REMOVE)) {
-		prepare_physics_server();
+		_prepare_physics_server();
 	}
 }
 
diff --git a/scene/3d/soft_dynamic_body_3d.h b/scene/3d/soft_dynamic_body_3d.h
index 5e7fbfe29e..57e116aa05 100644
--- a/scene/3d/soft_dynamic_body_3d.h
+++ b/scene/3d/soft_dynamic_body_3d.h
@@ -90,7 +90,7 @@ private:
 
 	DisableMode disable_mode = DISABLE_MODE_REMOVE;
 
-	bool mesh_owner = false;
+	RID owned_mesh;
 	uint32_t collision_mask = 1;
 	uint32_t collision_layer = 1;
 	NodePath parent_collision_ignore;
@@ -106,6 +106,12 @@ private:
 
 	void _update_pickable();
 
+	void _update_physics_server();
+	void _draw_soft_mesh();
+
+	void _prepare_physics_server();
+	void _become_mesh_owner();
+
 protected:
 	bool _set(const StringName &p_name, const Variant &p_value);
 	bool _get(const StringName &p_name, Variant &r_ret) const;
@@ -120,14 +126,7 @@ protected:
 
 	TypedArray<String> get_configuration_warnings() const override;
 
-protected:
-	void _update_physics_server();
-	void _draw_soft_mesh();
-
 public:
-	void prepare_physics_server();
-	void become_mesh_owner();
-
 	RID get_physics_rid() const { return physics_rid; }
 
 	void set_collision_mask(uint32_t p_mask);