Allow double-axis lock in RigidBody and KinematicBody

4 files changed, 22 insertions, 39 deletions

diff --git a/modules/bullet/bullet_physics_server.cpp b/modules/bullet/bullet_physics_server.cpp
index 7f95d16ba6..63a4fe5bc6 100644
--- a/modules/bullet/bullet_physics_server.cpp
+++ b/modules/bullet/bullet_physics_server.cpp
@@ -723,15 +723,15 @@ void BulletPhysicsServer::body_set_axis_velocity(RID p_body, const Vector3 &p_ax
 	body->set_linear_velocity(v);
 }
 
-void BulletPhysicsServer::body_set_axis_lock(RID p_body, PhysicsServer::BodyAxisLock p_lock) {
+void BulletPhysicsServer::body_set_axis_lock(RID p_body, int axis, bool p_lock) {
 	RigidBodyBullet *body = rigid_body_owner.get(p_body);
 	ERR_FAIL_COND(!body);
-	body->set_axis_lock(p_lock);
+	body->set_axis_lock(axis, p_lock);
 }
 
-PhysicsServer::BodyAxisLock BulletPhysicsServer::body_get_axis_lock(RID p_body) const {
+bool BulletPhysicsServer::body_get_axis_lock(RID p_body) const {
 	const RigidBodyBullet *body = rigid_body_owner.get(p_body);
-	ERR_FAIL_COND_V(!body, BODY_AXIS_LOCK_DISABLED);
+	ERR_FAIL_COND_V(!body, 0);
 	return body->get_axis_lock();
 }
 
diff --git a/modules/bullet/bullet_physics_server.h b/modules/bullet/bullet_physics_server.h
index ad8137ee2f..ed5acb9041 100644
--- a/modules/bullet/bullet_physics_server.h
+++ b/modules/bullet/bullet_physics_server.h
@@ -226,8 +226,8 @@ public:
 	virtual void body_apply_torque_impulse(RID p_body, const Vector3 &p_impulse);
 	virtual void body_set_axis_velocity(RID p_body, const Vector3 &p_axis_velocity);
 
-	virtual void body_set_axis_lock(RID p_body, BodyAxisLock p_lock);
-	virtual BodyAxisLock body_get_axis_lock(RID p_body) const;
+	virtual void body_set_axis_lock(RID p_body, int axis, bool p_lock);
+	virtual bool body_get_axis_lock(RID p_body) const;
 
 	virtual void body_add_collision_exception(RID p_body, RID p_body_b);
 	virtual void body_remove_collision_exception(RID p_body, RID p_body_b);
diff --git a/modules/bullet/rigid_body_bullet.cpp b/modules/bullet/rigid_body_bullet.cpp
index 98ae82bc5f..2bac863843 100644
--- a/modules/bullet/rigid_body_bullet.cpp
+++ b/modules/bullet/rigid_body_bullet.cpp
@@ -277,7 +277,7 @@ RigidBodyBullet::RigidBodyBullet()
 	setupBulletCollisionObject(btBody);
 
 	set_mode(PhysicsServer::BODY_MODE_RIGID);
-	set_axis_lock(PhysicsServer::BODY_AXIS_LOCK_DISABLED);
+	set_axis_lock(0, locked_axis[0]);
 
 	areasWhereIam.resize(maxAreasWhereIam);
 	for (int i = areasWhereIam.size() - 1; 0 <= i; --i) {
@@ -498,24 +498,24 @@ void RigidBodyBullet::set_mode(PhysicsServer::BodyMode p_mode) {
 	switch (p_mode) {
 		case PhysicsServer::BODY_MODE_KINEMATIC:
 			mode = PhysicsServer::BODY_MODE_KINEMATIC;
-			set_axis_lock(axis_lock); // Reload axis lock
+			set_axis_lock(0, locked_axis[0]); // Reload axis lock
 			_internal_set_mass(0);
 			init_kinematic_utilities();
 			break;
 		case PhysicsServer::BODY_MODE_STATIC:
 			mode = PhysicsServer::BODY_MODE_STATIC;
-			set_axis_lock(axis_lock); // Reload axis lock
+			set_axis_lock(0, locked_axis[0]); // Reload axis lock
 			_internal_set_mass(0);
 			break;
 		case PhysicsServer::BODY_MODE_RIGID: {
 			mode = PhysicsServer::BODY_MODE_RIGID;
-			set_axis_lock(axis_lock); // Reload axis lock
+			set_axis_lock(0, locked_axis[0]); // Reload axis lock
 			_internal_set_mass(0 == mass ? 1 : mass);
 			break;
 		}
 		case PhysicsServer::BODY_MODE_CHARACTER: {
 			mode = PhysicsServer::BODY_MODE_CHARACTER;
-			set_axis_lock(axis_lock); // Reload axis lock
+			set_axis_lock(0, locked_axis[0]); // Reload axis lock
 			_internal_set_mass(0 == mass ? 1 : mass);
 			break;
 		}
@@ -653,22 +653,14 @@ Vector3 RigidBodyBullet::get_applied_torque() const {
 	return gTotTorq;
 }
 
-void RigidBodyBullet::set_axis_lock(PhysicsServer::BodyAxisLock p_lock) {
-	axis_lock = p_lock;
+void RigidBodyBullet::set_axis_lock(int axis, bool p_lock) {
+	locked_axis[axis] = p_lock;
 
-	if (PhysicsServer::BODY_AXIS_LOCK_DISABLED == axis_lock) {
-		btBody->setLinearFactor(btVector3(1., 1., 1.));
+	btBody->setLinearFactor(btVector3(locked_axis[0] ? 0 : 1., locked_axis[1] ? 0 : 1., locked_axis[2] ? 0 : 1.));
+	if (locked_axis[0] || locked_axis[1] || locked_axis[2])
+		btBody->setAngularFactor(btVector3(locked_axis[0] ? 1. : 0, locked_axis[1] ? 1. : 0, locked_axis[2] ? 1. : 0));
+	else
 		btBody->setAngularFactor(btVector3(1., 1., 1.));
-	} else if (PhysicsServer::BODY_AXIS_LOCK_X == axis_lock) {
-		btBody->setLinearFactor(btVector3(0., 1., 1.));
-		btBody->setAngularFactor(btVector3(1., 0., 0.));
-	} else if (PhysicsServer::BODY_AXIS_LOCK_Y == axis_lock) {
-		btBody->setLinearFactor(btVector3(1., 0., 1.));
-		btBody->setAngularFactor(btVector3(0., 1., 0.));
-	} else if (PhysicsServer::BODY_AXIS_LOCK_Z == axis_lock) {
-		btBody->setLinearFactor(btVector3(1., 1., 0.));
-		btBody->setAngularFactor(btVector3(0., 0., 1.));
-	}
 
 	if (PhysicsServer::BODY_MODE_CHARACTER == mode) {
 		/// When character lock angular
@@ -676,17 +668,8 @@ void RigidBodyBullet::set_axis_lock(PhysicsServer::BodyAxisLock p_lock) {
 	}
 }
 
-PhysicsServer::BodyAxisLock RigidBodyBullet::get_axis_lock() const {
-	btVector3 vec = btBody->getLinearFactor();
-	if (0. == vec.x()) {
-		return PhysicsServer::BODY_AXIS_LOCK_X;
-	} else if (0. == vec.y()) {
-		return PhysicsServer::BODY_AXIS_LOCK_Y;
-	} else if (0. == vec.z()) {
-		return PhysicsServer::BODY_AXIS_LOCK_Z;
-	} else {
-		return PhysicsServer::BODY_AXIS_LOCK_DISABLED;
-	}
+bool RigidBodyBullet::get_axis_lock() const {
+	return locked_axis;
 }
 
 void RigidBodyBullet::set_continuous_collision_detection(bool p_enable) {
diff --git a/modules/bullet/rigid_body_bullet.h b/modules/bullet/rigid_body_bullet.h
index ab3c3e58b2..634ecf48cf 100644
--- a/modules/bullet/rigid_body_bullet.h
+++ b/modules/bullet/rigid_body_bullet.h
@@ -184,7 +184,7 @@ private:
 	KinematicUtilities *kinematic_utilities;
 
 	PhysicsServer::BodyMode mode;
-	PhysicsServer::BodyAxisLock axis_lock;
+	bool locked_axis[3] = { false, false, false };
 	GodotMotionState *godotMotionState;
 	btRigidBody *btBody;
 	real_t mass;
@@ -269,8 +269,8 @@ public:
 	void set_applied_torque(const Vector3 &p_torque);
 	Vector3 get_applied_torque() const;
 
-	void set_axis_lock(PhysicsServer::BodyAxisLock p_lock);
-	PhysicsServer::BodyAxisLock get_axis_lock() const;
+	void set_axis_lock(int axis, bool p_lock);
+	bool get_axis_lock() const;
 
 	/// Doc:
 	/// http://www.bulletphysics.org/mediawiki-1.5.8/index.php?title=Anti_tunneling_by_Motion_Clamping