Support aligning lanes to x or y axis during creation

crates/rmf_site_editor/src/interaction/mod.rs

@@ -151,6 +151,7 @@ impl Plugin for InteractionPlugin {
             .init_resource::<GizmoState>()
             .init_resource::<CurrentEditDrawing>()
             .init_resource::<CurrentLevel>()
+            .init_resource::<CreationSettings>()
             .insert_resource(HighlightAnchors(false))
             .add_event::<MoveTo>()
             .add_event::<GizmoClicked>()
@@ -263,7 +264,11 @@ impl Plugin for InteractionPlugin {
             .add_systems(OnExit(InteractionState::Enable), hide_cursor)
             .add_systems(
                 PostUpdate,
-                (move_anchor.before(update_anchor_transforms), move_pose)
+                (
+                    move_anchor.before(update_anchor_transforms),
+                    move_pose,
+                    apply_creation_settings.after(select_anchor_cursor_transform),
+                )
                     .run_if(in_state(InteractionState::Enable)),
             );
         }

crates/rmf_site_editor/src/interaction/select_impl/create_edges.rs

@@ -367,21 +367,38 @@ pub fn on_select_for_create_edges(
 }
 
 pub fn on_keyboard_for_create_edges(
-    In((button, key)): In<(KeyCode, BufferKey<CreateEdges>)>,
+    In(((button, input_type), key)): In<((KeyCode, ButtonInputType), BufferKey<CreateEdges>)>,
     mut access: BufferAccessMut<CreateEdges>,
     mut edges: Query<&mut Edge<Entity>>,
     cursor: Res<Cursor>,
+    mut creation_settings: ResMut<CreationSettings>,
     mut commands: Commands,
+    transform: Query<&'static Transform>,
 ) -> SelectionNodeResult {
-    if !matches!(button, KeyCode::Escape) {
-        // The button was not the escape key, so there's nothing for us to do
+    if !matches!(button, KeyCode::Escape) && !matches!(button, KeyCode::ShiftLeft) {
+        // The button was not the escape or shift key, so there's nothing for us to do
         // here.
         return Ok(());
     }
 
     let mut access = access.get_mut(&key).or_broken_buffer()?;
     let state = access.newest_mut().or_broken_state()?;
-    backout(state, &mut edges, &cursor, &mut commands)
+
+    if matches!(button, KeyCode::Escape) {
+        return backout(state, &mut edges, &cursor, &mut commands);
+    }
+    if matches!(button, KeyCode::ShiftLeft) {
+        return align(
+            state,
+            &mut edges,
+            &cursor,
+            &mut creation_settings,
+            &input_type,
+            &transform,
+        );
+    }
+
+    Ok(())
 }
 
 fn backout(
@@ -434,12 +451,59 @@ fn backout(
     Ok(())
 }
 
+fn align(
+    state: &mut CreateEdges,
+    edges: &mut Query<&mut Edge<Entity>>,
+    cursor: &Res<Cursor>,
+    creation_settings: &mut ResMut<CreationSettings>,
+    input_type: &ButtonInputType,
+    transform: &Query<&Transform>,
+) -> SelectionNodeResult {
+    if let Some(preview) = &mut state.preview_edge {
+        if preview.side == Side::end() {
+            // We currently have an active preview edge and are selecting for
+            // the second point in the edge. Shift means we want the current lane
+            // to align with either the X- or Y- axis.
+            let edge = edges.get(preview.edge).or_broken_query()?;
+            let end_anchor = edge.right();
+            if !cursor.is_placement_anchor(end_anchor) {
+                // We do not want to modify an existing anchor
+                return Ok(());
+            }
+
+            match input_type {
+                ButtonInputType::Pressed => {
+                    let Ok(delta) = transform.get(edge.left()).and_then(|tf| {
+                        transform
+                            .get(cursor.frame)
+                            .map(|c_tf| c_tf.translation - tf.translation)
+                    }) else {
+                        return Ok(());
+                    };
+                    if delta.x.abs() > delta.y.abs() {
+                        creation_settings.direction_alignment = vec![Vec2::new(0.0, -delta.y)];
+                    } else {
+                        creation_settings.direction_alignment = vec![Vec2::new(-delta.x, 0.0)];
+                    }
+                }
+                ButtonInputType::JustReleased => {
+                    creation_settings.reset();
+                }
+                ButtonInputType::JustPressed => {}
+            }
+        }
+    }
+
+    Ok(())
+}
+
 pub fn cleanup_create_edges(
     In(key): In<BufferKey<CreateEdges>>,
     mut access: BufferAccessMut<CreateEdges>,
     edges: Query<&'static Edge<Entity>>,
     mut commands: Commands,
     cursor: Res<Cursor>,
+    mut creation_settings: ResMut<CreationSettings>,
     mut dependents: Query<&mut Dependents>,
 ) -> SelectionNodeResult {
     let mut access = access.get_mut(&key).or_broken_buffer()?;
@@ -458,5 +522,7 @@ pub fn cleanup_create_edges(
         deps.0.clear();
     }
 
+    creation_settings.reset();
+
     Ok(())
 }

crates/rmf_site_editor/src/interaction/select_impl/mod.rs

@@ -26,6 +26,7 @@ pub use select_anchor::*;
 
 use anyhow::Error as Anyhow;
 
+use crate::interaction::*;
 use bevy::{
     ecs::{relationship::AncestorIter, system::SystemParam},
     prelude::*,
@@ -88,3 +89,27 @@ pub fn are_anchors_siblings(
 
     Ok(are_siblings)
 }
+
+#[derive(Default, Resource)]
+pub struct CreationSettings {
+    pub direction_alignment: Vec<Vec2>,
+}
+
+impl CreationSettings {
+    pub fn reset(&mut self) {
+        self.direction_alignment = Vec::new();
+    }
+}
+
+pub fn apply_creation_settings(
+    creation_settings: Res<CreationSettings>,
+    cursor: Res<Cursor>,
+    mut transform: Query<&mut Transform>,
+) {
+    let Ok(mut frame_tf) = transform.get_mut(cursor.frame) else {
+        return;
+    };
+    for alignment in &creation_settings.direction_alignment {
+        frame_tf.translation += alignment.extend(0.0);
+    }
+}

crates/rmf_site_editor/src/interaction/select_impl/place_object.rs

@@ -16,7 +16,9 @@
 */
 
 use crate::{
-    interaction::{spawn_place_object_2d_workflow, PlaceObject2d, PlaceObjectContinuity},
+    interaction::{
+        spawn_place_object_2d_workflow, ButtonInputType, PlaceObject2d, PlaceObjectContinuity,
+    },
     site::ModelInstance,
 };
 use bevy::{
@@ -44,7 +46,7 @@ impl Plugin for ObjectPlacementPlugin {
 pub struct ObjectPlacementServices {
     pub place_object_2d: Service<Option<Entity>, ()>,
     pub find_placement_2d: Service<(), Transform>,
-    pub on_key_code_2d: Service<KeyCode, SelectionNodeResult>,
+    pub on_key_code_2d: Service<(KeyCode, ButtonInputType), SelectionNodeResult>,
 }
 
 impl ObjectPlacementServices {

crates/rmf_site_editor/src/interaction/select_impl/place_object_2d.rs

@@ -42,10 +42,7 @@ pub fn spawn_place_object_2d_workflow(app: &mut App) -> ObjectPlacementServices
     let on_key_code_2d = app.spawn_service(on_keyboard_for_place_object_2d.into_blocking_service());
     let cleanup = app.spawn_service(place_object_2d_cleanup.into_blocking_service());
 
-    let keyboard_just_pressed = app
-        .world()
-        .resource::<KeyboardServices>()
-        .keyboard_just_pressed;
+    let keyboard_pressed = app.world().resource::<KeyboardServices>().keyboard_pressed;
 
     let place_object_2d = app
         .world_mut()
@@ -55,7 +52,7 @@ pub fn spawn_place_object_2d_workflow(app: &mut App) -> ObjectPlacementServices
             placement_chosen,
             on_key_code_2d,
             cleanup,
-            keyboard_just_pressed,
+            keyboard_pressed,
         ));
 
     ObjectPlacementServices {
@@ -69,9 +66,9 @@ pub fn build_2d_placement_workflow<State: 'static + Send + Sync>(
     setup: Service<BufferKey<State>, SelectionNodeResult>,
     find_placement: Service<(), Transform>,
     placement_chosen: Service<(Transform, BufferKey<State>), PlaceObjectResult>,
-    handle_key_code: Service<KeyCode, SelectionNodeResult>,
+    handle_key_code: Service<(KeyCode, ButtonInputType), SelectionNodeResult>,
     cleanup: Service<BufferKey<State>, ()>,
-    keyboard_just_pressed: Service<(), (), StreamOf<KeyCode>>,
+    keyboard_pressed: Service<(), (), StreamOf<(KeyCode, ButtonInputType)>>,
 ) -> impl FnOnce(Scope<Option<Entity>, ()>, &mut Builder) {
     move |scope, builder| {
         let buffer = builder.create_buffer::<State>(BufferSettings::keep_last(1));
@@ -114,7 +111,7 @@ pub fn build_2d_placement_workflow<State: 'static + Send + Sync>(
 
         let keyboard_node = setup_finished
             .clone_chain(builder)
-            .then_node(keyboard_just_pressed);
+            .then_node(keyboard_pressed);
         keyboard_node
             .streams
             .chain(builder)
@@ -218,8 +215,10 @@ pub fn place_object_2d_find_placement(
     }
 }
 
-pub fn on_keyboard_for_place_object_2d(In(key): In<KeyCode>) -> SelectionNodeResult {
-    if matches!(key, KeyCode::Escape) {
+pub fn on_keyboard_for_place_object_2d(
+    In((key, input)): In<(KeyCode, ButtonInputType)>,
+) -> SelectionNodeResult {
+    if matches!(key, KeyCode::Escape) && matches!(input, ButtonInputType::JustPressed) {
         // Simply end the workflow if the escape key was pressed
         info!("Exiting 2D object placement");
         return Err(None);

crates/rmf_site_editor/src/interaction/select_impl/select_anchor.rs

@@ -66,7 +66,7 @@ impl Plugin for AnchorSelectionPlugin {
 pub struct AnchorSelectionHelpers {
     pub anchor_select_stream: Service<(), (), SelectionStreams>,
     pub anchor_cursor_transform: Service<(), ()>,
-    pub keyboard_just_pressed: Service<(), (), StreamOf<KeyCode>>,
+    pub keyboard_pressed: Service<(), (), StreamOf<(KeyCode, ButtonInputType)>>,
     pub cleanup_anchor_selection: Service<(), ()>,
 }
 
@@ -83,15 +83,12 @@ impl AnchorSelectionHelpers {
             .world_mut()
             .spawn_service(cleanup_anchor_selection.into_blocking_service());
 
-        let keyboard_just_pressed = app
-            .world()
-            .resource::<KeyboardServices>()
-            .keyboard_just_pressed;
+        let keyboard_pressed = app.world().resource::<KeyboardServices>().keyboard_pressed;
 
         Self {
             anchor_select_stream,
             anchor_cursor_transform,
-            keyboard_just_pressed,
+            keyboard_pressed,
             cleanup_anchor_selection,
         }
     }
@@ -102,7 +99,10 @@ impl AnchorSelectionHelpers {
         state_setup: Service<BufferKey<State>, SelectionNodeResult>,
         update_preview: Service<(Hover, BufferKey<State>), SelectionNodeResult>,
         update_current: Service<(SelectionCandidate, BufferKey<State>), SelectionNodeResult>,
-        handle_key_code: Service<(KeyCode, BufferKey<State>), SelectionNodeResult>,
+        handle_key_code: Service<
+            ((KeyCode, ButtonInputType), BufferKey<State>),
+            SelectionNodeResult,
+        >,
         cleanup_state: Service<BufferKey<State>, SelectionNodeResult>,
         world: &mut World,
     ) -> Service<Option<Entity>, ()> {
@@ -115,7 +115,7 @@ impl AnchorSelectionHelpers {
             cleanup_state,
             self.anchor_cursor_transform,
             self.anchor_select_stream,
-            self.keyboard_just_pressed,
+            self.keyboard_pressed,
             self.cleanup_anchor_selection,
         ))
     }
@@ -375,11 +375,11 @@ pub fn build_anchor_selection_workflow<State: 'static + Send + Sync>(
     state_setup: Service<BufferKey<State>, SelectionNodeResult>,
     update_preview: Service<(Hover, BufferKey<State>), SelectionNodeResult>,
     update_current: Service<(SelectionCandidate, BufferKey<State>), SelectionNodeResult>,
-    handle_key_code: Service<(KeyCode, BufferKey<State>), SelectionNodeResult>,
+    handle_key_code: Service<((KeyCode, ButtonInputType), BufferKey<State>), SelectionNodeResult>,
     cleanup_state: Service<BufferKey<State>, SelectionNodeResult>,
     anchor_cursor_transform: Service<(), ()>,
     anchor_select_stream: Service<(), (), SelectionStreams>,
-    keyboard_just_pressed: Service<(), (), StreamOf<KeyCode>>,
+    keyboard_pressed: Service<(), (), StreamOf<(KeyCode, ButtonInputType)>>,
     cleanup_anchor_selection: Service<(), ()>,
 ) -> impl FnOnce(Scope<Option<Entity>, ()>, &mut Builder) {
     move |scope, builder| {
@@ -442,7 +442,7 @@ pub fn build_anchor_selection_workflow<State: 'static + Send + Sync>(
 
         let keyboard = begin_input_services
             .clone_chain(builder)
-            .then_node(keyboard_just_pressed);
+            .then_node(keyboard_pressed);
         keyboard
             .streams
             .chain(builder)
@@ -623,7 +623,11 @@ impl<'w, 's> SelectionFilter for AnchorFilter<'w, 's> {
 
         // There was no anchor currently hovered which means we need to create
         // a new provisional anchor.
-        let Ok(tf) = self.transforms.get(self.cursor.frame) else {
+        let tf = if let Ok(placement_tf) = self.transforms.get(self.cursor.level_anchor_placement) {
+            placement_tf
+        } else if let Ok(cursor_tf) = self.transforms.get(self.cursor.frame) {
+            cursor_tf
+        } else {
             error!("Cannot find cursor transform");
             return None;
         };
@@ -781,8 +785,10 @@ fn compute_parent_inverse_pose(
     Some(pose.align_with(&Transform::from_matrix((inv_tf * goal_tf).into())))
 }
 
-pub fn exit_on_esc<T>(In((button, _)): In<(KeyCode, BufferKey<T>)>) -> SelectionNodeResult {
-    if matches!(button, KeyCode::Escape) {
+pub fn exit_on_esc<T>(
+    In(((button, input_type), _)): In<((KeyCode, ButtonInputType), BufferKey<T>)>,
+) -> SelectionNodeResult {
+    if matches!(button, KeyCode::Escape) && matches!(input_type, ButtonInputType::JustPressed) {
         // The escape key was pressed so we should exit this mode
         return Err(None);
     }