Add Particle Shader Userdata

* Adds optional vec4 USERDATA1 .. USERDATA6 to particles, allowing to store custom data.
* This data is allocated on demand, so shaders that do not use it do not cost more.

4 files changed, 121 insertions, 35 deletions

diff --git a/servers/rendering/renderer_rd/renderer_storage_rd.cpp b/servers/rendering/renderer_rd/renderer_storage_rd.cpp
index b42591f9a3..e3829eb5ed 100644
--- a/servers/rendering/renderer_rd/renderer_storage_rd.cpp
+++ b/servers/rendering/renderer_rd/renderer_storage_rd.cpp
@@ -4522,6 +4522,8 @@ void RendererStorageRD::_particles_free_data(Particles *particles) {
 		particles->particle_instance_buffer = RID();
 	}
 
+	particles->userdata_count = 0;
+
 	if (particles->frame_params_buffer.is_valid()) {
 		RD::get_singleton()->free(particles->frame_params_buffer);
 		particles->frame_params_buffer = RID();
@@ -4716,6 +4718,14 @@ void RendererStorageRD::particles_set_process_material(RID p_particles, RID p_ma
 	ERR_FAIL_COND(!particles);
 
 	particles->process_material = p_material;
+	particles->dependency.changed_notify(DEPENDENCY_CHANGED_PARTICLES); //the instance buffer may have changed
+}
+
+RID RendererStorageRD::particles_get_process_material(RID p_particles) const {
+	Particles *particles = particles_owner.get_or_null(p_particles);
+	ERR_FAIL_COND_V(!particles, RID());
+
+	return particles->process_material;
 }
 
 void RendererStorageRD::particles_set_draw_order(RID p_particles, RS::ParticlesDrawOrder p_order) {
@@ -4852,10 +4862,13 @@ AABB RendererStorageRD::particles_get_current_aabb(RID p_particles) {
 	if (buffer.size()) {
 		bool first = true;
 
-		const ParticleData *particle_data = reinterpret_cast<const ParticleData *>(buffer.ptr());
+		const uint8_t *data_ptr = (const uint8_t *)buffer.ptr();
+		uint32_t particle_data_size = sizeof(ParticleData) + sizeof(float) * particles->userdata_count;
+
 		for (int i = 0; i < total_amount; i++) {
-			if (particle_data[i].active) {
-				Vector3 pos = Vector3(particle_data[i].xform[12], particle_data[i].xform[13], particle_data[i].xform[14]);
+			const ParticleData &particle_data = *(const ParticleData *)&data_ptr[particle_data_size * i];
+			if (particle_data.active) {
+				Vector3 pos = Vector3(particle_data.xform[12], particle_data.xform[13], particle_data.xform[14]);
 				if (!particles->use_local_coords) {
 					pos = inv.xform(pos);
 				}
@@ -5420,6 +5433,7 @@ void RendererStorageRD::particles_set_view_axis(RID p_particles, const Vector3 &
 
 	copy_push_constant.frame_remainder = particles->interpolate ? particles->frame_remainder : 0.0;
 	copy_push_constant.total_particles = particles->amount;
+	copy_push_constant.copy_mode_2d = false;
 
 	Vector3 axis = -p_axis; // cameras look to z negative
 
@@ -5440,7 +5454,7 @@ void RendererStorageRD::particles_set_view_axis(RID p_particles, const Vector3 &
 	if (do_sort) {
 		RD::ComputeListID compute_list = RD::get_singleton()->compute_list_begin();
 
-		RD::get_singleton()->compute_list_bind_compute_pipeline(compute_list, particles_shader.copy_pipelines[ParticlesShader::COPY_MODE_FILL_SORT_BUFFER]);
+		RD::get_singleton()->compute_list_bind_compute_pipeline(compute_list, particles_shader.copy_pipelines[ParticlesShader::COPY_MODE_FILL_SORT_BUFFER + particles->userdata_count * ParticlesShader::COPY_MODE_MAX]);
 		RD::get_singleton()->compute_list_bind_uniform_set(compute_list, particles->particles_copy_uniform_set, 0);
 		RD::get_singleton()->compute_list_bind_uniform_set(compute_list, particles->particles_sort_uniform_set, 1);
 		RD::get_singleton()->compute_list_bind_uniform_set(compute_list, particles->trail_bind_pose_uniform_set, 2);
@@ -5455,7 +5469,10 @@ void RendererStorageRD::particles_set_view_axis(RID p_particles, const Vector3 &
 	copy_push_constant.total_particles *= copy_push_constant.total_particles;
 
 	RD::ComputeListID compute_list = RD::get_singleton()->compute_list_begin();
-	RD::get_singleton()->compute_list_bind_compute_pipeline(compute_list, particles_shader.copy_pipelines[do_sort ? ParticlesShader::COPY_MODE_FILL_INSTANCES_WITH_SORT_BUFFER : (particles->mode == RS::PARTICLES_MODE_2D ? ParticlesShader::COPY_MODE_FILL_INSTANCES_2D : ParticlesShader::COPY_MODE_FILL_INSTANCES)]);
+	uint32_t copy_pipeline = do_sort ? ParticlesShader::COPY_MODE_FILL_INSTANCES_WITH_SORT_BUFFER : ParticlesShader::COPY_MODE_FILL_INSTANCES;
+	copy_pipeline += particles->userdata_count * ParticlesShader::COPY_MODE_MAX;
+	copy_push_constant.copy_mode_2d = particles->mode == RS::PARTICLES_MODE_2D ? 1 : 0;
+	RD::get_singleton()->compute_list_bind_compute_pipeline(compute_list, particles_shader.copy_pipelines[copy_pipeline]);
 	RD::get_singleton()->compute_list_bind_uniform_set(compute_list, particles->particles_copy_uniform_set, 0);
 	if (do_sort) {
 		RD::get_singleton()->compute_list_bind_uniform_set(compute_list, particles->particles_sort_uniform_set, 1);
@@ -5470,6 +5487,19 @@ void RendererStorageRD::particles_set_view_axis(RID p_particles, const Vector3 &
 }
 
 void RendererStorageRD::_particles_update_buffers(Particles *particles) {
+	uint32_t userdata_count = 0;
+
+	const Material *material = material_owner.get_or_null(particles->process_material);
+	if (material && material->shader && material->shader->data) {
+		const ParticlesShaderData *shader_data = static_cast<const ParticlesShaderData *>(material->shader->data);
+		userdata_count = shader_data->userdata_count;
+	}
+
+	if (userdata_count != particles->userdata_count) {
+		// Mismatch userdata, re-create buffers.
+		_particles_free_data(particles);
+	}
+
 	if (particles->amount > 0 && particles->particle_buffer.is_null()) {
 		int total_amount = particles->amount;
 		if (particles->trails_enabled && particles->trail_bind_poses.size() > 1) {
@@ -5478,7 +5508,9 @@ void RendererStorageRD::_particles_update_buffers(Particles *particles) {
 
 		uint32_t xform_size = particles->mode == RS::PARTICLES_MODE_2D ? 2 : 3;
 
-		particles->particle_buffer = RD::get_singleton()->storage_buffer_create(sizeof(ParticleData) * total_amount);
+		particles->particle_buffer = RD::get_singleton()->storage_buffer_create((sizeof(ParticleData) + userdata_count * sizeof(float) * 4) * total_amount);
+
+		particles->userdata_count = userdata_count;
 
 		particles->particle_instance_buffer = RD::get_singleton()->storage_buffer_create(sizeof(float) * 4 * (xform_size + 1 + 1) * total_amount);
 		//needs to clear it
@@ -5702,7 +5734,8 @@ void RendererStorageRD::update_particles() {
 			copy_push_constant.lifetime_reverse = particles->draw_order == RS::PARTICLES_DRAW_ORDER_REVERSE_LIFETIME;
 
 			RD::ComputeListID compute_list = RD::get_singleton()->compute_list_begin();
-			RD::get_singleton()->compute_list_bind_compute_pipeline(compute_list, particles_shader.copy_pipelines[particles->mode == RS::PARTICLES_MODE_2D ? ParticlesShader::COPY_MODE_FILL_INSTANCES_2D : ParticlesShader::COPY_MODE_FILL_INSTANCES]);
+			copy_push_constant.copy_mode_2d = particles->mode == RS::PARTICLES_MODE_2D ? 1 : 0;
+			RD::get_singleton()->compute_list_bind_compute_pipeline(compute_list, particles_shader.copy_pipelines[ParticlesShader::COPY_MODE_FILL_INSTANCES + particles->userdata_count * ParticlesShader::COPY_MODE_MAX]);
 			RD::get_singleton()->compute_list_bind_uniform_set(compute_list, particles->particles_copy_uniform_set, 0);
 			RD::get_singleton()->compute_list_bind_uniform_set(compute_list, particles->trail_bind_pose_uniform_set, 2);
 			RD::get_singleton()->compute_list_set_push_constant(compute_list, &copy_push_constant, sizeof(ParticlesShader::CopyPushConstant));
@@ -5754,6 +5787,12 @@ void RendererStorageRD::ParticlesShaderData::set_code(const String &p_code) {
 
 	actions.usage_flag_pointers["COLLIDED"] = &uses_collision;
 
+	userdata_count = 0;
+	for (uint32_t i = 0; i < ParticlesShader::MAX_USERDATAS; i++) {
+		userdatas_used[i] = false;
+		actions.usage_flag_pointers["USERDATA" + itos(i + 1)] = &userdatas_used[i];
+	}
+
 	actions.uniforms = &uniforms;
 
 	Error err = base_singleton->particles_shader.compiler.compile(RS::SHADER_PARTICLES, code, &actions, path, gen_code);
@@ -5763,6 +5802,12 @@ void RendererStorageRD::ParticlesShaderData::set_code(const String &p_code) {
 		version = base_singleton->particles_shader.shader.version_create();
 	}
 
+	for (uint32_t i = 0; i < ParticlesShader::MAX_USERDATAS; i++) {
+		if (userdatas_used[i]) {
+			userdata_count++;
+		}
+	}
+
 	base_singleton->particles_shader.shader.version_set_compute_code(version, gen_code.code, gen_code.uniforms, gen_code.stage_globals[ShaderCompiler::STAGE_COMPUTE], gen_code.defines);
 	ERR_FAIL_COND(!base_singleton->particles_shader.shader.version_is_valid(version));
 
@@ -9930,6 +9975,11 @@ RendererStorageRD::RendererStorageRD() {
 		actions.renames["ACTIVE"] = "particle_active";
 		actions.renames["RESTART"] = "restart";
 		actions.renames["CUSTOM"] = "PARTICLE.custom";
+		for (int i = 0; i < ParticlesShader::MAX_USERDATAS; i++) {
+			String udname = "USERDATA" + itos(i + 1);
+			actions.renames[udname] = "PARTICLE.userdata" + itos(i + 1);
+			actions.usage_defines[udname] = "#define USERDATA" + itos(i + 1) + "_USED\n";
+		}
 		actions.renames["TRANSFORM"] = "PARTICLE.xform";
 		actions.renames["TIME"] = "frame_history.data[0].time";
 		actions.renames["PI"] = _MKSTR(Math_PI);
@@ -10034,17 +10084,26 @@ void process() {
 
 	{
 		Vector<String> copy_modes;
-		copy_modes.push_back("\n#define MODE_FILL_INSTANCES\n");
-		copy_modes.push_back("\n#define MODE_FILL_INSTANCES\n#define MODE_2D\n");
-		copy_modes.push_back("\n#define MODE_FILL_SORT_BUFFER\n#define USE_SORT_BUFFER\n");
-		copy_modes.push_back("\n#define MODE_FILL_INSTANCES\n#define USE_SORT_BUFFER\n");
+		for (int i = 0; i <= ParticlesShader::MAX_USERDATAS; i++) {
+			if (i == 0) {
+				copy_modes.push_back("\n#define MODE_FILL_INSTANCES\n");
+				copy_modes.push_back("\n#define MODE_FILL_SORT_BUFFER\n#define USE_SORT_BUFFER\n");
+				copy_modes.push_back("\n#define MODE_FILL_INSTANCES\n#define USE_SORT_BUFFER\n");
+			} else {
+				copy_modes.push_back("\n#define MODE_FILL_INSTANCES\n#define USERDATA_COUNT " + itos(i) + "\n");
+				copy_modes.push_back("\n#define MODE_FILL_SORT_BUFFER\n#define USE_SORT_BUFFER\n#define USERDATA_COUNT " + itos(i) + "\n");
+				copy_modes.push_back("\n#define MODE_FILL_INSTANCES\n#define USE_SORT_BUFFER\n#define USERDATA_COUNT " + itos(i) + "\n");
+			}
+		}
 
 		particles_shader.copy_shader.initialize(copy_modes);
 
 		particles_shader.copy_shader_version = particles_shader.copy_shader.version_create();
 
-		for (int i = 0; i < ParticlesShader::COPY_MODE_MAX; i++) {
-			particles_shader.copy_pipelines[i] = RD::get_singleton()->compute_pipeline_create(particles_shader.copy_shader.version_get_shader(particles_shader.copy_shader_version, i));
+		for (int i = 0; i <= ParticlesShader::MAX_USERDATAS; i++) {
+			for (int j = 0; j < ParticlesShader::COPY_MODE_MAX; j++) {
+				particles_shader.copy_pipelines[i * ParticlesShader::COPY_MODE_MAX + j] = RD::get_singleton()->compute_pipeline_create(particles_shader.copy_shader.version_get_shader(particles_shader.copy_shader_version, i * ParticlesShader::COPY_MODE_MAX + j));
+			}
 		}
 	}
 
diff --git a/servers/rendering/renderer_rd/renderer_storage_rd.h b/servers/rendering/renderer_rd/renderer_storage_rd.h
index 6febb71e93..ee4d18210a 100644
--- a/servers/rendering/renderer_rd/renderer_storage_rd.h
+++ b/servers/rendering/renderer_rd/renderer_storage_rd.h
@@ -751,6 +751,8 @@ private:
 		RID particle_instance_buffer;
 		RID frame_params_buffer;
 
+		uint32_t userdata_count = 0;
+
 		RID particles_material_uniform_set;
 		RID particles_copy_uniform_set;
 		RID particles_transforms_buffer_uniform_set;
@@ -849,12 +851,14 @@ private:
 			uint32_t order_by_lifetime;
 			uint32_t lifetime_split;
 			uint32_t lifetime_reverse;
-			uint32_t pad;
+			uint32_t copy_mode_2d;
 		};
 
 		enum {
+			MAX_USERDATAS = 6
+		};
+		enum {
 			COPY_MODE_FILL_INSTANCES,
-			COPY_MODE_FILL_INSTANCES_2D,
 			COPY_MODE_FILL_SORT_BUFFER,
 			COPY_MODE_FILL_INSTANCES_WITH_SORT_BUFFER,
 			COPY_MODE_MAX,
@@ -862,7 +866,7 @@ private:
 
 		ParticlesCopyShaderRD copy_shader;
 		RID copy_shader_version;
-		RID copy_pipelines[COPY_MODE_MAX];
+		RID copy_pipelines[COPY_MODE_MAX * (MAX_USERDATAS + 1)];
 
 		LocalVector<float> pose_update_buffer;
 
@@ -888,7 +892,10 @@ private:
 
 		RID pipeline;
 
-		bool uses_time;
+		bool uses_time = false;
+
+		bool userdatas_used[ParticlesShader::MAX_USERDATAS] = {};
+		uint32_t userdata_count = 0;
 
 		virtual void set_code(const String &p_Code);
 		virtual void set_default_texture_param(const StringName &p_name, RID p_texture, int p_index);
@@ -2162,6 +2169,8 @@ public:
 	void particles_set_speed_scale(RID p_particles, double p_scale);
 	void particles_set_use_local_coordinates(RID p_particles, bool p_enable);
 	void particles_set_process_material(RID p_particles, RID p_material);
+	RID particles_get_process_material(RID p_particles) const;
+
 	void particles_set_fixed_fps(RID p_particles, int p_fps);
 	void particles_set_interpolate(RID p_particles, bool p_enable);
 	void particles_set_fractional_delta(RID p_particles, bool p_enable);
diff --git a/servers/rendering/renderer_rd/shaders/particles.glsl b/servers/rendering/renderer_rd/shaders/particles.glsl
index d691ea2fdf..1b1051ecfa 100644
--- a/servers/rendering/renderer_rd/shaders/particles.glsl
+++ b/servers/rendering/renderer_rd/shaders/particles.glsl
@@ -112,6 +112,24 @@ struct ParticleData {
 	uint flags;
 	vec4 color;
 	vec4 custom;
+#ifdef USERDATA1_USED
+	vec4 userdata1;
+#endif
+#ifdef USERDATA2_USED
+	vec4 userdata2;
+#endif
+#ifdef USERDATA3_USED
+	vec4 userdata3;
+#endif
+#ifdef USERDATA4_USED
+	vec4 userdata4;
+#endif
+#ifdef USERDATA5_USED
+	vec4 userdata5;
+#endif
+#ifdef USERDATA6_USED
+	vec4 userdata6;
+#endif
 };
 
 layout(set = 1, binding = 1, std430) restrict buffer Particles {
diff --git a/servers/rendering/renderer_rd/shaders/particles_copy.glsl b/servers/rendering/renderer_rd/shaders/particles_copy.glsl
index bb11e4c78d..b991880cd9 100644
--- a/servers/rendering/renderer_rd/shaders/particles_copy.glsl
+++ b/servers/rendering/renderer_rd/shaders/particles_copy.glsl
@@ -16,6 +16,9 @@ struct ParticleData {
 	uint flags;
 	vec4 color;
 	vec4 custom;
+#ifdef USERDATA_COUNT
+	vec4 userdata[USERDATA_COUNT];
+#endif
 };
 
 layout(set = 0, binding = 1, std430) restrict readonly buffer Particles {
@@ -57,7 +60,7 @@ layout(push_constant, std430) uniform Params {
 	bool order_by_lifetime;
 	uint lifetime_split;
 	bool lifetime_reverse;
-	uint pad;
+	bool copy_mode_2d;
 }
 params;
 
@@ -201,25 +204,22 @@ void main() {
 		txform = mat4(vec4(0.0), vec4(0.0), vec4(0.0), vec4(0.0)); //zero scale, becomes invisible
 	}
 
-#ifdef MODE_2D
-
-	uint write_offset = gl_GlobalInvocationID.x * (2 + 1 + 1); //xform + color + custom
+	if (params.copy_mode_2d) {
+		uint write_offset = gl_GlobalInvocationID.x * (2 + 1 + 1); //xform + color + custom
 
-	instances.data[write_offset + 0] = txform[0];
-	instances.data[write_offset + 1] = txform[1];
-	instances.data[write_offset + 2] = particles.data[particle].color;
-	instances.data[write_offset + 3] = particles.data[particle].custom;
-
-#else
-
-	uint write_offset = gl_GlobalInvocationID.x * (3 + 1 + 1); //xform + color + custom
+		instances.data[write_offset + 0] = txform[0];
+		instances.data[write_offset + 1] = txform[1];
+		instances.data[write_offset + 2] = particles.data[particle].color;
+		instances.data[write_offset + 3] = particles.data[particle].custom;
+	} else {
+		uint write_offset = gl_GlobalInvocationID.x * (3 + 1 + 1); //xform + color + custom
 
-	instances.data[write_offset + 0] = txform[0];
-	instances.data[write_offset + 1] = txform[1];
-	instances.data[write_offset + 2] = txform[2];
-	instances.data[write_offset + 3] = particles.data[particle].color;
-	instances.data[write_offset + 4] = particles.data[particle].custom;
-#endif //MODE_2D
+		instances.data[write_offset + 0] = txform[0];
+		instances.data[write_offset + 1] = txform[1];
+		instances.data[write_offset + 2] = txform[2];
+		instances.data[write_offset + 3] = particles.data[particle].color;
+		instances.data[write_offset + 4] = particles.data[particle].custom;
+	}
 
 #endif
 }