rush_sh/
state.rs

1use std::collections::{HashMap, HashSet};
2use std::rc::Rc;
3use std::cell::RefCell;
4use super::parser::Ast;
5use std::env;
6use std::io::IsTerminal;
7
8#[derive(Debug, Clone)]
9pub struct ColorScheme {
10    /// ANSI color code for prompt
11    pub prompt: String,
12    /// ANSI color code for error messages
13    pub error: String,
14    /// ANSI color code for success messages
15    pub success: String,
16    /// ANSI color code for builtin command output
17    pub builtin: String,
18    /// ANSI color code for directory listings
19    pub directory: String,
20}
21
22impl Default for ColorScheme {
23    fn default() -> Self {
24        Self {
25            prompt: "\x1b[32m".to_string(),    // Green
26            error: "\x1b[31m".to_string(),     // Red
27            success: "\x1b[32m".to_string(),   // Green
28            builtin: "\x1b[36m".to_string(),   // Cyan
29            directory: "\x1b[34m".to_string(), // Blue
30        }
31    }
32}
33
34#[derive(Debug, Clone)]
35pub struct ShellState {
36    /// Shell variables (local to the shell session)
37    pub variables: HashMap<String, String>,
38    /// Which variables are exported to child processes
39    pub exported: HashSet<String>,
40    /// Last exit code ($?)
41    pub last_exit_code: i32,
42    /// Shell process ID ($$)
43    pub shell_pid: u32,
44    /// Script name or command ($0)
45    pub script_name: String,
46    /// Directory stack for pushd/popd
47    pub dir_stack: Vec<String>,
48    /// Command aliases
49    pub aliases: HashMap<String, String>,
50    /// Whether colors are enabled
51    pub colors_enabled: bool,
52    /// Current color scheme
53    pub color_scheme: ColorScheme,
54    /// Positional parameters ($1, $2, $3, ...)
55    pub positional_params: Vec<String>,
56    /// Function definitions
57    pub functions: HashMap<String, Ast>,
58    /// Local variable stack for function scoping
59    pub local_vars: Vec<HashMap<String, String>>,
60    /// Function call depth for local scope management
61    pub function_depth: usize,
62    /// Maximum allowed recursion depth
63    pub max_recursion_depth: usize,
64    /// Flag to indicate if we're currently returning from a function
65    pub returning: bool,
66    /// Return value when returning from a function
67    pub return_value: Option<i32>,
68    /// Output capture buffer for command substitution
69    pub capture_output: Option<Rc<RefCell<Vec<u8>>>>,
70    /// Whether to use condensed cwd display in prompt
71    pub condensed_cwd: bool,
72}
73
74impl ShellState {
75    pub fn new() -> Self {
76        let shell_pid = std::process::id();
77
78        // Check NO_COLOR environment variable (respects standard)
79        let no_color = std::env::var("NO_COLOR").is_ok();
80
81        // Check RUSH_COLORS environment variable for explicit control
82        let rush_colors = std::env::var("RUSH_COLORS")
83            .map(|v| v.to_lowercase())
84            .unwrap_or_else(|_| "auto".to_string());
85
86        let colors_enabled = match rush_colors.as_str() {
87            "1" | "true" | "on" | "enable" => !no_color && std::io::stdout().is_terminal(),
88            "0" | "false" | "off" | "disable" => false,
89            "auto" => !no_color && std::io::stdout().is_terminal(),
90            _ => !no_color && std::io::stdout().is_terminal(),
91        };
92
93        // Check RUSH_CONDENSED environment variable for cwd display preference
94        let rush_condensed = std::env::var("RUSH_CONDENSED")
95            .map(|v| v.to_lowercase())
96            .unwrap_or_else(|_| "true".to_string());
97
98        let condensed_cwd = match rush_condensed.as_str() {
99            "1" | "true" | "on" | "enable" => true,
100            "0" | "false" | "off" | "disable" => false,
101            _ => true, // Default to condensed for backward compatibility
102        };
103
104        Self {
105            variables: HashMap::new(),
106            exported: HashSet::new(),
107            last_exit_code: 0,
108            shell_pid,
109            script_name: "rush".to_string(),
110            dir_stack: Vec::new(),
111            aliases: HashMap::new(),
112            colors_enabled,
113            color_scheme: ColorScheme::default(),
114            positional_params: Vec::new(),
115            functions: HashMap::new(),
116            local_vars: Vec::new(),
117            function_depth: 0,
118            max_recursion_depth: 500, // Default recursion limit (reduced to avoid Rust stack overflow)
119            returning: false,
120            return_value: None,
121            capture_output: None,
122            condensed_cwd,
123        }
124    }
125
126    /// Get a variable value, checking local scopes first, then shell variables, then environment
127    pub fn get_var(&self, name: &str) -> Option<String> {
128        // Handle special variables (these are never local)
129        match name {
130            "?" => Some(self.last_exit_code.to_string()),
131            "$" => Some(self.shell_pid.to_string()),
132            "0" => Some(self.script_name.clone()),
133            "*" => {
134                // $* - all positional parameters as single string (space-separated)
135                if self.positional_params.is_empty() {
136                    Some("".to_string())
137                } else {
138                    Some(self.positional_params.join(" "))
139                }
140            }
141            "@" => {
142                // $@ - all positional parameters as separate words (but returns as single string for compatibility)
143                if self.positional_params.is_empty() {
144                    Some("".to_string())
145                } else {
146                    Some(self.positional_params.join(" "))
147                }
148            }
149            "#" => Some(self.positional_params.len().to_string()),
150            _ => {
151                // Handle positional parameters $1, $2, $3, etc. (these are never local)
152                if let Ok(index) = name.parse::<usize>()
153                    && index > 0
154                    && index <= self.positional_params.len()
155                {
156                    return Some(self.positional_params[index - 1].clone());
157                }
158
159                // Check local scopes first, then shell variables, then environment
160                // Search local scopes from innermost to outermost
161                for scope in self.local_vars.iter().rev() {
162                    if let Some(value) = scope.get(name) {
163                        return Some(value.clone());
164                    }
165                }
166
167                // Check shell variables
168                if let Some(value) = self.variables.get(name) {
169                    Some(value.clone())
170                } else {
171                    // Fall back to environment variables
172                    env::var(name).ok()
173                }
174            }
175        }
176    }
177
178    /// Set a shell variable (updates local scope if variable exists there, otherwise sets globally)
179    pub fn set_var(&mut self, name: &str, value: String) {
180        // Check if this variable exists in any local scope
181        // If it does, update it there instead of setting globally
182        for scope in self.local_vars.iter_mut().rev() {
183            if scope.contains_key(name) {
184                scope.insert(name.to_string(), value);
185                return;
186            }
187        }
188        
189        // Variable doesn't exist in local scopes, set it globally
190        self.variables.insert(name.to_string(), value);
191    }
192
193    /// Remove a shell variable
194    pub fn unset_var(&mut self, name: &str) {
195        self.variables.remove(name);
196        self.exported.remove(name);
197    }
198
199    /// Mark a variable as exported
200    pub fn export_var(&mut self, name: &str) {
201        if self.variables.contains_key(name) {
202            self.exported.insert(name.to_string());
203        }
204    }
205
206    /// Set and export a variable
207    pub fn set_exported_var(&mut self, name: &str, value: String) {
208        self.set_var(name, value);
209        self.export_var(name);
210    }
211
212    /// Get all environment variables for child processes (exported + inherited)
213    pub fn get_env_for_child(&self) -> HashMap<String, String> {
214        let mut child_env = HashMap::new();
215
216        // Add all current environment variables
217        for (key, value) in env::vars() {
218            child_env.insert(key, value);
219        }
220
221        // Override with exported shell variables
222        for var_name in &self.exported {
223            if let Some(value) = self.variables.get(var_name) {
224                child_env.insert(var_name.clone(), value.clone());
225            }
226        }
227
228        child_env
229    }
230
231    /// Update the last exit code
232    pub fn set_last_exit_code(&mut self, code: i32) {
233        self.last_exit_code = code;
234    }
235
236    /// Set the script name ($0)
237    pub fn set_script_name(&mut self, name: &str) {
238        self.script_name = name.to_string();
239    }
240
241    /// Get the condensed current working directory for the prompt
242    pub fn get_condensed_cwd(&self) -> String {
243        match std::env::current_dir() {
244            Ok(path) => {
245                let path_str = path.to_string_lossy();
246                let components: Vec<&str> = path_str.split('/').collect();
247                if components.is_empty() || (components.len() == 1 && components[0].is_empty()) {
248                    return "/".to_string();
249                }
250                let mut result = String::new();
251                for (i, comp) in components.iter().enumerate() {
252                    if comp.is_empty() {
253                        continue; // skip leading empty component
254                    }
255                    if i == components.len() - 1 {
256                        result.push('/');
257                        result.push_str(comp);
258                    } else {
259                        result.push('/');
260                        if let Some(first) = comp.chars().next() {
261                            result.push(first);
262                        }
263                    }
264                }
265                if result.is_empty() {
266                    "/".to_string()
267                } else {
268                    result
269                }
270            }
271            Err(_) => "/?".to_string(), // fallback if can't get cwd
272        }
273    }
274
275    /// Get the full current working directory for the prompt
276    pub fn get_full_cwd(&self) -> String {
277        match std::env::current_dir() {
278            Ok(path) => path.to_string_lossy().to_string(),
279            Err(_) => "/?".to_string(), // fallback if can't get cwd
280        }
281    }
282
283    /// Get the user@hostname string for the prompt
284    pub fn get_user_hostname(&self) -> String {
285        let user = env::var("USER").unwrap_or_else(|_| "user".to_string());
286        let hostname = env::var("HOSTNAME").unwrap_or_else(|_| "hostname".to_string());
287        format!("{}@{}", user, hostname)
288    }
289
290    /// Get the full prompt string
291    pub fn get_prompt(&self) -> String {
292        let user = env::var("USER").unwrap_or_else(|_| "user".to_string());
293        let prompt_char = if user == "root" { "#" } else { "$" };
294        let cwd = if self.condensed_cwd {
295            self.get_condensed_cwd()
296        } else {
297            self.get_full_cwd()
298        };
299        format!(
300            "{}:{} {} ",
301            self.get_user_hostname(),
302            cwd,
303            prompt_char
304        )
305    }
306
307    /// Set an alias
308    pub fn set_alias(&mut self, name: &str, value: String) {
309        self.aliases.insert(name.to_string(), value);
310    }
311
312    /// Get an alias value
313    pub fn get_alias(&self, name: &str) -> Option<&String> {
314        self.aliases.get(name)
315    }
316
317    /// Remove an alias
318    pub fn remove_alias(&mut self, name: &str) {
319        self.aliases.remove(name);
320    }
321
322    /// Get all aliases
323    pub fn get_all_aliases(&self) -> &HashMap<String, String> {
324        &self.aliases
325    }
326
327    /// Set positional parameters
328    pub fn set_positional_params(&mut self, params: Vec<String>) {
329        self.positional_params = params;
330    }
331
332    /// Get positional parameters
333    #[allow(dead_code)]
334    pub fn get_positional_params(&self) -> &[String] {
335        &self.positional_params
336    }
337
338    /// Shift positional parameters (remove first n parameters)
339    pub fn shift_positional_params(&mut self, count: usize) {
340        if count > 0 {
341            for _ in 0..count {
342                if !self.positional_params.is_empty() {
343                    self.positional_params.remove(0);
344                }
345            }
346        }
347    }
348
349    /// Add a positional parameter at the end
350    #[allow(dead_code)]
351    pub fn push_positional_param(&mut self, param: String) {
352        self.positional_params.push(param);
353    }
354
355    /// Define a function
356    pub fn define_function(&mut self, name: String, body: Ast) {
357        self.functions.insert(name, body);
358    }
359
360    /// Get a function definition
361    pub fn get_function(&self, name: &str) -> Option<&Ast> {
362        self.functions.get(name)
363    }
364
365    /// Remove a function definition
366    #[allow(dead_code)]
367    pub fn remove_function(&mut self, name: &str) {
368        self.functions.remove(name);
369    }
370
371    /// Get all function names
372    #[allow(dead_code)]
373    pub fn get_function_names(&self) -> Vec<&String> {
374        self.functions.keys().collect()
375    }
376
377    /// Push a new local variable scope
378    pub fn push_local_scope(&mut self) {
379        self.local_vars.push(HashMap::new());
380    }
381
382    /// Pop the current local variable scope
383    pub fn pop_local_scope(&mut self) {
384        if !self.local_vars.is_empty() {
385            self.local_vars.pop();
386        }
387    }
388
389    /// Set a local variable in the current scope
390    pub fn set_local_var(&mut self, name: &str, value: String) {
391        if let Some(current_scope) = self.local_vars.last_mut() {
392            current_scope.insert(name.to_string(), value);
393        } else {
394            // If no local scope exists, set as global variable
395            self.set_var(name, value);
396        }
397    }
398
399
400    /// Enter a function context (push local scope if needed)
401    pub fn enter_function(&mut self) {
402        self.function_depth += 1;
403        if self.function_depth > self.local_vars.len() {
404            self.push_local_scope();
405        }
406    }
407
408    /// Exit a function context (pop local scope if needed)
409    pub fn exit_function(&mut self) {
410        if self.function_depth > 0 {
411            self.function_depth -= 1;
412            if self.function_depth == self.local_vars.len() - 1 {
413                self.pop_local_scope();
414            }
415        }
416    }
417
418    /// Set return state for function returns
419    pub fn set_return(&mut self, value: i32) {
420        self.returning = true;
421        self.return_value = Some(value);
422    }
423
424    /// Clear return state
425    pub fn clear_return(&mut self) {
426        self.returning = false;
427        self.return_value = None;
428    }
429
430    /// Check if currently returning
431    pub fn is_returning(&self) -> bool {
432        self.returning
433    }
434
435    /// Get return value if returning
436    pub fn get_return_value(&self) -> Option<i32> {
437        self.return_value
438    }
439
440
441}
442
443
444impl Default for ShellState {
445    fn default() -> Self {
446        Self::new()
447    }
448}
449
450#[cfg(test)]
451mod tests {
452    use super::*;
453
454    #[test]
455    fn test_shell_state_basic() {
456        let mut state = ShellState::new();
457        state.set_var("TEST_VAR", "test_value".to_string());
458        assert_eq!(state.get_var("TEST_VAR"), Some("test_value".to_string()));
459    }
460
461    #[test]
462    fn test_special_variables() {
463        let mut state = ShellState::new();
464        state.set_last_exit_code(42);
465        state.set_script_name("test_script");
466
467        assert_eq!(state.get_var("?"), Some("42".to_string()));
468        assert_eq!(state.get_var("$"), Some(state.shell_pid.to_string()));
469        assert_eq!(state.get_var("0"), Some("test_script".to_string()));
470    }
471
472    #[test]
473    fn test_export_variable() {
474        let mut state = ShellState::new();
475        state.set_var("EXPORT_VAR", "export_value".to_string());
476        state.export_var("EXPORT_VAR");
477
478        let child_env = state.get_env_for_child();
479        assert_eq!(
480            child_env.get("EXPORT_VAR"),
481            Some(&"export_value".to_string())
482        );
483    }
484
485    #[test]
486    fn test_unset_variable() {
487        let mut state = ShellState::new();
488        state.set_var("UNSET_VAR", "value".to_string());
489        state.export_var("UNSET_VAR");
490
491        assert!(state.variables.contains_key("UNSET_VAR"));
492        assert!(state.exported.contains("UNSET_VAR"));
493
494        state.unset_var("UNSET_VAR");
495
496        assert!(!state.variables.contains_key("UNSET_VAR"));
497        assert!(!state.exported.contains("UNSET_VAR"));
498    }
499
500    #[test]
501    fn test_get_user_hostname() {
502        let state = ShellState::new();
503        let user_hostname = state.get_user_hostname();
504        // Should contain @ since it's user@hostname format
505        assert!(user_hostname.contains('@'));
506    }
507
508    #[test]
509    fn test_get_prompt() {
510        let state = ShellState::new();
511        let prompt = state.get_prompt();
512        // Should end with $ and contain @
513        assert!(prompt.ends_with(" $ "));
514        assert!(prompt.contains('@'));
515    }
516
517    #[test]
518    fn test_positional_parameters() {
519        let mut state = ShellState::new();
520        state.set_positional_params(vec![
521            "arg1".to_string(),
522            "arg2".to_string(),
523            "arg3".to_string(),
524        ]);
525
526        assert_eq!(state.get_var("1"), Some("arg1".to_string()));
527        assert_eq!(state.get_var("2"), Some("arg2".to_string()));
528        assert_eq!(state.get_var("3"), Some("arg3".to_string()));
529        assert_eq!(state.get_var("4"), None);
530        assert_eq!(state.get_var("#"), Some("3".to_string()));
531        assert_eq!(state.get_var("*"), Some("arg1 arg2 arg3".to_string()));
532        assert_eq!(state.get_var("@"), Some("arg1 arg2 arg3".to_string()));
533    }
534
535    #[test]
536    fn test_positional_parameters_empty() {
537        let mut state = ShellState::new();
538        state.set_positional_params(vec![]);
539
540        assert_eq!(state.get_var("1"), None);
541        assert_eq!(state.get_var("#"), Some("0".to_string()));
542        assert_eq!(state.get_var("*"), Some("".to_string()));
543        assert_eq!(state.get_var("@"), Some("".to_string()));
544    }
545
546    #[test]
547    fn test_shift_positional_params() {
548        let mut state = ShellState::new();
549        state.set_positional_params(vec![
550            "arg1".to_string(),
551            "arg2".to_string(),
552            "arg3".to_string(),
553        ]);
554
555        assert_eq!(state.get_var("1"), Some("arg1".to_string()));
556        assert_eq!(state.get_var("2"), Some("arg2".to_string()));
557        assert_eq!(state.get_var("3"), Some("arg3".to_string()));
558
559        state.shift_positional_params(1);
560
561        assert_eq!(state.get_var("1"), Some("arg2".to_string()));
562        assert_eq!(state.get_var("2"), Some("arg3".to_string()));
563        assert_eq!(state.get_var("3"), None);
564        assert_eq!(state.get_var("#"), Some("2".to_string()));
565
566        state.shift_positional_params(2);
567
568        assert_eq!(state.get_var("1"), None);
569        assert_eq!(state.get_var("#"), Some("0".to_string()));
570    }
571
572    #[test]
573    fn test_push_positional_param() {
574        let mut state = ShellState::new();
575        state.set_positional_params(vec!["arg1".to_string()]);
576
577        assert_eq!(state.get_var("1"), Some("arg1".to_string()));
578        assert_eq!(state.get_var("#"), Some("1".to_string()));
579
580        state.push_positional_param("arg2".to_string());
581
582        assert_eq!(state.get_var("1"), Some("arg1".to_string()));
583        assert_eq!(state.get_var("2"), Some("arg2".to_string()));
584        assert_eq!(state.get_var("#"), Some("2".to_string()));
585    }
586
587    #[test]
588    fn test_local_variable_scoping() {
589        let mut state = ShellState::new();
590
591        // Set a global variable
592        state.set_var("global_var", "global_value".to_string());
593        assert_eq!(state.get_var("global_var"), Some("global_value".to_string()));
594
595        // Push local scope
596        state.push_local_scope();
597
598        // Set a local variable with the same name
599        state.set_local_var("global_var", "local_value".to_string());
600        assert_eq!(state.get_var("global_var"), Some("local_value".to_string()));
601
602        // Set another local variable
603        state.set_local_var("local_var", "local_only".to_string());
604        assert_eq!(state.get_var("local_var"), Some("local_only".to_string()));
605
606        // Pop local scope
607        state.pop_local_scope();
608
609        // Should be back to global variable
610        assert_eq!(state.get_var("global_var"), Some("global_value".to_string()));
611        assert_eq!(state.get_var("local_var"), None);
612    }
613
614    #[test]
615    fn test_nested_local_scopes() {
616        let mut state = ShellState::new();
617
618        // Set global variable
619        state.set_var("test_var", "global".to_string());
620
621        // Push first local scope
622        state.push_local_scope();
623        state.set_local_var("test_var", "level1".to_string());
624        assert_eq!(state.get_var("test_var"), Some("level1".to_string()));
625
626        // Push second local scope
627        state.push_local_scope();
628        state.set_local_var("test_var", "level2".to_string());
629        assert_eq!(state.get_var("test_var"), Some("level2".to_string()));
630
631        // Pop second scope
632        state.pop_local_scope();
633        assert_eq!(state.get_var("test_var"), Some("level1".to_string()));
634
635        // Pop first scope
636        state.pop_local_scope();
637        assert_eq!(state.get_var("test_var"), Some("global".to_string()));
638    }
639
640    #[test]
641    fn test_variable_set_in_local_scope() {
642        let mut state = ShellState::new();
643
644        // No local scope initially
645        state.set_var("test_var", "global".to_string());
646        assert_eq!(state.get_var("test_var"), Some("global".to_string()));
647
648        // Push local scope and set local variable
649        state.push_local_scope();
650        state.set_local_var("test_var", "local".to_string());
651        assert_eq!(state.get_var("test_var"), Some("local".to_string()));
652
653        // Pop scope
654        state.pop_local_scope();
655        assert_eq!(state.get_var("test_var"), Some("global".to_string()));
656    }
657
658    #[test]
659    fn test_condensed_cwd_environment_variable() {
660        // Test default behavior (should be true for backward compatibility)
661        let state = ShellState::new();
662        assert!(state.condensed_cwd);
663
664        // Test explicit true
665        unsafe {
666            std::env::set_var("RUSH_CONDENSED", "true");
667        }
668        let state = ShellState::new();
669        assert!(state.condensed_cwd);
670
671        // Test explicit false
672        unsafe {
673            std::env::set_var("RUSH_CONDENSED", "false");
674        }
675        let state = ShellState::new();
676        assert!(!state.condensed_cwd);
677
678        // Clean up
679        unsafe {
680            std::env::remove_var("RUSH_CONDENSED");
681        }
682    }
683
684    #[test]
685    fn test_get_full_cwd() {
686        let state = ShellState::new();
687        let full_cwd = state.get_full_cwd();
688        assert!(!full_cwd.is_empty());
689        // Should contain path separators (either / or \ depending on platform)
690        assert!(full_cwd.contains('/') || full_cwd.contains('\\'));
691    }
692
693    #[test]
694    fn test_prompt_with_condensed_setting() {
695        let mut state = ShellState::new();
696
697        // Test with condensed enabled (default)
698        assert!(state.condensed_cwd);
699        let prompt_condensed = state.get_prompt();
700        assert!(prompt_condensed.contains('@'));
701
702        // Test with condensed disabled
703        state.condensed_cwd = false;
704        let prompt_full = state.get_prompt();
705        assert!(prompt_full.contains('@'));
706
707        // Both should end with "$ " (or "# " for root)
708        assert!(prompt_condensed.ends_with("$ ") || prompt_condensed.ends_with("# "));
709        assert!(prompt_full.ends_with("$ ") || prompt_full.ends_with("# "));
710    }
711}
rush_sh/state.rs

rush_sh/
state.rs
