From 2c2e5144c9c034b3e15c480cab8711345671d23f Mon Sep 17 00:00:00 2001
From: Conrad Irwin <conrad.irwin@gmail.com>
Date: Mon, 4 Dec 2023 21:28:37 +0000
Subject: [PATCH] Fix context key matching

* You need to check all layers of the context stack
* When in command, the context should be based on where focus was (to
  match `available_actions`.
---
 .../command_palette2/src/command_palette.rs   |  6 ++++-
 crates/gpui2/src/key_dispatch.rs              | 19 +++++++++++---
 crates/gpui2/src/window.rs                    | 26 ++++++++++++++++---
 crates/ui2/src/components/keybinding.rs       | 15 ++++++++---
 4 files changed, 54 insertions(+), 12 deletions(-)

diff --git a/crates/command_palette2/src/command_palette.rs b/crates/command_palette2/src/command_palette.rs
index 04688b0549..a2abadd5fd 100644
--- a/crates/command_palette2/src/command_palette.rs
+++ b/crates/command_palette2/src/command_palette.rs
@@ -311,7 +311,11 @@ impl PickerDelegate for CommandPaletteDelegate {
                         command.name.clone(),
                         r#match.positions.clone(),
                     ))
-                    .children(KeyBinding::for_action(&*command.action, cx)),
+                    .children(KeyBinding::for_action_in(
+                        &*command.action,
+                        &self.previous_focus_handle,
+                        cx,
+                    )),
             ),
         )
     }
diff --git a/crates/gpui2/src/key_dispatch.rs b/crates/gpui2/src/key_dispatch.rs
index 1ab99ec487..95915b98ed 100644
--- a/crates/gpui2/src/key_dispatch.rs
+++ b/crates/gpui2/src/key_dispatch.rs
@@ -16,7 +16,7 @@ pub struct DispatchNodeId(usize);
 
 pub(crate) struct DispatchTree {
     node_stack: Vec<DispatchNodeId>,
-    context_stack: Vec<KeyContext>,
+    pub(crate) context_stack: Vec<KeyContext>,
     nodes: Vec<DispatchNode>,
     focusable_node_ids: HashMap<FocusId, DispatchNodeId>,
     keystroke_matchers: HashMap<SmallVec<[KeyContext; 4]>, KeystrokeMatcher>,
@@ -163,13 +163,24 @@ impl DispatchTree {
         actions
     }
 
-    pub fn bindings_for_action(&self, action: &dyn Action) -> Vec<KeyBinding> {
+    pub fn bindings_for_action(
+        &self,
+        action: &dyn Action,
+        context_stack: &Vec<KeyContext>,
+    ) -> Vec<KeyBinding> {
         self.keymap
             .lock()
             .bindings_for_action(action.type_id())
             .filter(|candidate| {
-                candidate.action.partial_eq(action)
-                    && candidate.matches_context(&self.context_stack)
+                if !candidate.action.partial_eq(action) {
+                    return false;
+                }
+                for i in 1..context_stack.len() {
+                    if candidate.matches_context(&context_stack[0..i]) {
+                        return true;
+                    }
+                }
+                return false;
             })
             .cloned()
             .collect()
diff --git a/crates/gpui2/src/window.rs b/crates/gpui2/src/window.rs
index 5724f1e070..b88e89ef55 100644
--- a/crates/gpui2/src/window.rs
+++ b/crates/gpui2/src/window.rs
@@ -1492,10 +1492,28 @@ impl<'a> WindowContext<'a> {
     }
 
     pub fn bindings_for_action(&self, action: &dyn Action) -> Vec<KeyBinding> {
-        self.window
-            .current_frame
-            .dispatch_tree
-            .bindings_for_action(action)
+        self.window.current_frame.dispatch_tree.bindings_for_action(
+            action,
+            &self.window.current_frame.dispatch_tree.context_stack,
+        )
+    }
+
+    pub fn bindings_for_action_in(
+        &self,
+        action: &dyn Action,
+        focus_handle: &FocusHandle,
+    ) -> Vec<KeyBinding> {
+        let dispatch_tree = &self.window.previous_frame.dispatch_tree;
+
+        let Some(node_id) = dispatch_tree.focusable_node_id(focus_handle.id) else {
+            return vec![];
+        };
+        let context_stack = dispatch_tree
+            .dispatch_path(node_id)
+            .into_iter()
+            .map(|node_id| dispatch_tree.node(node_id).context.clone())
+            .collect();
+        dispatch_tree.bindings_for_action(action, &context_stack)
     }
 
     pub fn listener_for<V: Render, E>(
diff --git a/crates/ui2/src/components/keybinding.rs b/crates/ui2/src/components/keybinding.rs
index 993e2f323e..c4054fa1a4 100644
--- a/crates/ui2/src/components/keybinding.rs
+++ b/crates/ui2/src/components/keybinding.rs
@@ -1,5 +1,5 @@
 use crate::{h_stack, prelude::*, Icon, IconElement, IconSize};
-use gpui::{relative, rems, Action, Div, IntoElement, Keystroke};
+use gpui::{relative, rems, Action, Div, FocusHandle, IntoElement, Keystroke};
 
 #[derive(IntoElement, Clone)]
 pub struct KeyBinding {
@@ -49,12 +49,21 @@ impl RenderOnce for KeyBinding {
 
 impl KeyBinding {
     pub fn for_action(action: &dyn Action, cx: &mut WindowContext) -> Option<Self> {
-        // todo! this last is arbitrary, we want to prefer users key bindings over defaults,
-        // and vim over normal (in vim mode), etc.
         let key_binding = cx.bindings_for_action(action).last().cloned()?;
         Some(Self::new(key_binding))
     }
 
+    // like for_action(), but lets you specify the context from which keybindings
+    // are matched.
+    pub fn for_action_in(
+        action: &dyn Action,
+        focus: &FocusHandle,
+        cx: &mut WindowContext,
+    ) -> Option<Self> {
+        let key_binding = cx.bindings_for_action_in(action, focus).last().cloned()?;
+        Some(Self::new(key_binding))
+    }
+
     fn icon_for_key(keystroke: &Keystroke) -> Option<Icon> {
         let mut icon: Option<Icon> = None;