chore: improve code quanity on ReplCommand/State

3 files changed, 64 insertions, 37 deletions

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 5220c5d..d6ea6db 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -18,7 +18,7 @@ use inquire::{Confirm, Select, Text};
 use is_terminal::IsTerminal;
 use parking_lot::RwLock;
 use serde::Deserialize;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::{
     env,
     fs::{create_dir_all, read_dir, read_to_string, remove_file, File, OpenOptions},
@@ -915,7 +915,7 @@ impl Keybindings {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum State {
     Normal,
     Role,
@@ -924,6 +924,45 @@ pub enum State {
     Session,
 }
 
+impl State {
+    pub fn all() -> Vec<Self> {
+        vec![
+            Self::Normal,
+            Self::Role,
+            Self::EmptySession,
+            Self::EmptySessionWithRole,
+            Self::Session,
+        ]
+    }
+
+    pub fn in_session() -> Vec<Self> {
+        vec![
+            Self::EmptySession,
+            Self::EmptySessionWithRole,
+            Self::Session,
+        ]
+    }
+
+    pub fn notin_session() -> Vec<Self> {
+        let excludes: HashSet<_> = Self::in_session().into_iter().collect();
+        Self::all()
+            .into_iter()
+            .filter(|v| !excludes.contains(v))
+            .collect()
+    }
+
+    pub fn can_change_role() -> Vec<Self> {
+        Self::all()
+            .into_iter()
+            .filter(|v| *v != Self::Session)
+            .collect()
+    }
+
+    pub fn in_role() -> Vec<Self> {
+        vec![Self::Role, Self::EmptySessionWithRole]
+    }
+}
+
 fn create_config_file(config_path: &Path) -> Result<()> {
     let ans = Confirm::new("No config file, create a new one?")
         .with_default(true)
diff --git a/src/repl/completer.rs b/src/repl/completer.rs
index 4d337df..aea1ee6 100644
--- a/src/repl/completer.rs
+++ b/src/repl/completer.rs
@@ -33,7 +33,7 @@ impl Completer for ReplCompleter {
             .commands
             .iter()
             .filter(|cmd| {
-                if cmd.unavailable(&state) {
+                if !cmd.is_valid(&state) {
                     return false;
                 }
                 let line = parts
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index 118ff0c..d6fa7f2 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -26,47 +26,35 @@ const MENU_NAME: &str = "completion_menu";
 
 lazy_static! {
     static ref REPL_COMMANDS: [ReplCommand; 13] = [
-        ReplCommand::new(".help", "Print this help message", vec![]),
-        ReplCommand::new(".info", "Print system info", vec![]),
-        ReplCommand::new(".model", "Switch LLM model", vec![]),
-        ReplCommand::new(".role", "Use a role", vec![State::Session]),
-        ReplCommand::new(
-            ".info role",
-            "Show role info",
-            vec![State::Normal, State::EmptySession, State::Session]
-        ),
-        ReplCommand::new(
-            ".exit role",
-            "Leave current role",
-            vec![State::Normal, State::EmptySession, State::Session]
-        ),
+        ReplCommand::new(".help", "Print this help message", State::all()),
+        ReplCommand::new(".info", "Print system info", State::all()),
+        ReplCommand::new(".model", "Switch LLM model", State::all()),
+        ReplCommand::new(".role", "Use a role", State::can_change_role()),
+        ReplCommand::new(".info role", "Show role info", State::in_role(),),
+        ReplCommand::new(".exit role", "Leave current role", State::in_role(),),
         ReplCommand::new(
             ".session",
             "Start a context-aware chat session",
-            vec![
-                State::EmptySession,
-                State::EmptySessionWithRole,
-                State::Session
-            ]
-        ),
-        ReplCommand::new(
-            ".info session",
-            "Show session info",
-            vec![State::Normal, State::Role]
+            State::notin_session(),
         ),
+        ReplCommand::new(".info session", "Show session info", State::in_session(),),
         ReplCommand::new(
             ".exit session",
             "End the current session",
-            vec![State::Normal, State::Role]
+            State::in_session(),
         ),
         ReplCommand::new(
             ".file",
             "Attach files to the message and then submit it",
-            vec![]
+            State::all()
+        ),
+        ReplCommand::new(".set", "Modify the configuration parameters", State::all()),
+        ReplCommand::new(
+            ".copy",
+            "Copy the last reply to the clipboard",
+            State::all()
         ),
-        ReplCommand::new(".set", "Modify the configuration parameters", vec![]),
-        ReplCommand::new(".copy", "Copy the last reply to the clipboard", vec![]),
-        ReplCommand::new(".exit", "Exit the REPL", vec![]),
+        ReplCommand::new(".exit", "Exit the REPL", State::all()),
     ];
     static ref COMMAND_RE: Regex = Regex::new(r"^\s*(\.\S*)\s*").unwrap();
     static ref MULTILINE_RE: Regex = Regex::new(r"(?s)^\s*:::\s*(.*)\s*:::\s*$").unwrap();
@@ -361,20 +349,20 @@ Type ".help" for more information.
 pub struct ReplCommand {
     name: &'static str,
     description: &'static str,
-    unavailable_states: Vec<State>,
+    valid_states: Vec<State>,
 }
 
 impl ReplCommand {
-    fn new(name: &'static str, desc: &'static str, unavailable_states: Vec<State>) -> Self {
+    fn new(name: &'static str, desc: &'static str, valid_states: Vec<State>) -> Self {
         Self {
             name,
             description: desc,
-            unavailable_states,
+            valid_states,
         }
     }
 
-    fn unavailable(&self, state: &State) -> bool {
-        self.unavailable_states.contains(state)
+    fn is_valid(&self, state: &State) -> bool {
+        self.valid_states.contains(state)
     }
 }