From f4fb1b9733f71fad1d66022896aafc4febbbe4b1 Mon Sep 17 00:00:00 2001
From: juliohq <JulioHenrique31501234@gmail.com>
Date: Sun, 28 Apr 2024 09:09:40 -0300
Subject: [PATCH] Add ability to lock on rotate/grab pan

---
 examples/advanced.rs |  4 ++++
 src/controller.rs    | 39 ++++++++++++++++++++++++++++++++++++---
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/examples/advanced.rs b/examples/advanced.rs
index fea8d69..59593ce 100644
--- a/examples/advanced.rs
+++ b/examples/advanced.rs
@@ -146,8 +146,12 @@ Press T to toggle controls (K and L will still work)"
             key_right: KeyCode::KeyD,
             // Rotate the camera with right click
             button_rotate: MouseButton::Right,
+            // Keep the mouse cursor in place when rotating
+            lock_on_rotate: true,
             // Drag pan with middle click
             button_drag: Some(MouseButton::Middle),
+            // Keep the mouse cursor in place when dragging
+            lock_on_drag: true,
             // Change the width of the area that triggers edge pan. 0.1 is 10% of the window height.
             edge_pan_width: 0.1,
             // Increase pan speed
diff --git a/src/controller.rs b/src/controller.rs
index e092d3d..ab09cac 100644
--- a/src/controller.rs
+++ b/src/controller.rs
@@ -1,8 +1,10 @@
+#![allow(clippy::too_many_arguments)]
+
 use crate::{Ground, RtsCamera, RtsCameraSystemSet};
 use bevy::input::mouse::{MouseMotion, MouseScrollUnit, MouseWheel};
 use bevy::input::ButtonInput;
 use bevy::prelude::*;
-use bevy::window::PrimaryWindow;
+use bevy::window::{CursorGrabMode, PrimaryWindow};
 use bevy_mod_raycast::immediate::{Raycast, RaycastSettings};
 use bevy_mod_raycast::{CursorRay, DefaultRaycastingPlugin};
 use std::f32::consts::PI;
@@ -66,9 +68,15 @@ pub struct RtsCameraControls {
     /// How fast the keys will rotate the camera.
     /// Defaults to `16.0`.
     pub key_rotate_speed: f32,
+    /// Whether to lock the mouse cursor in place while rotating.
+    /// Defaults to `false`.
+    pub lock_on_rotate: bool,
     /// The mouse button used to 'drag pan' the camera.
     /// Defaults to `None`.
     pub button_drag: Option<MouseButton>,
+    /// Whether to lock the mouse cursor in place while dragging.
+    /// Defaults to `false`.
+    pub lock_on_drag: bool,
     /// How far away from the side of the screen edge pan will kick in, defined as a percentage
     /// of the window's height. Set to `0.0` to disable edge panning.
     /// Defaults to `0.05` (5%).
@@ -92,7 +100,9 @@ impl Default for RtsCameraControls {
             key_rotate_left: KeyCode::KeyQ,
             key_rotate_right: KeyCode::KeyE,
             key_rotate_speed: 16.0,
+            lock_on_rotate: false,
             button_drag: None,
+            lock_on_drag: false,
             edge_pan_width: 0.05,
             pan_speed: 15.0,
             enabled: true,
@@ -199,6 +209,7 @@ pub fn grab_pan(
     cursor_ray: Res<CursorRay>,
     mut ray_hit: Local<Option<Vec3>>,
     ground_q: Query<Entity, With<Ground>>,
+    mut primary_window_q: Query<&mut Window, With<PrimaryWindow>>,
 ) {
     for (cam_tfm, mut cam, controller, camera, projection) in
         cam_q.iter_mut().filter(|(_, _, ctrl, _, _)| ctrl.enabled)
@@ -208,6 +219,13 @@ pub fn grab_pan(
         };
 
         if mouse_button.just_pressed(drag_button) {
+            if controller.lock_on_drag {
+                if let Ok(mut primary_window) = primary_window_q.get_single_mut() {
+                    primary_window.cursor.grab_mode = CursorGrabMode::Locked;
+                    primary_window.cursor.visible = false;
+                }
+            }
+
             if let Some(cursor_ray) = **cursor_ray {
                 *ray_hit = raycast
                     .cast_ray(
@@ -224,6 +242,11 @@ pub fn grab_pan(
 
         if mouse_button.just_released(drag_button) {
             *ray_hit = None;
+
+            if let Ok(mut primary_window) = primary_window_q.get_single_mut() {
+                primary_window.cursor.grab_mode = CursorGrabMode::None;
+                primary_window.cursor.visible = true;
+            }
         }
 
         if mouse_button.pressed(drag_button) {
@@ -257,10 +280,15 @@ pub fn rotate(
     mouse_input: Res<ButtonInput<MouseButton>>,
     keys: Res<ButtonInput<KeyCode>>,
     mut mouse_motion: EventReader<MouseMotion>,
-    primary_window_q: Query<&Window, With<PrimaryWindow>>,
+    mut primary_window_q: Query<&mut Window, With<PrimaryWindow>>,
 ) {
-    if let Ok(primary_window) = primary_window_q.get_single() {
+    if let Ok(mut primary_window) = primary_window_q.get_single_mut() {
         for (mut cam, controller) in cam_q.iter_mut().filter(|(_, ctrl)| ctrl.enabled) {
+            if mouse_input.just_pressed(controller.button_rotate) {
+                primary_window.cursor.grab_mode = CursorGrabMode::Locked;
+                primary_window.cursor.visible = false;
+            }
+
             if mouse_input.pressed(controller.button_rotate) {
                 let mouse_delta = mouse_motion.read().map(|e| e.delta).sum::<Vec2>();
                 // Adjust based on window size, so that moving mouse entire width of window
@@ -283,6 +311,11 @@ pub fn rotate(
                     (right - left) / primary_window.width() * PI * controller.key_rotate_speed,
                 );
             }
+
+            if mouse_input.just_released(controller.button_rotate) {
+                primary_window.cursor.grab_mode = CursorGrabMode::None;
+                primary_window.cursor.visible = true;
+            }
         }
     }
 }