2 files changed, 30 insertions, 7 deletions

diff --git a/scene/resources/particles_material.cpp b/scene/resources/particles_material.cpp
index 691d73c181..edca9e5171 100644
--- a/scene/resources/particles_material.cpp
+++ b/scene/resources/particles_material.cpp
@@ -45,6 +45,7 @@ void ParticlesMaterial::init_shaders() {
 
 	shader_names = memnew(ShaderNames);
 
+	shader_names->direction = "direction";
 	shader_names->spread = "spread";
 	shader_names->flatness = "flatness";
 	shader_names->initial_linear_velocity = "initial_linear_velocity";
@@ -144,6 +145,7 @@ void ParticlesMaterial::_update_shader() {
 
 	String code = "shader_type particles;\n";
 
+	code += "uniform vec3 direction;\n";
 	code += "uniform float spread;\n";
 	code += "uniform float flatness;\n";
 	code += "uniform float initial_linear_velocity;\n";
@@ -304,20 +306,20 @@ void ParticlesMaterial::_update_shader() {
 
 	if (flags[FLAG_DISABLE_Z]) {
 
-		code += "		float angle1_rad = rand_from_seed_m1_p1(alt_seed) * spread_rad;\n";
+		code += "		float angle1_rad = atan(direction.y, direction.x) + rand_from_seed_m1_p1(alt_seed) * spread_rad;\n";
 		code += "		vec3 rot = vec3(cos(angle1_rad), sin(angle1_rad), 0.0);\n";
 		code += "		VELOCITY = rot * initial_linear_velocity * mix(1.0, rand_from_seed(alt_seed), initial_linear_velocity_random);\n";
 
 	} else {
 		//initiate velocity spread in 3D
-		code += "		float angle1_rad = rand_from_seed_m1_p1(alt_seed) * spread_rad;\n";
-		code += "		float angle2_rad = rand_from_seed_m1_p1(alt_seed) * spread_rad * (1.0 - flatness);\n";
+		code += "		float angle1_rad = atan(direction.x, direction.z) + rand_from_seed_m1_p1(alt_seed) * spread_rad;\n";
+		code += "		float angle2_rad = atan(direction.y, abs(direction.z)) + rand_from_seed_m1_p1(alt_seed) * spread_rad * (1.0 - flatness);\n";
 		code += "		vec3 direction_xz = vec3(sin(angle1_rad), 0.0, cos(angle1_rad));\n";
 		code += "		vec3 direction_yz = vec3(0.0, sin(angle2_rad), cos(angle2_rad));\n";
 		code += "		direction_yz.z = direction_yz.z / max(0.0001,sqrt(abs(direction_yz.z))); // better uniform distribution\n";
-		code += "		vec3 direction = vec3(direction_xz.x * direction_yz.z, direction_yz.y, direction_xz.z * direction_yz.z);\n";
-		code += "		direction = normalize(direction);\n";
-		code += "		VELOCITY = direction * initial_linear_velocity * mix(1.0, rand_from_seed(alt_seed), initial_linear_velocity_random);\n";
+		code += "		vec3 vec_direction = vec3(direction_xz.x * direction_yz.z, direction_yz.y, direction_xz.z * direction_yz.z);\n";
+		code += "		vec_direction = normalize(vec_direction);\n";
+		code += "		VELOCITY = vec_direction * initial_linear_velocity * mix(1.0, rand_from_seed(alt_seed), initial_linear_velocity_random);\n";
 	}
 
 	code += "		float base_angle = (initial_angle + tex_angle) * mix(1.0, angle_rand, initial_angle_random);\n";
@@ -627,6 +629,17 @@ bool ParticlesMaterial::_is_shader_dirty() const {
 	return dirty;
 }
 
+void ParticlesMaterial::set_direction(Vector3 p_direction) {
+
+	direction = p_direction;
+	VisualServer::get_singleton()->material_set_param(_get_material(), shader_names->direction, direction);
+}
+
+Vector3 ParticlesMaterial::get_direction() const {
+
+	return direction;
+}
+
 void ParticlesMaterial::set_spread(float p_spread) {
 
 	spread = p_spread;
@@ -1042,6 +1055,9 @@ Shader::Mode ParticlesMaterial::get_shader_mode() const {
 
 void ParticlesMaterial::_bind_methods() {
 
+	ClassDB::bind_method(D_METHOD("set_direction", "degrees"), &ParticlesMaterial::set_direction);
+	ClassDB::bind_method(D_METHOD("get_direction"), &ParticlesMaterial::get_direction);
+
 	ClassDB::bind_method(D_METHOD("set_spread", "degrees"), &ParticlesMaterial::set_spread);
 	ClassDB::bind_method(D_METHOD("get_spread"), &ParticlesMaterial::get_spread);
 
@@ -1115,7 +1131,8 @@ void ParticlesMaterial::_bind_methods() {
 	ADD_PROPERTYI(PropertyInfo(Variant::BOOL, "flag_align_y"), "set_flag", "get_flag", FLAG_ALIGN_Y_TO_VELOCITY);
 	ADD_PROPERTYI(PropertyInfo(Variant::BOOL, "flag_rotate_y"), "set_flag", "get_flag", FLAG_ROTATE_Y);
 	ADD_PROPERTYI(PropertyInfo(Variant::BOOL, "flag_disable_z"), "set_flag", "get_flag", FLAG_DISABLE_Z);
-	ADD_GROUP("Spread", "");
+	ADD_GROUP("Direction", "");
+	ADD_PROPERTY(PropertyInfo(Variant::VECTOR3, "direction"), "set_direction", "get_direction");
 	ADD_PROPERTY(PropertyInfo(Variant::REAL, "spread", PROPERTY_HINT_RANGE, "0,180,0.01"), "set_spread", "get_spread");
 	ADD_PROPERTY(PropertyInfo(Variant::REAL, "flatness", PROPERTY_HINT_RANGE, "0,1,0.01"), "set_flatness", "get_flatness");
 	ADD_GROUP("Gravity", "");
@@ -1200,6 +1217,7 @@ void ParticlesMaterial::_bind_methods() {
 ParticlesMaterial::ParticlesMaterial() :
 		element(this) {
 
+	set_direction(Vector3(1, 0, 0));
 	set_spread(45);
 	set_flatness(0);
 	set_param(PARAM_INITIAL_LINEAR_VELOCITY, 0);
diff --git a/scene/resources/particles_material.h b/scene/resources/particles_material.h
index 42bf60ff48..7f11bd794b 100644
--- a/scene/resources/particles_material.h
+++ b/scene/resources/particles_material.h
@@ -129,6 +129,7 @@ private:
 	static SelfList<ParticlesMaterial>::List *dirty_materials;
 
 	struct ShaderNames {
+		StringName direction;
 		StringName spread;
 		StringName flatness;
 		StringName initial_linear_velocity;
@@ -194,6 +195,7 @@ private:
 	_FORCE_INLINE_ void _queue_shader_change();
 	_FORCE_INLINE_ bool _is_shader_dirty() const;
 
+	Vector3 direction;
 	float spread;
 	float flatness;
 
@@ -230,6 +232,9 @@ protected:
 	virtual void _validate_property(PropertyInfo &property) const;
 
 public:
+	void set_direction(Vector3 p_direction);
+	Vector3 get_direction() const;
+
 	void set_spread(float p_spread);
 	float get_spread() const;