From 3bb189d75b0118f0176a35f420d9b594acf08648 Mon Sep 17 00:00:00 2001
From: Jonathan Cornaz <jonathan.cornaz@gmail.com>
Date: Mon, 8 Mar 2021 22:59:17 +0100
Subject: [PATCH] Introduce rotation contraints

---
 core/src/lib.rs    | 38 ++++++++++++++++++++++++++++++++++++++
 rapier/src/body.rs | 23 +++++++++++++++++++++--
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/core/src/lib.rs b/core/src/lib.rs
index 7b5c6f36..6c00b572 100644
--- a/core/src/lib.rs
+++ b/core/src/lib.rs
@@ -79,6 +79,7 @@ impl Plugin for CorePlugin {
             .register_type::<BodyType>()
             .register_type::<PhysicMaterial>()
             .register_type::<Velocity>()
+            .register_type::<RotationConstraints>()
             .add_stage_after(bevy::app::stage::UPDATE, crate::stage::ROOT, {
                 let mut schedule = Schedule::default();
 
@@ -276,3 +277,40 @@ impl Default for PhysicMaterial {
         }
     }
 }
+
+/// Components that restrict rotations due to forces
+///
+/// Note that angular velocity may still be applied programmatically.
+#[derive(Debug, Copy, Clone, Reflect)]
+pub struct RotationConstraints {
+    /// Set to true to prevent rotations around the x axis
+    pub allow_x: bool,
+
+    /// Set to true to prevent rotations around the y axis
+    pub allow_y: bool,
+
+    /// Set to true to prevent rotations around the Z axis
+    pub allow_z: bool,
+}
+
+impl Default for RotationConstraints {
+    fn default() -> Self {
+        Self {
+            allow_x: true,
+            allow_y: true,
+            allow_z: true,
+        }
+    }
+}
+
+impl RotationConstraints {
+    /// Lock rotations around all axes
+    #[must_use]
+    pub fn lock() -> Self {
+        Self {
+            allow_x: false,
+            allow_y: false,
+            allow_z: false,
+        }
+    }
+}
diff --git a/rapier/src/body.rs b/rapier/src/body.rs
index 3219b271..424aa9a0 100644
--- a/rapier/src/body.rs
+++ b/rapier/src/body.rs
@@ -3,7 +3,7 @@ use bevy::math::prelude::*;
 use bevy::transform::prelude::*;
 use fnv::FnvHashMap;
 
-use heron_core::{Body, BodyType, PhysicMaterial, Velocity};
+use heron_core::{Body, BodyType, PhysicMaterial, RotationConstraints, Velocity};
 
 use crate::convert::{IntoBevy, IntoRapier};
 use crate::rapier::dynamics::{
@@ -29,17 +29,36 @@ pub(crate) fn create(
             Option<&BodyType>,
             Option<&Velocity>,
             Option<&PhysicMaterial>,
+            Option<&RotationConstraints>,
         ),
         Without<BodyHandle>,
     >,
 ) {
-    for (entity, body, transform, body_type, velocity, material) in query.iter() {
+    for (entity, body, transform, body_type, velocity, material, restrict_rotation) in query.iter()
+    {
         let body_type = body_type.cloned().unwrap_or_default();
 
         let mut builder = RigidBodyBuilder::new(body_status(body_type))
             .user_data(entity.to_bits().into())
             .position((transform.translation, transform.rotation).into_rapier());
 
+        #[allow(unused_variables)]
+        if let Some(RotationConstraints {
+            allow_x,
+            allow_y,
+            allow_z,
+        }) = restrict_rotation.cloned()
+        {
+            #[cfg(feature = "2d")]
+            if !allow_z {
+                builder = builder.lock_rotations();
+            }
+            #[cfg(feature = "3d")]
+            {
+                builder = builder.restrict_rotations(allow_x, allow_y, allow_z);
+            }
+        }
+
         if let Some(v) = velocity {
             #[cfg(feature = "2d")]
             {