refactor(keybinding): Optimize KeyBindingManager with Trie (B.3)

src/keybinding.cpp

@@ -247,7 +247,9 @@ KeyBinding::KeyBinding(const std::string& key_str, std::string cmd_name, std::st
 // ============================================================================
 
 KeyBindingManager::KeyBindingManager(CommandSystem* command_system) 
-    : sequence_timeout_(std::chrono::milliseconds(1000)), command_system_(command_system) {
+    : root_(std::make_unique<TrieNode>()), // Initialize root_
+      sequence_timeout_(std::chrono::milliseconds(1000)), 
+      command_system_(command_system) {
 }
 
 void KeyBindingManager::bind(const KeySequence& sequence, std::string cmd_name, const std::string& description) {
@@ -255,7 +257,15 @@ void KeyBindingManager::bind(const KeySequence& sequence, std::string cmd_name,
         return;
     }
     
-    bindings_[sequence] = KeyBinding(sequence, std::move(cmd_name), description);
+    TrieNode* current = root_.get();
+    for (const Key& key : sequence.keys()) {
+        if (current->children.find(key) == current->children.end()) {
+            current->children[key] = std::make_unique<TrieNode>();
+        }
+        current = current->children[key].get();
+        current->ref_count++; // Increment ref count for this path
+    }
+    current->binding = KeyBinding(sequence, std::move(cmd_name), description);
 }
 
 void KeyBindingManager::bind(const std::string& key_str, std::string cmd_name, const std::string& description) {
@@ -263,9 +273,47 @@ void KeyBindingManager::bind(const std::string& key_str, std::string cmd_name, c
 }
 
 void KeyBindingManager::unbind(const KeySequence& sequence) {
-    bindings_.erase(sequence);
+    if (sequence.empty()) return;
+
+    // Find the node corresponding to the end of the sequence
+    TrieNode* current = root_.get();
+    std::vector<TrieNode*> path_nodes; // To decrement ref_counts later
+    path_nodes.push_back(current);
+
+    for (const Key& key : sequence.keys()) {
+        if (current->children.find(key) == current->children.end()) {
+            return; // Sequence not found
+        }
+        current = current->children[key].get();
+        path_nodes.push_back(current);
+    }
+
+    if (!current->binding) {
+        return; // No exact binding at this node
+    }
+
+    // Remove the binding
+    current->binding.reset();
+
+    // Decrement ref counts along the path, and remove nodes if ref_count becomes 0
+    // Traverse backwards from the node just before the binding
+    for (int i = static_cast<int>(sequence.length()) - 1; i >= 0; --i) {
+        TrieNode* parent = path_nodes[i];
+        TrieNode* child = path_nodes[i+1];
+        
+        child->ref_count--; // Decrement ref count
+        
+        // If child's ref_count is 0 AND it has no own binding, we can remove it
+        if (child->ref_count == 0 && !child->binding.has_value() && child->children.empty()) {
+            parent->children.erase(sequence.keys()[i]);
+        } else {
+            // Stop if the node is still referenced or has a binding/children
+            break;
+        }
+    }
 }
 
+
 void KeyBindingManager::unbind(const std::string& key_str) {
     unbind(KeySequence(key_str));
 }
@@ -284,7 +332,7 @@ KeyResult KeyBindingManager::process_key(const Key& key) {
     sequence_start_time_ = std::chrono::steady_clock::now();
     
     // Check for exact binding
-    auto exact_binding = find_exact_binding(current_sequence_);
+    std::optional<KeyBinding> exact_binding = find_exact_binding(current_sequence_);
     if (exact_binding.has_value()) {
         // Found exact match
         clear_partial_sequence();
@@ -317,15 +365,21 @@ KeyResult KeyBindingManager::process_key(const std::string& key_str) {
 }
 
 bool KeyBindingManager::has_binding(const KeySequence& sequence) const {
-    return has_exact_binding(sequence) || has_prefix_bindings(sequence);
+    // A sequence has a binding if it's an exact binding or a prefix for other bindings
+    TrieNode* node = find_node(sequence);
+    if (!node) return false;
+    
+    return node->binding.has_value() || !node->children.empty();
 }
 
 bool KeyBindingManager::has_exact_binding(const KeySequence& sequence) const {
-    return bindings_.find(sequence) != bindings_.end();
+    TrieNode* node = find_node(sequence);
+    return node && node->binding.has_value();
 }
 
 bool KeyBindingManager::has_prefix_bindings(const KeySequence& sequence) const {
-    return has_prefix_bindings_impl(sequence);
+    TrieNode* node = find_node(sequence);
+    return node && !node->children.empty();
 }
 
 void KeyBindingManager::clear_partial_sequence() {
@@ -347,12 +401,7 @@ std::string KeyBindingManager::current_sequence_display() const {
 
 std::vector<KeyBinding> KeyBindingManager::get_all_bindings() const {
     std::vector<KeyBinding> result;
-    result.reserve(bindings_.size());
-    
-    for (const auto& [sequence, binding] : bindings_) {
-        result.push_back(binding);
-    }
-    
+    collect_bindings(root_.get(), KeySequence(), result);
     return result;
 }
 
@@ -369,22 +418,46 @@ bool KeyBindingManager::has_sequence_timed_out() const {
     return (now - sequence_start_time_) > sequence_timeout_;
 }
 
+// Helper function to find a node in the trie
+KeyBindingManager::TrieNode* KeyBindingManager::find_node(const KeySequence& sequence) const {
+    TrieNode* current = root_.get();
+    for (const Key& key : sequence.keys()) {
+        if (current->children.find(key) == current->children.end()) {
+            return nullptr;
+        }
+        current = current->children[key].get();
+    }
+    return current;
+}
+
+// Helper to traverse trie and collect all bindings
+void KeyBindingManager::collect_bindings(TrieNode* node, KeySequence current_prefix, std::vector<KeyBinding>& result) const {
+    if (!node) return;
+
+    if (node->binding.has_value()) {
+        result.push_back(node->binding.value());
+    }
+
+    for (const auto& pair : node->children) {
+        KeySequence next_prefix = current_prefix;
+        next_prefix.add_key(pair.first);
+        collect_bindings(pair.second.get(), next_prefix, result);
+    }
+}
+
 std::optional<KeyBinding> KeyBindingManager::find_exact_binding(const KeySequence& sequence) const {
-    auto it = bindings_.find(sequence);
-    if (it != bindings_.end()) {
-        return it->second;
+    TrieNode* node = find_node(sequence);
+    if (node && node->binding.has_value()) {
+        return node->binding;
     }
     return std::nullopt;
 }
 
 bool KeyBindingManager::has_prefix_bindings_impl(const KeySequence& sequence) const {
-    // Check if any registered binding has this sequence as a prefix
-    for (const auto& [binding_sequence, binding] : bindings_) {
-        if (sequence.is_prefix_of(binding_sequence) && sequence != binding_sequence) {
-            return true;
-        }
-    }
-    return false;
+    TrieNode* node = find_node(sequence);
+    return node && !node->children.empty();
 }
 
+} // namespace lumacs
+
 } // namespace lumacs